"""
Evolutionary search over step function PROGRAMS, inspired by alphaevolve.

The key insight: alphaevolve evolves CODE that generates step functions.
This means the optimal function might have a specific ALGORITHMIC structure.

Let's search over step function programs parameterized by:
1. Number of steps K
2. Height formula: h_i = f(i, K, params) for some simple formula
3. Width formula: w_i = g(i, K, params)

Examples of formulas:
- h_i = a * sin(b * i/K + c)²  (sinusoidal)
- h_i = a * (i/K)^b             (power law)
- h_i = a * exp(-b * (i/K - c)²) (Gaussian)
- h_i = a if i in S else 0       (indicator of set S)
"""
import numpy as np
from scipy.optimize import minimize, differential_evolution
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(params, N, dx, alpha, maxiter):
    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})
    return result.x

def full_opt_at_N(f_init, N):
    dx = 0.5 / N
    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]:
        params = opt(params, N, dx, alpha, 2000)
    # Alpha cycling
    for _ in range(3):
        for alpha in [0.5, 2.0, 10.0]:
            params = opt(params, N, dx, alpha, 300)
        for alpha in [100.0, 1000.0, 10000.0, 100000.0]:
            params = opt(params, N, dx, alpha, 1000)
    f_out = params ** 2
    return f_out, compute_c1_fft(f_out, dx)


# Generate step functions from various formulas
N_eval = 5000  # Resolution for evaluation
t0 = time.time()
best_c1 = np.inf
best_f = None

# Formula 1: Equally spaced steps with sinusoidal heights
# h_i = A + B * sin(2*pi*freq*i/K + phase)^2
print("=== Sinusoidal step heights ===")
for K in [20, 30, 50, 75, 100]:
    for freq in range(1, 10):
        bounds = [(0.01, 5.0), (0.01, 5.0), (0, 2*np.pi)]  # A, B, phase

        def make_f(params, K=K, freq=freq):
            A, B, phase = params
            i = np.arange(K)
            heights = A + B * np.sin(2*np.pi*freq*i/K + phase)**2
            heights = np.maximum(heights, 0.0)
            # Map to N_eval points
            f = np.repeat(heights, N_eval // K + 1)[:N_eval]
            return f

        def obj(params, K=K, freq=freq):
            f = make_f(params, K, freq)
            return compute_c1_fft(f, 0.5/N_eval)

        result = differential_evolution(obj, bounds, maxiter=200, seed=42, tol=1e-12)
        if result.fun < best_c1:
            f_step = make_f(result.x)
            f_opt, c1 = full_opt_at_N(f_step, N_eval)
            if c1 < best_c1:
                best_c1 = c1
                best_f = f_opt.copy()
                print(f"  K={K} freq={freq}: C1 = {c1:.10f} ***")

# Formula 2: Power law heights with varying support
print("\n=== Power law + support fraction ===")
for K in [30, 50, 100]:
    for power in np.linspace(0.2, 3.0, 8):
        for frac in np.linspace(0.3, 1.0, 8):
            i = np.arange(K)
            heights = (i + 1.0) ** (-power)
            # Randomly shuffle (use a specific pattern)
            # Actually, let's use different orderings
            for ordering in ['asc', 'desc', 'interleave', 'random']:
                if ordering == 'asc':
                    h = np.sort(heights)
                elif ordering == 'desc':
                    h = np.sort(heights)[::-1]
                elif ordering == 'interleave':
                    h_sorted = np.sort(heights)
                    h = np.zeros_like(h_sorted)
                    h[::2] = h_sorted[:len(h[::2])]
                    h[1::2] = h_sorted[len(h[::2]):][::-1]
                else:
                    np.random.seed(42 + K + int(power*10))
                    h = heights.copy()
                    np.random.shuffle(h)

                # Zero out (1-frac) of the steps
                n_active = max(1, int(K * frac))
                h[n_active:] = 0.0

                f = np.repeat(h, N_eval // K + 1)[:N_eval]
                c1_raw = compute_c1_fft(f, 0.5/N_eval)

                if c1_raw < 2.0:  # Only optimize promising ones
                    f_opt, c1 = full_opt_at_N(f, N_eval)
                    if c1 < best_c1:
                        best_c1 = c1
                        best_f = f_opt.copy()
                        print(f"  K={K} power={power:.1f} frac={frac:.1f} {ordering}: C1 = {c1:.10f} ***")
                        break  # Move to next K/power/frac

elapsed = time.time() - t0
print(f"\nPhase 1 complete: best C1 = {best_c1:.10f} ({elapsed:.0f}s)")

# Formula 3: Binary patterns from number theory
print("\n=== Number-theoretic binary patterns ===")
for p in [31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97]:
    # Quadratic residues
    qr = set()
    for a in range(p):
        qr.add((a * a) % p)

    # Pattern: indicator of QR
    h = np.array([1.0 if i in qr else 0.0 for i in range(p)])
    f = np.repeat(h, N_eval // p + 1)[:N_eval]
    f_opt, c1 = full_opt_at_N(f, N_eval)
    if c1 < best_c1:
        best_c1 = c1
        best_f = f_opt.copy()
        print(f"  QR p={p}: C1 = {c1:.10f} ***")

    # Also try: indicator of non-QR
    h_nqr = 1.0 - h
    h_nqr[0] = 0.0  # 0 is neither QR nor NQR
    f = np.repeat(h_nqr, N_eval // p + 1)[:N_eval]
    f_opt, c1 = full_opt_at_N(f, N_eval)
    if c1 < best_c1:
        best_c1 = c1
        best_f = f_opt.copy()
        print(f"  NQR p={p}: C1 = {c1:.10f} ***")

    # Try with height proportional to Legendre symbol value
    for h_val in [0.5, 1.0, 2.0]:
        h_leg = np.zeros(p)
        for i in range(p):
            if i in qr and i > 0:
                h_leg[i] = h_val
            elif i > 0:
                h_leg[i] = 1.0
        f = np.repeat(h_leg, N_eval // p + 1)[:N_eval]
        f_opt, c1 = full_opt_at_N(f, N_eval)
        if c1 < best_c1:
            best_c1 = c1
            best_f = f_opt.copy()
            print(f"  Leg p={p} h={h_val}: C1 = {c1:.10f} ***")

elapsed = time.time() - t0
print(f"\nFinal best: C1 = {best_c1:.10f} ({elapsed:.0f}s)")
print(f"Score: {1.5052939684401607 / best_c1:.10f}")

# Compare with existing
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 at N=5000: C1 = {c1_existing:.10f}")
if best_c1 < c1_existing:
    np.save('/workspace/best_f_5000_evolve.npy', best_f)
    print("NEW BEST! Saved.")