"""
1. Run seed 100 at N=4000, save params
2. Upsample to N=8000, refine
3. Perturbation restarts from best
4. Also scan seeds 80-130 for better starts
"""
import sys
import jax
import jax.numpy as jnp
import numpy as np
from scipy.optimize import minimize as scipy_minimize
import optax


def compute_c1_numpy(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 optimize(N, seed=None, init_params_np=None, adam_steps=80000, lr=0.005, temp=300.0):
    dx = 0.5 / N

    @jax.jit
    def obj_smooth(params, t):
        f = jnp.exp(jnp.clip(params, -8, 4))
        padded = jnp.zeros(2 * N)
        padded = padded.at[:N].set(f)
        fft_f = jnp.fft.rfft(padded)
        conv = jnp.fft.irfft(fft_f * fft_f, n=2 * N) * dx
        integral_sq = (jnp.sum(f) * dx) ** 2
        smooth_max = jax.nn.logsumexp(t * conv) / t
        return smooth_max / integral_sq

    @jax.jit
    def obj_hard(params):
        f = jnp.exp(jnp.clip(params, -8, 4))
        padded = jnp.zeros(2 * N)
        padded = padded.at[:N].set(f)
        fft_f = jnp.fft.rfft(padded)
        conv = jnp.fft.irfft(fft_f * fft_f, n=2 * N) * dx
        integral_sq = (jnp.sum(f) * dx) ** 2
        return jnp.max(conv) / integral_sq

    grad_smooth = jax.jit(jax.grad(obj_smooth))
    grad_hard = jax.jit(jax.grad(obj_hard))

    if init_params_np is not None:
        if len(init_params_np) != N:
            # Upsample
            old_f = np.exp(np.clip(init_params_np, -8, 4))
            new_f = np.interp(np.linspace(0, 1, N), np.linspace(0, 1, len(init_params_np)), old_f)
            params = jnp.array(np.log(np.maximum(new_f, 1e-6)))
        else:
            params = jnp.array(init_params_np)
    else:
        np.random.seed(seed)
        init_f = np.ones(N) * 0.5 + 0.02 * np.random.randn(N)
        params = jnp.array(np.log(np.maximum(init_f, 1e-6)))

    lr_schedule = optax.warmup_cosine_decay_schedule(
        init_value=0.0, peak_value=lr, warmup_steps=min(2000, adam_steps//5),
        decay_steps=adam_steps - min(2000, adam_steps//5), end_value=lr * 1e-4,
    )
    optimizer = optax.adam(learning_rate=lr_schedule)
    opt_state = optimizer.init(params)

    best_c1 = float('inf')
    best_params = params

    for step in range(adam_steps):
        loss, grads = jax.value_and_grad(obj_smooth)(params, temp)
        updates, opt_state = optimizer.update(grads, opt_state, params)
        params = optax.apply_updates(params, updates)

        if step >= adam_steps - 2000:
            hc = float(obj_hard(params))
            if hc < best_c1:
                best_c1 = hc
                best_params = params

    # L-BFGS polish
    params_np = np.array(best_params, dtype=np.float64)
    for t in [1000.0, 10000.0]:
        def scipy_obj(p):
            p_jax = jnp.array(p)
            val = float(obj_smooth(p_jax, t))
            g = np.array(grad_smooth(p_jax, t), dtype=np.float64)
            return val, g
        result = scipy_minimize(
            scipy_obj, params_np, method='L-BFGS-B', jac=True,
            options={'maxiter': 5000, 'ftol': 1e-15, 'gtol': 1e-14, 'maxcor': 100},
        )
        params_np = result.x

    for _ in range(3):
        def scipy_hard(p):
            p_jax = jnp.array(p)
            val = float(obj_hard(p_jax))
            g = np.array(grad_hard(p_jax), dtype=np.float64)
            return val, g
        result = scipy_minimize(
            scipy_hard, params_np, method='L-BFGS-B', jac=True,
            options={'maxiter': 20000, 'ftol': 1e-16, 'gtol': 1e-15, 'maxcor': 100},
        )
        params_np = result.x

    f_final = np.exp(np.clip(params_np, -8, 4))
    c1_final = compute_c1_numpy(f_final, N)
    return params_np, f_final, c1_final


def run():
    N = 4000
    best_c1_overall = float('inf')
    best_f_overall = None
    best_params_overall = None

    # Scan seeds 80-130 (neighborhood of best seed 100)
    sys.stdout.write("Phase 1: Scanning seeds at N=4000\n")
    sys.stdout.flush()

    for seed in range(80, 131):
        params, f, c1 = optimize(N, seed=seed, adam_steps=60000)
        sys.stdout.write(f"  Seed {seed:3d}: C1={c1:.10f}")
        if c1 < best_c1_overall:
            best_c1_overall = c1
            best_f_overall = f
            best_params_overall = params
            sys.stdout.write(" ***")
        sys.stdout.write("\n")
        sys.stdout.flush()

    sys.stdout.write(f"\nBest after scan: C1={best_c1_overall:.10f}\n")

    # Phase 2: Upsample best to N=8000 and refine
    sys.stdout.write("Phase 2: Upsample to N=8000\n")
    sys.stdout.flush()
    N2 = 8000
    params2, f2, c1_2 = optimize(N2, init_params_np=best_params_overall,
                                   adam_steps=40000, lr=0.002)
    sys.stdout.write(f"  N=8000: C1={c1_2:.10f}\n")
    sys.stdout.flush()

    if c1_2 < best_c1_overall:
        best_c1_overall = c1_2
        best_f_overall = f2
        best_params_overall = params2
        N = N2

    # Phase 3: Perturbation restarts
    sys.stdout.write("Phase 3: Perturbation restarts\n")
    sys.stdout.flush()
    for i in range(5):
        key = jax.random.PRNGKey(i * 17)
        noise = 0.03 * jax.random.normal(key, shape=(len(best_params_overall),))
        perturbed = best_params_overall + np.array(noise)
        p, f, c1 = optimize(len(best_f_overall), init_params_np=perturbed,
                            adam_steps=30000, lr=0.002)
        sys.stdout.write(f"  Perturb {i}: C1={c1:.10f}")
        if c1 < best_c1_overall:
            best_c1_overall = c1
            best_f_overall = f
            best_params_overall = p
            N = len(f)
            sys.stdout.write(" ***")
        sys.stdout.write("\n")
        sys.stdout.flush()

    sys.stdout.write(f"\nFinal C1: {best_c1_overall:.10f}\n")
    sys.stdout.flush()
    return best_f_overall, best_c1_overall, best_c1_overall, len(best_f_overall)