"""
Mathematical constructions for low C1.

Key insight: we need functions where the sumset A+A is uniformly distributed.
This is a problem in additive combinatorics.

Constructions to try:
1. Quadratic residue indicator functions
2. Functions based on polynomial evaluation over finite fields
3. B-spline wavelets
4. Optimized step functions via DE at small K
5. Functions that equalize the autoconvolution
"""
import numpy as np
from scipy.optimize import minimize, differential_evolution
from scipy.interpolate import interp1d
import time

def compute_c1_fft(f_values, dx):
    f = np.maximum(f_values, 0.0)
    N = len(f)
    M = 2 * N
    fft_f = np.fft.rfft(f, n=M)
    conv = np.fft.irfft(fft_f * fft_f, n=M) * dx
    integral_sq = (np.sum(f) * dx) ** 2
    if integral_sq < 1e-20:
        return 1e10
    return float(np.max(conv) / integral_sq)

def compute_c1_smooth_and_grad(f, N, dx, alpha=200.0):
    M = 2 * N
    fft_f = np.fft.rfft(f, n=M)
    conv = np.fft.irfft(fft_f * fft_f, n=M) * dx
    integral = np.sum(f) * dx
    if integral < 1e-15:
        return 1e10, np.zeros(N)
    integral_sq = integral ** 2
    max_val = np.max(conv)
    shifted = conv - max_val
    mask = shifted > -50.0 / alpha
    weights = np.zeros_like(conv)
    weights[mask] = np.exp(alpha * shifted[mask])
    sum_w = np.sum(weights)
    if sum_w < 1e-30:
        weights[np.argmax(conv)] = 1.0
        sum_w = 1.0
    smooth_max = max_val + np.log(sum_w) / alpha
    softmax_w = weights / sum_w
    c1 = smooth_max / integral_sq
    fft_sw = np.fft.rfft(softmax_w, n=M)
    fft_fp = np.fft.rfft(f, n=M)
    corr = np.fft.irfft(fft_sw * np.conj(fft_fp), n=M)
    grad_f = 2.0 * corr[:N] * dx / integral_sq - 2.0 * smooth_max * dx / (integral**3)
    return c1, grad_f

def opt_pipeline(f_init, N, dx, maxiter=2000):
    """Full optimization pipeline."""
    params = np.sqrt(np.maximum(f_init, 0.0) + 1e-12)
    for alpha in [0.5, 5.0, 50.0, 500.0, 5000.0, 50000.0]:
        def obj(p, a=alpha):
            f = p ** 2
            c1, g = compute_c1_smooth_and_grad(f, N, dx, a)
            return c1, g * 2 * p
        result = minimize(obj, params, jac=True, method='L-BFGS-B',
                         options={'maxiter': maxiter, 'ftol': 1e-16, 'gtol': 1e-15})
        params = result.x
    f_out = params ** 2
    return f_out, compute_c1_fft(f_out, dx)

def upscale(f, N_new):
    N_old = len(f)
    x_old = np.linspace(0, 1, N_old)
    x_new = np.linspace(0, 1, N_new)
    interp = interp1d(x_old, f, kind='linear', fill_value=0.0, bounds_error=False)
    return np.maximum(interp(x_new), 0.0)

# =====================================================
# Construction 1: Differential evolution on step heights
# =====================================================
print("=== DE on step functions ===")
t0 = time.time()

best_results = []

for K in [10, 15, 20, 25, 30, 40, 50, 75, 100]:
    N_k = K * 10  # Use higher resolution per step
    dx_k = 0.5 / N_k
    step_size = N_k // K

    def obj_de(heights):
        f = np.repeat(np.maximum(heights, 0.0), step_size)[:N_k]
        return compute_c1_fft(f, dx_k)

    bounds = [(0.0, 5.0)] * K
    result = differential_evolution(obj_de, bounds, maxiter=500, seed=42,
                                     popsize=20, tol=1e-12)
    c1_raw = result.fun

    # Refine with gradient optimization
    f_step = np.repeat(np.maximum(result.x, 0.0), step_size)[:N_k]
    f_opt, c1_opt = opt_pipeline(f_step, N_k, dx_k, maxiter=1000)

    best_results.append((K, c1_opt, f_opt, N_k, dx_k))
    elapsed = time.time() - t0
    print(f"  K={K:3d}: DE C1={c1_raw:.8f}, opt C1={c1_opt:.8f} ({elapsed:.0f}s)")

# =====================================================
# Construction 2: Functions from finite field arithmetic
# =====================================================
print("\n=== Finite field constructions ===")

# For prime p, define f(x) = indicator of {x^2 mod p : x in [0,p)} scaled to [0, 0.5]
for p in [31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 127, 131]:
    N_p = p * 10
    dx_p = 0.5 / N_p

    # Quadratic residues
    qr = set()
    for a in range(p):
        qr.add((a * a) % p)

    f = np.zeros(N_p)
    for r in qr:
        block_start = int(r * N_p / p)
        block_end = int((r + 1) * N_p / p)
        f[block_start:block_end] = 1.0

    f_opt, c1 = opt_pipeline(f, N_p, dx_p, maxiter=1000)
    if c1 < 1.515:
        best_results.append(('QR', c1, f_opt, N_p, dx_p))
        print(f"  p={p}: C1 = {c1:.8f} ***")

# =====================================================
# Construction 3: Autoconvolution equalization
# Iteratively adjust f to make autoconvolution more uniform
# =====================================================
print("\n=== Autoconvolution equalization ===")

for N_eq in [200, 500]:
    dx_eq = 0.5 / N_eq

    # Start from various initializations
    for seed in range(20):
        np.random.seed(seed + 3000)
        f = np.random.exponential(1.0, N_eq) + 0.1

        # Iterative equalization
        for iteration in range(50):
            f_nn = np.maximum(f, 0.0)
            M = 2 * N_eq
            fft_f = np.fft.rfft(f_nn, n=M)
            conv = np.fft.irfft(fft_f * fft_f, n=M) * dx_eq

            # Identify high points in autoconvolution
            max_val = np.max(conv)
            threshold = max_val * 0.99
            high_mask = conv > threshold

            # The autoconvolution at position t is sum_j f[j]*f[t-j]*dx
            # To reduce it, we need to reduce f at positions that contribute
            # This is related to the gradient

            # Simple approach: multiply f by a correction factor
            # based on how much each position contributes to the max
            power_spectrum = np.abs(fft_f) ** 2
            # Reduce the highest frequency components
            cutoff = len(power_spectrum) // 2
            damping = np.ones(len(power_spectrum))
            top_freqs = np.argsort(power_spectrum[1:cutoff])[::-1][:10] + 1
            damping[top_freqs] *= 0.95
            fft_mod = fft_f * damping
            f = np.fft.irfft(fft_mod, n=M)[:N_eq]
            f = np.maximum(f, 0.0) + 0.001

        f_opt, c1 = opt_pipeline(f, N_eq, dx_eq, maxiter=1000)
        if c1 < 1.515:
            best_results.append(('EQ', c1, f_opt, N_eq, dx_eq))
            print(f"  N={N_eq}, seed={seed}: C1 = {c1:.8f}")

# =====================================================
# Construction 4: Optimized B-spline basis
# =====================================================
print("\n=== B-spline basis ===")

for n_knots in [5, 10, 15, 20, 30]:
    N_bs = 500
    dx_bs = 0.5 / N_bs
    x = np.linspace(0, 0.5, N_bs, endpoint=False)

    # B-spline basis with n_knots uniformly spaced knots
    knots = np.linspace(0, 0.5, n_knots + 2)
    basis = np.zeros((N_bs, n_knots))
    for i in range(n_knots):
        # Triangular basis function
        center = knots[i + 1]
        width = knots[i + 2] - knots[i]
        basis[:, i] = np.maximum(1 - np.abs(x - center) / (width/2), 0)

    def obj_bs(coeffs):
        f = np.maximum(basis @ coeffs, 0.0)
        return compute_c1_fft(f, dx_bs)

    for seed in range(20):
        np.random.seed(seed + 4000)
        c0 = np.random.exponential(1.0, n_knots)

        result = minimize(obj_bs, c0, method='Nelder-Mead',
                         options={'maxiter': 10000, 'xatol': 1e-10, 'fatol': 1e-10})
        c1 = result.fun
        if c1 < 1.515:
            f_bs = np.maximum(basis @ result.x, 0.0)
            f_opt, c1_opt = opt_pipeline(f_bs, N_bs, dx_bs, maxiter=1000)
            best_results.append(('BS', c1_opt, f_opt, N_bs, dx_bs))
            print(f"  n_knots={n_knots}, seed={seed}: C1 = {c1:.4f} → {c1_opt:.8f}")

# =====================================================
# Sort all results and upscale the best to N=5000
# =====================================================
# Filter to tuples with the right structure
valid = [(r[1], r[2], r[3], r[4]) for r in best_results if isinstance(r[0], (int, str))]
valid.sort(key=lambda x: x[0])

print(f"\n=== Top 5 results ===")
for i, (c1, f, N_r, dx_r) in enumerate(valid[:5]):
    print(f"  #{i+1}: C1 = {c1:.10f} at N={N_r}")

if valid:
    best_c1, best_f, best_N, best_dx = valid[0]

    # Upscale to N=5000 and refine
    print(f"\nUpscaling best (C1={best_c1:.10f}) to N=5000...")
    N_final = 5000
    dx_final = 0.5 / N_final
    f_up = upscale(best_f, N_final)
    f_final, c1_final = opt_pipeline(f_up, N_final, dx_final, maxiter=5000)

    # Alpha cycling
    params = np.sqrt(np.maximum(f_final, 0.0) + 1e-12)
    for cycle in range(5):
        for alpha in [0.5, 2.0, 10.0]:
            def obj(p, a=alpha):
                f = p ** 2
                c1, g = compute_c1_smooth_and_grad(f, N_final, dx_final, a)
                return c1, g * 2 * p
            result = minimize(obj, params, jac=True, method='L-BFGS-B',
                             options={'maxiter': 500, 'ftol': 1e-16, 'gtol': 1e-15})
            params = result.x
        for alpha in [100.0, 1000.0, 10000.0, 100000.0]:
            def obj(p, a=alpha):
                f = p ** 2
                c1, g = compute_c1_smooth_and_grad(f, N_final, dx_final, a)
                return c1, g * 2 * p
            result = minimize(obj, params, jac=True, method='L-BFGS-B',
                             options={'maxiter': 2000, 'ftol': 1e-16, 'gtol': 1e-15})
            params = result.x
        f_out = params ** 2
        c1 = compute_c1_fft(f_out, dx_final)
        if c1 < c1_final:
            c1_final = c1
            f_final = f_out.copy()

    print(f"Final at N=5000: C1 = {c1_final:.10f}")
    print(f"Score: {1.5052939684401607 / c1_final:.10f}")

    # Compare with existing best
    try:
        f_existing = np.load('/workspace/best_f_5000.npy')
        c1_existing = compute_c1_fft(f_existing, 0.5/len(f_existing))
        print(f"Existing best: C1 = {c1_existing:.10f}")
        if c1_final < c1_existing:
            np.save('/workspace/best_f_5000_math.npy', f_final)
            print("NEW BEST! Saved to best_f_5000_math.npy")
    except:
        pass

print(f"\nTotal time: {time.time()-t0:.0f}s")