import json, math, random
from pathlib import Path
import numpy as np

# Graph-certified switching SSM MVP.  The one-node graph is path-complete
# when it contains one self-edge for every mode; multiple-node certificates
# are supported by the loss implementation below.

def edge_loss(a, rho, p=1.0, tau=0.02):
    """Mean softplus edge penalty for V(z)=p z^2, evaluated at z=1."""
    # softplus argument is the normalized violation; this is the exact
    # expression apart from the harmless tau smoothing.
    e = p * (np.asarray(a) ** 2 - rho)
    return np.logaddexp(0.0, e / tau) * tau

def core_verification():
    rho = 0.81
    boundary = math.sqrt(rho)
    # Prediction 1: violation switches at |a|=sqrt(rho).
    amps = np.linspace(0.0, 1.2, 1201)
    hard_violation = amps**2 > rho
    transition = amps[np.flatnonzero(hard_violation)[0]]
    # Prediction 2: arbitrary switching norm is product; worst sequence is max amp^T.
    modes = np.array([0.55, 0.72, 0.89, 1.03])
    T = 80
    predicted_log_growth = T * math.log(np.max(np.abs(modes)))
    rng = np.random.default_rng(7)
    seq = rng.integers(0, len(modes), size=T)
    observed_log_norm = np.sum(np.log(np.abs(modes[seq])))
    worst_log_norm = np.sum(np.log(np.max(np.abs(modes))) * np.ones(T))
    # Prediction 3: unsmoothed violation is exactly quadratic in amplitude excess.
    excess = np.array([0.01, 0.03, 0.07, 0.12])
    exact = (boundary + excess)**2 - rho
    slopes = np.diff(exact) / np.diff(excess)
    predicted_slope_at_boundary = 2 * boundary
    # Confirm direct rollout agrees with product formula.
    z = 1.0
    for a in modes[np.argmax(np.abs(modes))] * np.ones(T):
        z *= a
    return {
        "rho": rho, "predicted_boundary": boundary,
        "observed_boundary_grid": float(transition),
        "boundary_abs_error": float(abs(transition-boundary)),
        "modes": modes.tolist(), "T": T,
        "predicted_worst_log_norm": float(predicted_log_growth),
        "observed_worst_log_norm": float(math.log(abs(z))),
        "random_sequence_log_norm": float(observed_log_norm),
        "quadratic_excess": excess.tolist(), "exact_violation": exact.tolist(),
        "local_slopes": slopes.tolist(),
        "predicted_boundary_slope": predicted_slope_at_boundary,
        "slope_relative_error_last": float(abs(slopes[-1]-predicted_slope_at_boundary)/predicted_slope_at_boundary),
        "all_predictions_confirmed": bool(abs(transition-boundary)<0.002 and abs(math.log(abs(z))-predicted_log_growth)<1e-9 and abs(slopes[-1]-predicted_slope_at_boundary)/predicted_slope_at_boundary<0.15)
    }

def train_tiny(seed, regularized, steps=450):
    # Torch is optional; numpy fallback still produces a useful comparison.
    try:
        import torch
        torch.manual_seed(seed)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        try:
            torch.tensor([0.], device=device)
        except Exception:
            device = "cpu"
        dtype = torch.float32
        # A switching scalar recurrence learns a discounted accumulator.
        M, batch, T = 4, 64, 20
        rho, tau, lam = .81, .03, .25
        raw_a = torch.nn.Parameter(torch.tensor([0.15, 0.25, 0.35, 0.45], device=device))
        raw_b = torch.nn.Parameter(torch.zeros(M, device=device))
        out_w = torch.nn.Parameter(torch.tensor(1., device=device))
        opt = torch.optim.Adam([raw_a, raw_b, out_w], lr=.025)
        rng = np.random.default_rng(seed)
        for _ in range(steps):
            x = torch.tensor(rng.normal(size=(batch,T)), dtype=dtype, device=device)
            # random modes are used during training and all modes in penalty
            modes = torch.tensor(rng.integers(0,M,size=(batch,T)), dtype=torch.long, device=device)
            z = torch.zeros(batch, device=device)
            target = torch.zeros(batch, device=device)
            discount = 1.
            for t in range(T):
                target = target + discount*x[:,t]
                discount *= .8
                ai, bi = raw_a[modes[:,t]], raw_b[modes[:,t]]
                z = ai*z + bi*x[:,t]
            pred = out_w*z
            task = ((pred-target)**2).mean()
            # V(z)=z^2, one graph node, every mode self-edge. Use sampled z.
            zs = torch.randn(128, device=device)
            ai = raw_a[None,:]
            violations = ai**2 * zs[:,None]**2 - rho*zs[:,None]**2
            cert = torch.nn.functional.softplus(violations/tau).mean()*tau
            loss = task + (lam*cert if regularized else 0.)
            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_([raw_a,raw_b,out_w], 10.); opt.step()
        with torch.no_grad():
            a = raw_a.detach().cpu().numpy()
            # Worst arbitrary switching rollout, no inputs, initialized at one.
            worst = float(1000*np.log(max(np.max(np.abs(a)), 1e-30)))
            violation_rate = float(np.mean(a*a > rho))
            final_task = float(task.detach().cpu())
        return {"task_mse":final_task, "max_abs_a":float(np.max(np.abs(a))), "edge_violation_rate":violation_rate, "norm_after_1000_worst":worst, "device":device}
    except Exception as exc:
        return {"error": repr(exc)}

def main():
    result = {"verification": core_verification(), "mini_experiment": {
        "baseline_unregularized": train_tiny(11, False),
        "graph_certified": train_tiny(11, True)}}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == "__main__": main()