"""
Step function optimization with EXACT autoconvolution computation.

For a step function f = sum_i h_i * 1_{[a_i, b_i)}, the autoconvolution
(f*f)(t) = sum_{i,j} h_i * h_j * overlap(t, [a_i,b_i), [a_j,b_j))
where overlap is the length of intersection of [a_i,b_i) and [t-b_j, t-a_j).

This is piecewise linear in t, so we can find the exact max at breakpoints.
"""
import numpy as np
from scipy.optimize import minimize, differential_evolution
import time

def autoconv_step(heights, widths):
    """Compute exact autoconvolution of a step function at all breakpoints.

    f(x) = heights[i] for x in [cum_widths[i], cum_widths[i+1])
    where cum_widths[0] = 0.

    Returns: (breakpoints, autoconv_values) - the autoconvolution evaluated at breakpoints.
    """
    K = len(heights)
    # Compute interval endpoints
    endpoints = np.zeros(K + 1)
    endpoints[1:] = np.cumsum(widths)

    # Collect all breakpoints: t = a_i + a_j, a_i + b_j, b_i + a_j, b_i + b_j
    bp_set = set()
    for i in range(K):
        for j in range(K):
            ai, bi = endpoints[i], endpoints[i+1]
            aj, bj = endpoints[j], endpoints[j+1]
            # Overlap of [ai, bi) and [t-bj, t-aj) is nonzero when:
            # t-bj < bi and t-aj > ai, i.e., ai+aj < t < bi+bj
            # Breakpoints where the overlap function changes slope:
            for bp in [ai+aj, ai+bj, bi+aj, bi+bj]:
                bp_set.add(bp)

    breakpoints = np.array(sorted(bp_set))
    # Also add midpoints between consecutive breakpoints (for quadratic terms)
    # Actually for piecewise constant f, autoconv is piecewise LINEAR, so
    # max is at a breakpoint. But let's also check midpoints to be safe.
    midpoints = (breakpoints[:-1] + breakpoints[1:]) / 2
    all_t = np.concatenate([breakpoints, midpoints])
    all_t = np.sort(all_t)

    # Evaluate autoconvolution at each t
    values = np.zeros(len(all_t))
    for idx, t in enumerate(all_t):
        val = 0.0
        for i in range(K):
            for j in range(K):
                ai, bi = endpoints[i], endpoints[i+1]
                aj, bj = endpoints[j], endpoints[j+1]
                # Overlap of [ai, bi) and [t-bj, t-aj)
                lo = max(ai, t - bj)
                hi = min(bi, t - aj)
                if hi > lo:
                    val += heights[i] * heights[j] * (hi - lo)
        values[idx] = val

    return all_t, values


def compute_c1_step(heights, widths):
    """Compute C1 for a step function."""
    t_vals, conv_vals = autoconv_step(heights, widths)
    integral = np.sum(heights * widths)
    if integral < 1e-15:
        return 1e10
    return float(np.max(conv_vals) / integral**2)


def optimize_step_scipy(K, n_trials=100, method='de'):
    """Optimize step function with K equal-width steps."""
    width = 0.5 / K
    widths = np.full(K, width)

    best_c1 = np.inf
    best_heights = None

    if method == 'de':
        # Differential evolution on heights
        bounds = [(0.0, 10.0)] * K

        def obj(h):
            h = np.maximum(h, 0.0)
            return compute_c1_step(h, widths)

        result = differential_evolution(obj, bounds, maxiter=1000,
                                         seed=42, tol=1e-12,
                                         popsize=30, mutation=(0.5, 1.5),
                                         recombination=0.9)
        best_heights = np.maximum(result.x, 0.0)
        best_c1 = result.fun

    else:
        # Multi-start L-BFGS-B
        for trial in range(n_trials):
            np.random.seed(trial)
            h0 = np.abs(np.random.randn(K)) + 0.5

            def obj(h):
                h = np.maximum(h, 0.0)
                return compute_c1_step(h, widths)

            try:
                result = minimize(obj, h0, method='Nelder-Mead',
                                options={'maxiter': 10000, 'xatol': 1e-12, 'fatol': 1e-12})
                c1 = result.fun
                if c1 < best_c1:
                    best_c1 = c1
                    best_heights = np.maximum(result.x, 0.0)
            except:
                pass

    return best_heights, best_c1


def optimize_step_variable_width(K, n_trials=50):
    """Optimize both heights and widths of step function."""
    best_c1 = np.inf
    best_h = None
    best_w = None

    for trial in range(n_trials):
        np.random.seed(trial * 137 + 42)

        # Parameterize widths via softmax: w_i = exp(p_i) / sum(exp(p_j)) * 0.5
        # Heights: h_i = q_i^2
        p0 = np.random.randn(K) * 0.5
        q0 = np.abs(np.random.randn(K)) + 0.5
        x0 = np.concatenate([p0, q0])

        def obj(x):
            p = x[:K]
            q = x[K:]
            # Softmax for widths
            ep = np.exp(p - np.max(p))
            widths = ep / np.sum(ep) * 0.5
            heights = q ** 2
            return compute_c1_step(heights, widths)

        try:
            result = minimize(obj, x0, method='Nelder-Mead',
                            options={'maxiter': 50000, 'xatol': 1e-12, 'fatol': 1e-12})
            if result.fun < best_c1:
                best_c1 = result.fun
                x = result.x
                p = x[:K]
                ep = np.exp(p - np.max(p))
                best_w = ep / np.sum(ep) * 0.5
                best_h = x[K:] ** 2
                print(f"  Trial {trial}: C1 = {best_c1:.10f}")
        except:
            pass

    return best_h, best_w, best_c1


# Main search
print("=== Step Function Optimization ===\n")

# Equal-width steps
print("Equal-width steps (DE):")
for K in [5, 8, 10, 15, 20, 30, 50]:
    t0 = time.time()
    h, c1 = optimize_step_scipy(K, method='de')
    elapsed = time.time() - t0
    print(f"  K={K:3d}: C1 = {c1:.10f} ({elapsed:.1f}s)")
    if c1 < 1.51:
        print(f"    Heights: {h}")

# Variable-width steps
print("\nVariable-width steps:")
for K in [5, 8, 10, 15, 20]:
    t0 = time.time()
    h, w, c1 = optimize_step_variable_width(K, n_trials=30)
    elapsed = time.time() - t0
    print(f"  K={K:3d}: C1 = {c1:.10f} ({elapsed:.1f}s)")
    if c1 < 1.51:
        print(f"    Heights: {h}")
        print(f"    Widths: {w}")