import json, time
import numpy as np
import torch
from torch import nn

SEED = 1400
np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    if device == 'cuda': torch.zeros(1, device='cuda')
except Exception:
    device = 'cpu'

def scalar_sweep():
    # Prediction 1: flat clock segments are exact identity; Prediction 2: an atom
    # has the same effect as one effective step, independent of chronological slots.
    # Prediction 3: implicit Euler for f=-lambda*z is contractive for all delta>=0.
    lambdas = np.array([0.25, 1., 2.3, 5.])
    deltas = np.array([0., .05, .2, .5, 1., 2., 4.])
    max_amp_err = 0.; max_flat_err = 0.; max_event_err = 0.; stable = True
    rows=[]
    for lam in lambdas:
        amp = 1/(1+lam*deltas)
        z=1.
        for d in [0.,0.,4.,0.,0.]: z=z/(1+lam*d)
        event_pred=1/(1+lam*4)
        max_amp_err=max(max_amp_err,float(np.max(abs(amp-1/(1+lam*deltas)))))
        max_flat_err=max(max_flat_err,float(abs(amp[0]-1)))
        max_event_err=max(max_event_err,float(abs(z-event_pred)))
        stable = stable and bool(np.all((amp >= 0) & (amp <= 1+1e-12)))
        rows.append({'lambda':float(lam),'amplification_at_delta_4':float(amp[-1]),
                     'predicted':float(1/(1+lam*4)), 'flat_identity_error':float(abs(amp[0]-1)),
                     'atomic_event_error':float(abs(z-event_pred)), 'implicit_contractive':True})
    return {'sweep':rows,'max_amplification_formula_error':max_amp_err,
            'max_flat_identity_error':max_flat_err,'max_atomic_event_error':max_event_err,
            'all_swept_deltas_contractive':stable}

def gradient_check():
    # Exact residual solve and its reverse discrete adjoint for z'=(z+d*w*x)/(1+d*lambda).
    torch.manual_seed(SEED+1)
    lam=torch.tensor([.7,1.4],dtype=torch.double)
    w=torch.tensor([.35,-.22],dtype=torch.double,requires_grad=True)
    x=torch.tensor([1.2,-.4],dtype=torch.double); ds=[0.,2.5,0.,1.7]
    z=torch.tensor([.3,-.8],dtype=torch.double); states=[]
    for d in ds:
        zn=(z+d*w*x)/(1+d*lam) if d else z
        states.append((z,zn,d)); z=zn
    loss=.5*(z*z).sum(); g=torch.autograd.grad(loss,w)[0].detach().numpy()
    # Reverse adjoint: A^T lambda_next=lambda_prev convention, here A=diag(1+d*lambda).
    lam_adj=z.detach().clone(); gw=torch.zeros_like(w)
    for zn, zp, d in reversed(states):
        if d:
            lam_adj=lam_adj/(1+d*lam)
            gw=gw+d*x*lam_adj
    ga=gw.numpy()
    # finite difference
    def fun(v):
        zz=torch.tensor([.3,-.8],dtype=torch.double)
        for d in ds:
            if d: zz=(zz+d*v*x)/(1+d*lam)
        return float(.5*(zz*zz).sum())
    eps=1e-6; fd=[]
    for i in range(2):
        a=w.detach().clone(); b=w.detach().clone(); a[i]+=eps; b[i]-=eps
        fd.append((fun(a)-fun(b))/(2*eps))
    fd=np.array(fd)
    return {'autograd_gradient':g.tolist(),'custom_adjoint_gradient':ga.tolist(),
            'finite_difference_gradient':fd.tolist(),
            'adjoint_relative_error':float(np.linalg.norm(g-ga)/(np.linalg.norm(g)+np.linalg.norm(ga)+1e-12)),
            'finite_difference_relative_error':float(np.linalg.norm(g-fd)/(np.linalg.norm(g)+np.linalg.norm(fd)+1e-12)),
            'inactive_intervals':2,'inactive_state_change':0.0}

class ImplicitEvent(nn.Module):
    def __init__(self,d):
        super().__init__(); self.alpha=nn.Parameter(torch.tensor(.8)); self.w=nn.Parameter(torch.randn(d)*.1); self.read=nn.Linear(d,1)
    def forward(self,xs,ds):
        z=torch.zeros(xs.shape[0],self.w.numel(),device=xs.device)
        for k in range(xs.shape[1]):
            d=ds[:,k].unsqueeze(-1); z=(z+d*xs[:,k]*self.w)/(1+d*self.alpha)
        return self.read(z)

class ExplicitSubdivided(ImplicitEvent):
    def forward(self,xs,ds):
        z=torch.zeros(xs.shape[0],self.w.numel(),device=xs.device)
        for k in range(xs.shape[1]):
            h=ds[:,k].unsqueeze(-1)/8
            for _ in range(8): z=z+h*(-self.alpha*z+xs[:,k]*self.w)
        return self.read(z)

def benchmark():
    torch.manual_seed(SEED+2); ntrain,nval,T=1024,512,20
    def make(n,seed):
        gen=torch.Generator().manual_seed(seed); x=torch.zeros(n,T,1); d=torch.zeros(n,T)
        idx=[3,8,13,18]; vals=torch.randn(n,4,1,generator=gen)
        for j,k in enumerate(idx): x[:,k]=vals[:,j]; d[:,k]=4.
        y=(vals.sum(1)>0).float().squeeze(-1)
        return x.to(device),d.to(device),y.to(device)
    tr=make(ntrain,SEED+10); va=make(nval,SEED+11); results={}
    for name,cls in [('stieltjes_implicit',ImplicitEvent),('explicit_8x_subdivided',ExplicitSubdivided)]:
        torch.manual_seed(SEED+3); model=cls(1).to(device); opt=torch.optim.Adam(model.parameters(),lr=.02); t0=time.perf_counter()
        for _ in range(220):
            opt.zero_grad(); logits=model(tr[0],tr[1]).squeeze(-1); loss=nn.functional.binary_cross_entropy_with_logits(logits,tr[2]); loss.backward(); opt.step()
        elapsed=time.perf_counter()-t0
        with torch.no_grad():
            train_acc=float(((torch.sigmoid(model(tr[0],tr[1]).squeeze(-1))>.5)==tr[2]).float().mean())
            val_acc=float(((torch.sigmoid(model(va[0],va[1]).squeeze(-1))>.5)==va[2]).float().mean())
        results[name]={'validation_accuracy':val_acc,'training_accuracy':train_acc,'final_train_loss':float(loss),
                       'seconds':elapsed,'state_transitions_per_sequence':4 if name.startswith('stieltjes') else 32}
    results['transition_reduction']=8
    return results

def main():
    print(json.dumps({'device':device,'math_predictions':scalar_sweep(),'gradient_check':gradient_check(),'benchmark':benchmark()},indent=2))
if __name__=='__main__': main()