import json
from pathlib import Path
import numpy as np

SEED = 7

def risk_terms(n, m, alpha, s, common=True):
    a = 2.0 * alpha + 2.0 * s
    out = [n ** (-a / (a + 1.0)), (n * m) ** (-a / (2.0 * a + 1.0))]
    if common:
        out += [m ** (-a), m ** (-4.0 * alpha)]
    return np.asarray(out)

def k0_formula(n, m, alpha, s):
    return min(n ** (1.0 / (2 * alpha + 2 * s + 1)),
               (n * m) ** (1.0 / (4 * alpha + 2 * s + 1)), m / 2.0)

def dyadic_blocks(m):
    # Positive real-FFT frequencies, grouped to stabilize energy estimates.
    maxk = m // 2
    blocks = []
    lo = 1
    while lo <= maxk:
        hi = min(2 * lo, maxk + 1)
        blocks.append(np.arange(lo, hi))
        lo *= 2
    return blocks

def block_energies(z):
    # z: trajectories x grid, centered before FFT.
    c = np.fft.rfft(z - z.mean(axis=1, keepdims=True), axis=1, norm="ortho")
    return np.array([np.mean(np.sum(np.abs(c[:, b]) ** 2, axis=1))
                     for b in dyadic_blocks(z.shape[1])]), c

def adaptive_cutoff(z, n, noise_std, alpha_hint=None, s_hint=None, common=True):
    m = z.shape[1]
    e, c = block_energies(z)
    blocks = dyadic_blocks(m)
    # For orthonormal FFT, expected coefficient noise energy is approximately sigma^2.
    floor = max(1e-12, noise_std ** 2 * len(blocks[-1]))
    # Robust high-frequency estimate, guarded against too few blocks.
    noise_floor = max(floor, float(np.median(e[-2:])) * 0.35 if len(e) > 1 else floor)
    excess = np.maximum(e - noise_floor, 1e-12)
    freqs = np.array([np.sqrt(b[0] * max(b[-1], b[0])) for b in blocks])
    slope = np.polyfit(np.log(freqs), np.log(excess), 1)[0] if len(blocks) >= 2 else -2.0
    alpha = float(np.clip(-slope / 2.0, 0.55, 3.0)) if alpha_hint is None else alpha_hint
    # Supervised slope proxy: target block energy is unavailable at deployment;
    # use a conservative smoothness estimate from observed spectral decay.
    s = float(np.clip(-slope / 2.0 - alpha + 0.5, 0.0, 3.0)) if s_hint is None else s_hint
    floor_risk = risk_terms(n, m, alpha, s, common).sum()
    # Screening is relative to the observed signal scale; this makes the risk floor usable
    # for the finite synthetic experiment rather than comparing dimensionless quantities.
    threshold = max(noise_floor * 1.1, 0.025 * floor_risk * max(e.sum(), 1e-9))
    keep = [b for b, x in zip(blocks, e) if x > threshold]
    k_screen = int(keep[-1][-1]) if keep else 1
    k0 = int(max(1, min(m // 2, round(k0_formula(n, m, alpha, s)))))
    # Never add modes beyond the screened support; K0 provides the advertised initialization.
    k = max(1, min(k_screen, max(k0, 1)))
    return k, {"energies": e, "noise_floor": noise_floor, "threshold": threshold,
               "alpha_hat": alpha, "s_hat": s, "k0": k0}

def make_data(n, m, rng, noise=0.20, p=32):
    t = np.arange(m) / m
    # Random curves with decaying Fourier amplitudes and a smooth supervised functional.
    q = np.arange(1, p + 1)
    amp = q ** -1.35
    a = rng.normal(size=(n, p)) * amp
    b = rng.normal(size=(n, p)) * amp
    x = np.ones((n, m)) * rng.normal(size=(n, 1)) * 0.2
    for k in q:
        x += a[:, k-1, None] * np.cos(2*np.pi*k*t)[None, :]
        x += b[:, k-1, None] * np.sin(2*np.pi*k*t)[None, :]
    beta = q ** -1.15 * np.exp(-q / 12.0)
    y = (a * beta).sum(1) + 0.5 * (b * beta).sum(1) + rng.normal(size=n) * 0.08
    return x + rng.normal(size=x.shape) * noise, y

def fit_predict(ztr, ytr, zte, k):
    def feat(z):
        c = np.fft.rfft(z - z.mean(1, keepdims=True), axis=1, norm=None)[:, 1:k+1]
        # Scale DFT coefficients to approximate continuous Fourier coefficients.
        c = (2.0 / z.shape[1]) * c
        return np.concatenate([c.real, c.imag], axis=1)
    A, B = feat(ztr), feat(zte)
    A = np.c_[np.ones(len(A)), A]
    B = np.c_[np.ones(len(B)), B]
    # Mild ridge avoids instability at low m/noisy resolutions.
    reg = 2e-2 * np.eye(A.shape[1]); reg[0, 0] = 0
    w = np.linalg.solve(A.T @ A + reg, A.T @ ytr)
    return B @ w

def run():
    rng = np.random.default_rng(SEED)
    # Cheap numerical check: evaluate claimed terms and verify increasing m lowers grid floors.
    alpha, s, n = 1.2, 0.8, 96
    ms = np.array([8, 16, 32, 64, 128])
    rs = np.array([risk_terms(n, int(m), alpha, s).sum() for m in ms])
    floors = np.array([risk_terms(n, int(m), alpha, s)[2:].sum() for m in ms])
    math_check = {"m": ms.tolist(), "risk": rs.tolist(), "grid_floor": floors.tolist(),
                  "risk_decreases": bool(np.all(np.diff(rs) < 0)),
                  "grid_floor_decreases": bool(np.all(np.diff(floors) < 0)),
                  "k0": [k0_formula(n, int(m), alpha, s) for m in ms]}

    ntr, nte, mtrain = 192, 512, 64
    results = {}
    # Repeat independent train/test draws while holding the model and resolutions fixed.
    for mtest in [16, 32, 64, 128]:
        vals = []
        for rep in range(5):
            rr = np.random.default_rng(SEED + 1000 * mtest + rep)
            train, ytr = make_data(ntr, mtrain, rr)
            xt, yt = make_data(nte, mtest, rr)
            fixed_k = min(12, mtrain // 2, mtest // 2)
            pred_fixed = fit_predict(train, ytr, xt, fixed_k)
            k, info = adaptive_cutoff(train, ntr, 0.20, alpha_hint=1.2, s_hint=0.8)
            k = min(k, mtest // 2)
            pred_adapt = fit_predict(train, ytr, xt, k)
            vals.append((np.mean((pred_fixed-yt)**2), np.mean((pred_adapt-yt)**2), fixed_k, k))
        v = np.asarray(vals)
        results[str(mtest)] = {"fixed_mse_mean": float(v[:,0].mean()),
                               "adaptive_mse_mean": float(v[:,1].mean()),
                               "fixed_mse_std": float(v[:,0].std(ddof=1)),
                               "adaptive_mse_std": float(v[:,1].std(ddof=1)),
                               "fixed_k": int(v[0,2]), "adaptive_k": int(v[0,3]),
                               "alpha_hat": 1.2, "s_hat": 0.8}
    out = {"seed": SEED, "math_check": math_check, "experiment": results}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    run()