import json
from pathlib import Path
import numpy as np
from scipy.linalg import expm

SEED = 2273


def oscillator_matrix(damping=0.08, omega=1.7):
    return np.array([[damping, -omega], [omega, damping]], dtype=float)


def delay_row(A, c, t):
    return np.asarray(c @ expm(-A * float(t))).reshape(-1)


def sigma_min(O):
    return float(np.linalg.svd(O, compute_uv=False)[-1]) if O.shape[0] >= O.shape[1] else 0.0


def greedy_delays(A, c, candidates, budget, delta=0.0, initial=(0.0,)):
    selected = list(map(float, initial))
    O = np.stack([delay_row(A, c, t) for t in selected])
    margins = [sigma_min(O)]
    gains = []
    for _ in range(max(0, budget - len(selected))):
        remaining = [float(t) for t in candidates if all(abs(t-u) > 1e-10 for u in selected)]
        if not remaining:
            break
        scored = [(sigma_min(np.vstack((O, delay_row(A, c, t)))), t) for t in remaining]
        best, t = max(scored, key=lambda x: x[0])
        gain = best - margins[-1]
        if gain < delta:
            break
        selected.append(t)
        O = np.vstack((O, delay_row(A, c, t)))
        margins.append(float(best)); gains.append(float(gain))
    return np.array(selected), O, np.array(margins), np.array(gains)


def uniform_delays(tmax, count):
    return np.linspace(0.0, float(tmax), count)


def reconstruction_mse(O, noise_std, trials=12000, seed=991):
    rng = np.random.default_rng(seed)
    x = rng.normal(size=(O.shape[1], trials))
    y = O @ x + noise_std * rng.normal(size=(O.shape[0], trials))
    xhat = np.linalg.pinv(O) @ y
    return float(np.mean((xhat-x)**2))


def main():
    c = np.array([1.0, 0.0])
    candidates = np.unique(np.r_[np.linspace(0.0, 4.0, 161), np.geomspace(0.01, 12.0, 100)])
    rows = []
    A = oscillator_matrix(); selected, O, margins, _ = greedy_delays(A, c, candidates, 8)
    rows.append({'prediction':'row append monotonicity for every greedy step', 'observed_min_increment':float(np.min(np.diff(margins))), 'predicted':'>= 0', 'confirmed':bool(np.all(np.diff(margins) >= -1e-12))})

    # Frequency sweep: informative quadrature delay is pi/(2 omega).
    freq_sweep = []
    for omega in (0.8, 1.2, 1.7, 2.4, 3.2):
        A = oscillator_matrix(omega=omega)
        ts, _, _, _ = greedy_delays(A, c, candidates, 2)
        predicted = np.pi/(2*omega)
        observed = float(ts[1])
        freq_sweep.append({'omega':omega, 'predicted_quarter':predicted, 'observed_first_delay':observed, 'absolute_error':abs(observed-predicted)})
    rows.append({'prediction':'first greedy delay follows quarter-period pi/(2 omega)', 'sweep':freq_sweep, 'max_absolute_error':max(x['absolute_error'] for x in freq_sweep), 'confirmed':max(x['absolute_error'] for x in freq_sweep) < 0.15})

    # Noise sweep: fixed linear inverse predicts MSE proportional to noise variance.
    A = oscillator_matrix(); _, O4, margins4, _ = greedy_delays(A, c, candidates, 4)
    noise_sweep = []
    for noise in (0.005, 0.01, 0.02, 0.04, 0.08):
        mse = reconstruction_mse(O4, noise)
        noise_sweep.append({'noise_std':noise, 'mse':mse, 'mse_over_noise2':mse/(noise*noise)})
    ratios = np.array([x['mse_over_noise2'] for x in noise_sweep])
    rows.append({'prediction':'reconstruction MSE scales as noise_std^2 / sigma_min^2', 'sigma_min':float(margins4[-1]), 'sweep':noise_sweep, 'ratio_cv':float(np.std(ratios)/np.mean(ratios)), 'confirmed':bool(np.std(ratios)/np.mean(ratios) < 0.04)})

    # Margin sweep and inverse-margin relation across retained tap counts.
    checks=[]
    for budget in range(2,9):
        _, OO, mm, _ = greedy_delays(A, c, candidates, budget)
        mse = reconstruction_mse(OO, 0.03)
        checks.append((float(mm[-1]), mse, float(1/(mm[-1]**2 + 1e-12))))
    margins3=np.array([x[0] for x in checks]); mses=np.array([x[1] for x in checks]); inv=np.array([x[2] for x in checks])
    corr=float(np.corrcoef(mses, inv)[0,1])
    rows.append({'prediction':'error decreases with inverse stable margin squared', 'budgets':list(range(2,9)), 'sigma_min':margins3.tolist(), 'mse':mses.tolist(), 'inv_sigma2':inv.tolist(), 'correlation':corr, 'confirmed':corr > 0.85})

    compare=[]
    for budget in (4, 8, 16):
        gts, gO, gm, _ = greedy_delays(A,c,candidates,budget)
        uts=uniform_delays(4.0,budget); uO=np.stack([delay_row(A,c,t) for t in uts])
        compare.append({'budget':budget, 'greedy_delays':gts.tolist(), 'uniform_delays':uts.tolist(), 'greedy_sigma_min':sigma_min(gO), 'uniform_sigma_min':sigma_min(uO), 'greedy_mse':reconstruction_mse(gO,.03), 'uniform_mse':reconstruction_mse(uO,.03)})
    result={'seed':SEED,'system':{'damping':.08,'omega':1.7},'checks':rows,'comparison':compare}
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == '__main__': main()