import json
import numpy as np
from scipy.optimize import root

SEED = 1128
A, W0, H = 0.04, 1.0, 0.01

# Modal oscillator: z' = (-A + i W0) z - q z(t-tau).
def char_root(q, tau, guess=None):
    if guess is None:
        guess = complex(-A, W0)
    def fun(v):
        s = v[0] + 1j*v[1]
        f = s + A - 1j*W0 + q*np.exp(-s*tau)
        return [f.real, f.imag]
    sol = root(fun, [guess.real, guess.imag], method='hybr')
    if not sol.success:
        raise RuntimeError(sol.message)
    return complex(sol.x[0], sol.x[1])


def compensated_root(q, tau, guess=None):
    """Root for the implemented predictor: x[t-d]+d*(x[t-d]-x[t-d-1])."""
    if guess is None:
        guess = complex(-A, W0)
    d = int(round(tau / H))
    def fun(v):
        s = v[0] + 1j*v[1]
        # predictor transfer factor for exp(s t), with d samples of delay
        pred = np.exp(-s*d*H) * (1.0 + d*(1.0-np.exp(-s*H)))
        f = s + A - 1j*W0 + q*pred
        return [f.real, f.imag]
    sol = root(fun, [guess.real, guess.imag], method='hybr')
    if not sol.success:
        raise RuntimeError(sol.message)
    return complex(sol.x[0], sol.x[1])


def simulate_mode(q, tau, compensated=False, T=240.0):
    n, d = int(T/H), int(round(tau/H))
    t = np.arange(n+1)*H
    z = np.exp((-A+1j*W0)*(t-tau-2*H)).astype(complex)
    for k in range(1, n+1):
        idx = max(0, k-d)
        src = z[idx]
        if compensated and d > 0:
            src = src + d*(src-z[max(0, idx-1)])
        z[k] = z[k-1] + H*((-A+1j*W0)*z[k-1]-q*src)
    return z


def dominant_freq(z, discard=.45):
    """Phase-slope estimator avoids FFT-bin quantization for a single analytic mode."""
    x = z[int(len(z)*discard):]
    phase = np.unwrap(np.angle(x))
    slope = np.polyfit(np.arange(len(x))*H, phase, 1)[0]
    return float(slope)


def main():
    # Prediction 1: delay phase has slope -omega0.
    taus = np.array([0., .1, .2, .3, .4, .5])
    phases = np.unwrap(np.angle(np.exp(-1j*W0*taus)))
    phase_slope = float(np.polyfit(taus, phases, 1)[0])

    # Prediction 2: lambda=0 is delay invariant; a nonzero mode changes.
    z0, z0d = simulate_mode(0., .4), simulate_mode(0., .4, True)
    zero_err = float(np.max(abs(z0-z0d)))
    zbase, zdel = simulate_mode(.06, 0.), simulate_mode(.06, .4)
    nonzero_change = float(np.sqrt(np.mean(abs(zdel-zbase)**2)) /
                           np.sqrt(np.mean(abs(zbase)**2)))

    # Prediction 3: first-order frequency shift is linear in q.
    tau = .30
    qs = np.array([.005, .01, .02, .04, .06])
    roots = np.array([char_root(float(q), tau) for q in qs])
    shifts = roots.imag-W0
    fit_slope = float(np.polyfit(qs, shifts, 1)[0])
    pred_slope = float(np.exp(A*tau)*np.sin(W0*tau))

    # Four branches on a ring: Laplacian eigenvalues 0,2,2,4.
    lambdas, gamma, delay = [0.,2.,2.,4.], .03, .40
    rows = []
    for lam in lambdas:
        q = gamma*lam
        zd, zc = simulate_mode(q, delay), simulate_mode(q, delay, True)
        rb, rc = char_root(q, delay), compensated_root(q, delay)
        rows.append({'lambda':lam, 'baseline_freq':dominant_freq(zd),
            'compensated_freq':dominant_freq(zc),
            'baseline_root_freq':float(rb.imag), 'compensated_root_freq':float(rc.imag),
            'baseline_freq_error':abs(dominant_freq(zd)-W0),
            'compensated_freq_error':abs(dominant_freq(zc)-W0)})
    nz = [r for r in rows if r['lambda'] > 0]
    base_mean = float(np.mean([r['baseline_freq_error'] for r in nz]))
    comp_mean = float(np.mean([r['compensated_freq_error'] for r in nz]))
    result = {
      'seed':SEED, 'h':H, 'omega0':W0, 'damping':A,
      'prediction_1_phase': {'taus':taus.tolist(), 'observed_phase_slope':phase_slope,
                             'predicted_phase_slope':-W0},
      'prediction_2_zero_mode': {'max_abs_error_lambda0':zero_err,
                                 'relative_change_nonzero_q006_tau04':nonzero_change},
      'prediction_3_frequency_scaling': {'tau':tau, 'q_values':qs.tolist(),
        'measured_root_shifts':shifts.tolist(), 'fitted_shift_per_q':fit_slope,
        'predicted_first_order_shift_per_q':pred_slope,
        'relative_slope_error':abs(fit_slope-pred_slope)/abs(pred_slope)},
      'mini_experiment': {'gamma':gamma, 'delay':delay, 'rows':rows,
        'mean_nonzero_frequency_error_baseline':base_mean,
        'mean_nonzero_frequency_error_compensated':comp_mean,
        'relative_improvement':(base_mean-comp_mean)/base_mean},
      'checks': {'phase_ok':bool(abs(phase_slope+W0)<1e-10),
        'zero_mode_ok':bool(zero_err<1e-10 and nonzero_change>1e-3),
        'scaling_ok':bool(abs(fit_slope-pred_slope)/abs(pred_slope)<.08),
        'compensation_better':bool(comp_mean<base_mean)}
    }
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__ == '__main__': main()