import json
from pathlib import Path
import numpy as np
from carrier_probe import H, DT, Q_TRUE, carrier, gramians, rollout, fit_q


def main():
    rng = np.random.default_rng(123)
    a = 0.16
    delta = 1e-8
    alpha = 0.1
    candidates = []
    for w in [0.5, 1.0, 2.0, 4.0, 8.0]:
        candidates.append((f"sin_w{w}", carrier("sin", w)))
    t = np.arange(H) * DT
    candidates.append(("binary_w2", carrier("binary", 2.0)))
    # A short orthogonal-ish pulse, normalized to unit magnitude.
    pulse = np.zeros(H); pulse[H//3: H//3 + H//10] = 1.0
    candidates.append(("short_pulse", pulse))
    rows = []
    for name, m in candidates:
        Wo, Wr, eo, er = gramians(a, m)
        score = np.linalg.slogdet(Wo + delta*np.eye(2))[1] + alpha*np.linalg.slogdet(Wr + delta*np.eye(2))[1]
        rows.append({"candidate": name, "score": float(score),
                     "lambda_min_Wo": float(eo[0]), "lambda_min_Wr": float(er[0])})
    rows.sort(key=lambda x: x["score"], reverse=True)

    # Recovery curve under identical observation noise scale.
    recovery = []
    for aa in [0.0, 0.01, 0.02, 0.04, 0.08, 0.16, 0.32]:
        m = carrier("sin", 2.0)
        obs = rollout(aa, m)[:, 0] + rng.normal(0, 0.01, H+1)
        qhat = fit_q(obs, m, aa)
        recovery.append({"a": aa, "q_hat": qhat, "abs_error": abs(qhat-Q_TRUE)})
    out = {"amplitude": a, "alpha": alpha, "candidates_by_score": rows,
           "recovery_curve": recovery,
           "selected": rows[0]["candidate"]}
    Path("candidate_results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()