import json, math, random
import numpy as np
import torch
from torch import nn

SEED = 7
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cuda":
        torch.cuda.manual_seed_all(SEED)
except Exception:
    device = "cpu"


def tangent_project(x, q, lo=-1.0, hi=1.0):
    """Projection of q onto the tangent cone of [lo,hi]^d at x."""
    p = q.clone()
    p = torch.where(x <= lo, torch.clamp(p, min=0), p)
    p = torch.where(x >= hi, torch.clamp(p, max=0), p)
    return p


def numerical_math_check():
    # Same input-dependent forcing cancels in the difference.  The drift is -Mx,
    # M positive semidefinite, and the projected Euler flow stays in the box.
    rng = np.random.default_rng(SEED)
    d, trials, steps = 5, 1000, 2000
    R = rng.normal(size=(d, d)); M = R.T @ R / d + .15*np.eye(d)
    ratios = []; max_box = 0.; violations = 0
    h = 0.01
    for _ in range(trials):
        x = rng.uniform(-1, 1, d); z = rng.uniform(-1, 1, d)
        # common bounded forcing represents identical v(t)
        force = rng.normal(size=d) * .2
        initial = np.linalg.norm(x-z)
        for _ in range(steps):
            qx = -M @ x + force; qz = -M @ z + force
            px = qx.copy(); pz = qz.copy()
            px[x <= -1] = np.maximum(px[x <= -1], 0); px[x >= 1] = np.minimum(px[x >= 1], 0)
            pz[z <= -1] = np.maximum(pz[z <= -1], 0); pz[z >= 1] = np.minimum(pz[z >= 1], 0)
            x = np.clip(x + h*px, -1, 1); z = np.clip(z + h*pz, -1, 1)
            max_box = max(max_box, float(np.max(np.abs(x))), float(np.max(np.abs(z))))
        final = np.linalg.norm(x-z)
        ratios.append(final / (initial + 1e-12))
        violations += int(final > initial + 2e-7)
    ratios = np.asarray(ratios)
    return {"trials": trials, "max_distance_ratio": float(ratios.max()),
            "mean_distance_ratio": float(ratios.mean()), "violations": violations,
            "max_abs_state": max_box, "step": h, "duration": steps*h}


class VanillaResidual(nn.Module):
    def __init__(self, inp, hid):
        super().__init__(); self.hid = hid
        self.wx = nn.Linear(hid, hid, bias=False); self.wv = nn.Linear(inp, hid)
        nn.init.normal_(self.wx.weight, std=.35/math.sqrt(hid))
    def forward(self, seq):
        x = torch.zeros(seq.size(0), self.hid, device=seq.device)
        for t in range(seq.size(1)):
            x = x + .08*torch.tanh(self.wx(x) + self.wv(seq[:,t]))
        return x

class ProjectedDissipative(nn.Module):
    def __init__(self, inp, hid):
        super().__init__(); self.hid=hid; self.h=.08
        self.L = nn.Parameter(torch.randn(hid, hid)*(.18/math.sqrt(hid)))
        self.B = nn.Linear(inp, hid); self.bias=nn.Parameter(torch.zeros(hid))
    def forward(self, seq):
        x=torch.zeros(seq.size(0), self.hid, device=seq.device)
        # positive lambda makes the drift strictly dissipative
        M=self.L.T@self.L + .12*torch.eye(self.hid, device=seq.device)
        for t in range(seq.size(1)):
            q=-(x@M.T) + self.B(seq[:,t]) + self.bias
            p=tangent_project(x,q)
            x=torch.clamp(x+self.h*p,-1,1)
        return x


def train_model(model, train, target, test, ytest, steps=300):
    model.to(device); opt=torch.optim.Adam(model.parameters(), lr=3e-3); lossfn=nn.BCEWithLogitsLoss()
    head=nn.Linear(model.hid,1).to(device); opt=torch.optim.Adam(list(model.parameters())+list(head.parameters()),lr=3e-3)
    xtr=train.to(device); y=target.to(device)
    losses=[]
    for k in range(steps):
        opt.zero_grad(); pred=head(model(xtr)).squeeze(-1); loss=lossfn(pred,y); loss.backward(); opt.step()
        if k in (0, steps-1): losses.append(float(loss.detach().cpu()))
    with torch.no_grad():
        acc=((torch.sigmoid(head(model(test.to(device))).squeeze(-1))>.5)==ytest.to(device)).float().mean().item()
        # response to a small initial-state perturbation, using identical sequence
        def run_with_initial(s, initial):
            x=initial
            if isinstance(model, VanillaResidual):
                for t in range(s.size(1)): x=x+.08*torch.tanh(model.wx(x)+model.wv(s[:,t]))
            else:
                M=model.L.T@model.L+.12*torch.eye(model.hid,device=device)
                for t in range(s.size(1)):
                    q=-(x@M.T)+model.B(s[:,t])+model.bias; x=torch.clamp(x+model.h*tangent_project(x,q),-1,1)
            return x
        s=test[:32].to(device); a=torch.zeros(32,model.hid,device=device); b=a+1e-3*torch.randn_like(a)
        ratio=(run_with_initial(s,b)-run_with_initial(s,a)).norm(dim=1).div((b-a).norm(dim=1)).mean().item()
        final=model(test[:128].to(device)); state_norm=final.norm(dim=1).mean().item(); max_abs=final.abs().max().item()
    return {"loss_start":losses[0],"loss_final":losses[-1],"accuracy":acc,"perturbation_ratio":ratio,"mean_final_norm":state_norm,"max_abs_final":max_abs,"parameters":sum(p.numel() for p in model.parameters())+sum(p.numel() for p in head.parameters())}


def main():
    math_result=numerical_math_check()
    # Long-horizon sign-of-sum task; one scalar input per time step.
    ntr, nte, T = 256, 128, 200
    train=torch.randn(ntr,T,1); test=torch.randn(nte,T,1)
    y=(train.sum(1).squeeze(1)>0).float(); yt=(test.sum(1).squeeze(1)>0).float()
    baseline=train_model(VanillaResidual(1,16),train,y,test,yt)
    idea=train_model(ProjectedDissipative(1,16),train,y,test,yt)
    out={"device":device,"math_check":math_result,"baseline":baseline,"idea":idea}
    print(json.dumps(out, indent=2))
    with open("results.json","w") as f: json.dump(out,f,indent=2)

if __name__ == "__main__":
    try: main()
    except Exception as e:
        if device == "cuda":
            device="cpu"; print("CUDA failed, rerun with CPU", repr(e)); main()
        else: raise