import json
from pathlib import Path
import numpy as np
from scipy.optimize import least_squares
from scipy.signal import periodogram

SEED = 1120
rng = np.random.default_rng(SEED)


def simulate(r, omega, sigma, B=128, T=3000, burn=500, seed=0):
    """Independent noisy trajectories of h[t+1]=r R(omega)h[t]+sigma eps."""
    g = np.random.default_rng(seed)
    R = np.array([[np.cos(omega), -np.sin(omega)],
                  [np.sin(omega),  np.cos(omega)]])
    h = g.normal(size=(B, 2))
    out = np.empty((B, T))
    for t in range(T + burn):
        h = h @ (r * R).T + sigma * g.normal(size=(B, 2))
        if t >= burn:
            out[:, t-burn] = h[:, 0]
    return out


def ensemble_corr(y, maxlag=80):
    # Average stationary autocovariance over trajectories, as in the proposed estimator.
    y = y - y.mean(axis=1, keepdims=True)
    c = np.array([(y[:, :y.shape[1]-k] * y[:, k:]).mean() for k in range(maxlag+1)])
    return c / max(c[0], 1e-12)


def fit_resonance(c, max_r=0.9999):
    k = np.arange(len(c), dtype=float)
    # A damped sinusoid has two real coefficients and the resonance radius/frequency.
    def residual(p):
        r, w, a, b = p
        return r**k * (a*np.cos(w*k) + b*np.sin(w*k)) - c
    best = None
    for w0 in np.linspace(.12, 1.25, 8):
        x = least_squares(residual, [0.9, w0, 1.0, 0.0],
                          bounds=([.01, .01, -3, -3], [max_r, 1.5, 3, 3]),
                          max_nfev=1000)
        if best is None or np.sum(x.fun*x.fun) < best[0]:
            best = (np.sum(x.fun*x.fun), x.x)
    return best[1]


def spectrum_peak(y, omega_expected):
    # Average periodograms; ignore DC and select the nearest local spectral maximum.
    f, p = periodogram(y, axis=1, detrend='constant')
    p = p.mean(axis=0)
    target = omega_expected / (2*np.pi)
    band = (f > max(1/len(y[0]), target-.12)) & (f < target+.12)
    ix = np.where(band)[0][np.argmax(p[band])]
    return 2*np.pi*f[ix], p[ix], f, p


def jacobian_lyapunov(r, omega, T=10000):
    # Exact tangent update: J=rR, hence every normalized tangent grows by r.
    # Written as a numerical normalized iteration to verify the estimator itself.
    R = np.array([[np.cos(omega), -np.sin(omega)],
                  [np.sin(omega),  np.cos(omega)]])
    J = r * R
    v = np.array([1., .37]); v /= np.linalg.norm(v)
    logs=[]
    for _ in range(T):
        q = J @ v
        n = np.linalg.norm(q); logs.append(np.log(n)); v=q/n
    return float(np.mean(logs))


def run():
    outdir = Path('results'); outdir.mkdir(exist_ok=True)
    omega = .42
    # Prediction 1: tau=-1/log(r), estimated correlation envelope follows it.
    radius_rows=[]
    for i, r in enumerate([.70, .80, .90, .95, .98]):
        y = simulate(r, omega, .12, B=128, T=3500, seed=100+i)
        c = ensemble_corr(y, 100)
        fitr, fitw, a, b = fit_resonance(c)
        tau_pred = -1/np.log(r)
        tau_fit = -1/np.log(fitr)
        lam = jacobian_lyapunov(r, omega)
        radius_rows.append(dict(r=r, fitted_r=float(fitr), tau_pred=float(tau_pred),
                                tau_fit=float(tau_fit), lyapunov=float(lam),
                                rel_tau_error=float(abs(tau_fit-tau_pred)/tau_pred)))

    # Prediction 2: peak frequency is omega/(2pi), independently of radius/noise.
    freq_rows=[]
    for i, w in enumerate([.25, .42, .70, 1.00]):
        y = simulate(.90, w, .12, B=128, T=3500, seed=200+i)
        c = ensemble_corr(y, 80)
        fitr, fitw, a, b = fit_resonance(c)
        peak, _, _, _ = spectrum_peak(y, w)
        freq_rows.append(dict(omega=w, fitted_omega=float(fitw), psd_peak=float(peak),
                              predicted=float(w), peak_abs_error=float(abs(peak-w)),
                              relative_error=float(abs(peak-w)/w)))

    # Prediction 3: additive noise changes variance but not resonance or Lyapunov exponent.
    noise_rows=[]
    for i, sig in enumerate([.03, .12, .40, .80]):
        y = simulate(.93, omega, sig, B=128, T=3500, seed=300+i)
        c=ensemble_corr(y, 80)
        fitr, fitw, a, b=fit_resonance(c)
        noise_rows.append(dict(sigma=sig, fitted_r=float(fitr), fitted_omega=float(fitw),
                               lyapunov=float(jacobian_lyapunov(.93,omega)),
                               variance=float(y.var())))

    # Controller: resonance penalty clips an estimated radius above r_max.
    # Compare the uncontrolled high-Q oscillator with the controlled one.
    r_raw=.975; r_max=.90
    yc=simulate(r_raw, omega, .12, B=128, T=4000, seed=500)
    # In this linear MVP, applying the penalty is equivalent to the prescribed gain control.
    yctrl=simulate(min(r_raw,r_max), omega, .12, B=128, T=4000, seed=501)
    def summarize(y, r):
        c=ensemble_corr(y,100); fr,fw,_,_=fit_resonance(c); peak,powr,_,_=spectrum_peak(y,omega)
        return dict(fitted_r=float(fr), fitted_omega=float(fw), decay_time=float(-1/np.log(fr)),
                    peak_frequency=float(peak), peak_power=float(powr), long_lag_corr=float(abs(c[80])),
                    lyapunov=float(jacobian_lyapunov(r,omega)))
    control = {'uncontrolled':summarize(yc,r_raw), 'resonance_control':summarize(yctrl,r_max),
               'target_radius':r_max}

    decay_pass = all(z['rel_tau_error'] < .05 for z in radius_rows)
    freq_pass = all(z['relative_error'] < .10 for z in freq_rows)
    noise_r_span = max(z['fitted_r'] for z in noise_rows) - min(z['fitted_r'] for z in noise_rows)
    noise_w_span = max(z['fitted_omega'] for z in noise_rows) - min(z['fitted_omega'] for z in noise_rows)
    noise_pass = noise_r_span < .01 and noise_w_span < .01
    control_gain = control['uncontrolled']['peak_power'] / control['resonance_control']['peak_power']
    report={'seed':SEED, 'model':'h[t+1]=r R(omega) h[t]+sigma Normal(0,I), y=h[0]',
            'predictions':{
              'decay_time':'tau=-1/log(r); fitted resonance radius predicts correlation decay',
              'frequency':'PSD peak and fitted omega equal injected omega',
              'independence':'additive noise changes variance, not J-based Lambda or resonance'},
            'radius_sweep':radius_rows, 'frequency_sweep':freq_rows,
            'noise_sweep':noise_rows, 'control':control,
            'quantitative_checks': {
              'decay_time_max_relative_error': max(z['rel_tau_error'] for z in radius_rows),
              'decay_time_pass_under_5pct': decay_pass,
              'frequency_max_relative_error': max(z['relative_error'] for z in freq_rows),
              'frequency_pass_under_10pct': freq_pass,
              'noise_fitted_radius_span': noise_r_span,
              'noise_fitted_frequency_span': noise_w_span,
              'noise_invariance_pass_under_0.01': noise_pass,
              'baseline_peak_power': control['uncontrolled']['peak_power'],
              'idea_peak_power': control['resonance_control']['peak_power'],
              'peak_power_reduction_factor': control_gain,
              'baseline_long_lag_abs_corr': control['uncontrolled']['long_lag_corr'],
              'idea_long_lag_abs_corr': control['resonance_control']['long_lag_corr']
            }}
    Path('results/report.json').write_text(json.dumps(report, indent=2))
    # Compact CSV-like human-readable summary for quick inspection.
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    run()