import json
from pathlib import Path
import numpy as np

SEED = 1109
rng = np.random.default_rng(SEED)
D = 256
A = 1.0
u = (np.arange(D) + 0.5) / D
omega = 2*np.pi * u**(1/(A+1))
q = np.ones(D) / D


def return_proxy(lam, weights, ks):
    lam = np.asarray(lam, complex)
    weights = np.asarray(weights, float)
    weights = weights / weights.sum()
    amp = np.sum(weights[None, :] * lam[None, :] ** ks[:, None], axis=1)
    return np.abs(amp)**2


def fit_slope(ks, vals, lo=20, hi=180):
    sel = (ks >= lo) & (ks <= hi) & (vals > 1e-14)
    if sel.sum() < 3:
        return float("nan")
    return float(np.polyfit(np.log(ks[sel]), np.log(vals[sel]), 1)[0])

ks = np.arange(1, 401)
rho0 = 0.9998
teacher_lam = rho0 * np.exp(1j * omega)
R_teacher = return_proxy(teacher_lam, q, ks)


def select_frequency(M):
    return np.linspace(0, D - 1, M).round().astype(int)


def select_random(M):
    return np.sort(rng.choice(D, M, replace=False))


def select_magnitude(M):
    # Equal modal magnitudes make magnitude pruning unable to distinguish modes.
    return np.arange(M)


def relative_log_error(a, b):
    return float(np.sqrt(np.mean((np.log(a + 1e-16) - np.log(b + 1e-16)) ** 2)))

# Compression comparison. Selected modes inherit teacher weights and are renormalized.
compression = []
for M in [8, 16, 32, 64]:
    for name, selector in [("frequency", select_frequency), ("random", select_random), ("magnitude", select_magnitude)]:
        idx = selector(M)
        R = return_proxy(teacher_lam[idx], q[idx], ks)
        compression.append({
            "M": M, "method": name,
            "log_rmse": relative_log_error(R_teacher, R),
            "slope": fit_slope(ks, R),
            "teacher_slope": fit_slope(ks, R_teacher),
            "slope_relative_error": abs(fit_slope(ks, R) - fit_slope(ks, R_teacher)) / max(abs(fit_slope(ks, R_teacher)), 1e-12),
        })

# Prediction 1: stable radius gives an exponential envelope with time constant
# tau = -1/(2 log rho), because R(k) contains rho^(2k).
radii = [0.90, 0.95, 0.99, 0.999]
stability = []
for rho in radii:
    R = return_proxy(rho * np.exp(1j * omega[:32]), q[:32], ks)
    # Fit log envelope after averaging oscillations by a robust upper quantile in bins.
    bins = np.array_split(np.arange(len(ks)), 20)
    kb = np.array([np.mean(ks[b]) for b in bins])
    rb = np.array([np.quantile(R[b], 0.8) for b in bins])
    observed_slope = float(np.polyfit(kb, np.log(rb + 1e-30), 1)[0])
    predicted_slope = 2*np.log(rho)
    stability.append({"rho": rho, "predicted_log_slope": predicted_slope,
                      "observed_log_slope": observed_slope,
                      "relative_error": abs(observed_slope-predicted_slope)/abs(predicted_slope)})

# Prediction 2: at rho=1 there is no exponential decay; rho>1 grows at rate 2 log rho.
boundary = []
for rho in [0.999, 1.0, 1.001, 1.01]:
    R = return_proxy(rho * np.exp(1j * omega[:1]), np.array([1.0]), ks)
    observed = float(np.polyfit(ks, np.log(R), 1)[0])
    predicted = 2*np.log(rho)
    boundary.append({"rho": rho, "predicted_log_slope": predicted, "observed_log_slope": observed,
                     "absolute_error": abs(observed-predicted)})

# Prediction 3: increasing mode count improves quadrature of the pair-difference spectrum.
scaling = []
for M in [4, 8, 16, 32, 64, 128]:
    idx = select_frequency(M)
    R = return_proxy(teacher_lam[idx], q[idx], ks)
    scaling.append({"M": M, "log_rmse": relative_log_error(R_teacher, R),
                    "slope_relative_error": abs(fit_slope(ks, R)-fit_slope(ks, R_teacher))/max(abs(fit_slope(ks, R_teacher)),1e-12)})

out = {"seed": SEED, "D": D, "ks": [1, 400], "teacher_slope": fit_slope(ks, R_teacher),
       "stability_sweep": stability, "boundary_sweep": boundary,
       "compression": compression, "mode_scaling": scaling}
Path("results.json").write_text(json.dumps(out, indent=2))
print(json.dumps(out, indent=2))