import json
import numpy as np
from hidden_diffusion_monitor import (analytic_spectrum, paper_sigma, simulate,
                                      forward_reverse_action_gap, autocorr, expm)


def run():
    rng = np.random.default_rng(20250308)
    A = np.array([[1., 1.], [-1., 0.]])  # A_yy=0, stable, coupled
    Dx = Dm = 1.; dt = .03; n = 120000
    dcs = [-.75, 0., .75]
    omega = np.logspace(-2, 2, 250)
    spectra = []; acs = []; rows = []
    for dc in dcs:
        D = np.array([[Dx, dc], [dc, Dm]])
        spectra.append(analytic_spectrum(omega, A, D))
        path, C, F, Q = simulate(A, D, dt, n, rng)
        ac = autocorr(path, 80)[1:]
        acs.append(ac)
        sigma = paper_sigma(1., 1., Dx, Dm, dc)
        gap = forward_reverse_action_gap(path, F, Q, dt)
        exact = np.array([(expm(-A*t) @ C)[0, 0]/C[0, 0]
                          for t in np.arange(1, 81)*dt])
        rows.append({'D_xy': dc, 'spectrum_rel_to_baseline': float(
            np.max(np.abs(spectra[-1]-spectra[0]) / np.maximum(spectra[0], 1e-12))),
            'sigma_formula': float(sigma), 'monitor_gap': float(gap),
            'autocorrelation_max_error': float(np.max(np.abs(ac-exact)))})
    # Baseline has only observed x: its distinguishability is spectral/AC variation.
    baseline_pairwise_ac = max(float(np.max(np.abs(acs[i]-acs[j])))
                               for i in range(3) for j in range(i))
    baseline_pairwise_spectrum = max(float(np.max(np.abs(spectra[i]-spectra[j]) /
                                                  np.maximum(spectra[j], 1e-12)))
                                     for i in range(3) for j in range(i))
    monitor_gap_spread = max(r['monitor_gap'] for r in rows)-min(r['monitor_gap'] for r in rows)
    formula_spread = max(r['sigma_formula'] for r in rows)-min(r['sigma_formula'] for r in rows)
    result = {'rows': rows, 'baseline_observed_AC_pairwise_max': baseline_pairwise_ac,
              'baseline_observed_spectrum_pairwise_max': baseline_pairwise_spectrum,
              'monitor_gap_spread': monitor_gap_spread, 'formula_sigma_spread': formula_spread}
    print(json.dumps(result, indent=2))
    with open('results.json', 'w') as f: json.dump(result, f, indent=2)

if __name__ == '__main__': run()