import numpy as np
from scipy.optimize import minimize, differential_evolution


def compute_c1(f_values, n_points):
    dx = 0.5 / n_points
    f_nn = np.maximum(f_values, 0.0)
    autoconv = np.convolve(f_nn, f_nn, mode='full') * dx
    integral_sq = (np.sum(f_nn) * dx) ** 2
    if integral_sq < 1e-12:
        return 1e10
    return np.max(autoconv) / integral_sq


def heights_to_function(heights, n_eval):
    """Convert step heights to fine-grained function values"""
    n_steps = len(heights)
    f = np.zeros(n_eval)
    step_width = n_eval / n_steps
    for i in range(n_steps):
        start = int(i * step_width)
        end = int((i + 1) * step_width)
        end = min(end, n_eval)
        f[start:end] = max(heights[i], 0.0)
    return f


def optimize_step_function(n_steps, n_eval, seed=42):
    dx = 0.5 / n_eval

    def objective(heights):
        f = heights_to_function(np.maximum(heights, 0.0), n_eval)
        return compute_c1(f, n_eval)

    # Differential evolution
    bounds = [(0.01, 5.0)] * n_steps
    result = differential_evolution(
        objective, bounds, seed=seed,
        maxiter=2000, tol=1e-12, polish=True,
        popsize=40, mutation=(0.5, 1.5), recombination=0.9,
        workers=1,
    )

    heights = np.maximum(result.x, 0.0)
    f = heights_to_function(heights, n_eval)
    c1 = compute_c1(f, n_eval)

    # Polish with Nelder-Mead
    result2 = minimize(objective, heights, method='Nelder-Mead',
                       options={'maxiter': 50000, 'xatol': 1e-12, 'fatol': 1e-12})
    heights2 = np.maximum(result2.x, 0.0)
    f2 = heights_to_function(heights2, n_eval)
    c1_2 = compute_c1(f2, n_eval)

    if c1_2 < c1:
        return f2, c1_2, heights2
    return f, c1, heights


def run():
    best_c1 = float('inf')
    best_f = None
    best_n = None
    n_eval = 4000

    configs = [
        (10, 42), (15, 42), (20, 42), (25, 42), (30, 42),
        (40, 42), (50, 42), (60, 42),
        (20, 0), (30, 0), (40, 0), (50, 0),
        (20, 7), (30, 7), (40, 7),
    ]

    for n_steps, seed in configs:
        print(f"n_steps={n_steps}, seed={seed}...", end=" ", flush=True)
        f, c1, heights = optimize_step_function(n_steps, n_eval, seed)
        print(f"C1={c1:.10f}")

        if c1 < best_c1:
            best_c1 = c1
            best_f = f
            best_n = n_eval
            print(f"  *** NEW BEST: C1={c1:.10f}, heights={np.round(heights, 4)}")

    print(f"\nFinal best C1: {best_c1:.10f}")
    return best_f, best_c1, best_c1, best_n