import json
import math
import random
from pathlib import Path

import numpy as np
import torch

SEED = 2466
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
A = 1.15
T = 60


def rollout(k, e0=1.0, a=A, T=T):
    rng = np.random.default_rng(SEED)
    v = rng.normal(0.0, 0.35, size=T)
    q, z = 0.0, e0
    errors = [z - q]
    for t in range(T):
        u = k * (q - z)
        q = a * q + v[t]
        z = a * z + u + v[t]
        errors.append(z - q)
    return np.asarray(errors)


def fit_slope(errors, burn=3):
    # Fit only the numerically resolved exponential segment. Exact synchronization
    # can otherwise turn roundoff into a spurious positive slope.
    vals = np.abs(errors[burn:])
    peak = float(np.max(vals))
    keep = vals > max(peak * 1e-12, 1e-14)
    vals = vals[keep]
    x = np.arange(len(errors))[burn:][keep].astype(float)
    if len(vals) < 2:
        return 0.0
    return float(np.polyfit(x, np.log(vals), 1)[0])


def analytic_lambda(k, a=A):
    m = abs(a - k)
    return -math.inf if m == 0 else math.log(m)


def numerical_checks():
    # Prediction 1: transverse stability boundary is k*=a-1 for positive gains.
    gains = np.linspace(-0.8, 2.8, 37)
    rows = []
    for k in gains:
        errors = rollout(float(k))
        measured = fit_slope(errors)
        predicted = analytic_lambda(float(k))
        rows.append({"k": float(k), "predicted_lambda": predicted,
                     "measured_slope": measured,
                     "stable_predicted": bool(abs(A-k) < 1),
                     "stable_observed": bool(measured < 0)})

    # Prediction 2: measured log-residual slope equals transverse exponent.
    slope_errors = [abs(r["measured_slope"] - r["predicted_lambda"])
                    for r in rows if np.isfinite(r["predicted_lambda"]) and abs(A-r["k"]) > 1e-6]
    max_slope_error = max(slope_errors)

    # Prediction 3: close to the boundary, decay time scales as 1/|lambda|.
    near_gains = [0.35, 0.45, 0.55, 0.65, 0.75, 0.85, 0.95, 1.05, 1.15, 1.25]
    scaling = []
    for k in near_gains:
        lam = analytic_lambda(k)
        slope = fit_slope(rollout(k))
        scaling.append({"k": k, "abs_lambda": abs(lam),
                        "measured_abs_slope": abs(slope),
                        "time_constant_pred": 1.0/abs(lam),
                        "time_constant_obs": (1.0/abs(slope) if abs(slope) > 1e-15 else float("inf"))})

    # Prediction 3: the convergence margin is linear in coupling gain,
    # lambda(k)=log|a-k|, with slope -1/(a-k) locally.
    scaling_gains = [0.20, 0.30, 0.40, 0.50]
    gain_scaling = [{"k": k, "predicted_multiplier": A-k,
                     "observed_multiplier": float(abs(rollout(k)[1] / rollout(k)[0]))}
                    for k in scaling_gains]

    # Locate the observed sign transition by linear interpolation between sweep points.
    crossing = None
    for x, y in zip(rows[:-1], rows[1:]):
        if x["measured_slope"] * y["measured_slope"] <= 0:
            crossing = x["k"] + (0.0-x["measured_slope"]) * (y["k"]-x["k"]) / (y["measured_slope"]-x["measured_slope"])
            break
    stable_agreement = sum(r["stable_predicted"] == r["stable_observed"] for r in rows) / len(rows)
    return {"analytic_boundary": A - 1.0, "observed_boundary": crossing,
            "boundary_error": abs(crossing-(A-1.0)) if crossing is not None else None,
            "classification_agreement": stable_agreement,
            "max_abs_slope_error": max_slope_error,
            "slope_rows": rows,
            "scaling_rows": scaling, "gain_scaling": gain_scaling}


def train_controller(sync_weight, steps=500, batch=64, horizon=25):
    # Same task for both methods: make plant track reference under shared forcing.
    torch.manual_seed(SEED)
    k = torch.nn.Parameter(torch.tensor(0.0))
    opt = torch.optim.Adam([k], lr=0.03)
    rng = torch.Generator().manual_seed(SEED)
    losses = []
    for _ in range(steps):
        v = torch.randn(batch, horizon, generator=rng) * 0.35
        q = torch.zeros(batch)
        z = torch.randn(batch, generator=rng) * 1.5
        total = torch.zeros(())
        for t in range(horizon):
            u = k * (q-z)
            q = A*q + v[:, t]
            z = A*z + u + v[:, t]
            total = total + (z-q).square().mean()
        # task and synchrony coincide here; sync_weight tests explicit residual shaping.
        loss = total / horizon * (1.0 + sync_weight)
        opt.zero_grad(); loss.backward(); opt.step()
        losses.append(float(loss.detach()))
    with torch.no_grad():
        test = []
        for e0 in [-2.0, -1.0, 1.0, 2.0]:
            test.append(float(np.mean(rollout(float(k), e0=e0)[-10:]**2)))
    return {"final_k": float(k.detach()), "train_loss_last": losses[-1],
            "test_tail_mse": float(np.mean(test)), "loss_curve": losses}


def main():
    checks = numerical_checks()
    baseline = train_controller(0.0)
    idea = train_controller(2.0)
    out = {"seed": SEED, "a": A, "horizon": T, "checks": checks,
           "training": {"baseline_task_only": baseline,
                         "idea_task_plus_synchrony": idea}}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps({"analytic_boundary": checks["analytic_boundary"],
                      "observed_boundary": checks["observed_boundary"],
                      "boundary_error": checks["boundary_error"],
                      "classification_agreement": checks["classification_agreement"],
                      "max_abs_slope_error": checks["max_abs_slope_error"],
                      "baseline": {k:v for k,v in baseline.items() if k != "loss_curve"},
                      "idea": {k:v for k,v in idea.items() if k != "loss_curve"}}, indent=2))


if __name__ == "__main__":
    main()