import math, random, json
from pathlib import Path
import numpy as np
from scipy.linalg import expm

SEED=3154
np.random.seed(SEED); random.seed(SEED)

def hamiltonian(a,b,q,r):
    return np.array([[a, -b*b/r], [-q, -a]], dtype=float)

def feedback(a,b,q,r,pf,T):
    M=expm(hamiltonian(a,b,q,r)*T)
    D=float(M[0,0]+M[0,1]*pf)
    P=float((M[1,0]+M[1,1]*pf)/D) if abs(D)>1e-15 else np.nan
    return D,P,M

def math_check():
    # Direct determinant verification of the formula in the idea.
    samples=[(-1.,1.,1.,.2),(-1.,1.,.1,1.),(-1.,1.,1.,1.),(.3,2.,.7,.4)]
    rows=[]
    for a,b,q,r in samples:
        H=hamiltonian(a,b,q,r)
        stated=b*b*q/r-a*a
        actual=float(np.linalg.det(H))
        rows.append({'a':a,'b':b,'q':q,'r':r,'stated_det':stated,'actual_det':actual,
                     'identity_error':abs(actual-(-a*a-b*b*q/r))})
    # For valid q>=0,r>0 dynamics, roots can still occur in the hyperbolic case.
    a,b,q,r,pf=-1.,1.,1.,.2,1.
    H=hamiltonian(a,b,q,r)
    disc=a*a+b*b*q/r
    # exp(H T)=cosh(kT)I+sinh(kT)H/k, k^2=disc.
    k=math.sqrt(disc)
    coeff=(a-pf*b*b/r)/k
    # D=cosh(kT)+coeff*sinh(kT); solve tanh(kT)=-1/coeff if possible.
    predicted=None
    z=-1./coeff
    if coeff < -1 and 0 < z < 1:
        predicted=math.atanh(z)/k
    Ts=np.linspace(0,4,4001)
    Ds=np.array([feedback(a,b,q,r,pf,float(t))[0] for t in Ts])
    crossings=np.where(Ds[:-1]*Ds[1:]<=0)[0]
    numerical=float(Ts[crossings[0]]) if len(crossings) else None
    return {'determinant_rows':rows,
            'all_determinant_identities_correct':all(x['identity_error']<1e-12 for x in rows),
            'claimed_elliptic_det_for_first_case':rows[0]['stated_det'],
            'actual_det_for_first_case':rows[0]['actual_det'],
            'hyperbolic_predicted_conjugate':predicted,
            'hyperbolic_numerical_crossing_grid':numerical,
            'conjugate_point_observed':numerical is not None,
            'claim_formula_consistent':all(abs(x['stated_det']-x['actual_det'])<1e-12 for x in rows)}

def quadratic_run(kind, steps=250, lr=.08, T=.18):
    A=np.diag([.5,1.,3.,8.]); x=np.array([2.,-1.5,.8,-.5],dtype=float)
    losses=[]; horizons=[]; gains=[]; denoms=[]
    for _ in range(steps):
        grad=A@x
        if kind=='sgd':
            x=x-lr*grad; horizons.append(0.); gains.append(1.); denoms.append(1.)
        else:
            gs=[]; ds=[]; ts=[]
            for lam in np.diag(A):
                r=.2; t=T
                if kind=='adaptive':
                    while True:
                        d,p,_=feedback(-lam,1.,lam,r,lam,t)
                        if d>.05 and np.isfinite(p): break
                        t*=.8
                        if t<1e-5: r*=2.; t=T
                    d,p,_=feedback(-lam,1.,lam,r,lam,t)
                else: d,p,_=feedback(-lam,1.,lam,r,lam,T)
                if not np.isfinite(p) or d<=.005: p=np.sign(p)*min(abs(p) if np.isfinite(p) else 1e6,100.)
                gs.append(p/r); ds.append(d); ts.append(t)
            x=x-lr*np.array(gs)*x
            horizons.append(float(np.mean(ts))); gains.append(float(np.max(np.abs(gs)))); denoms.append(float(np.min(ds)))
        losses.append(.5*float(x@A@x))
    return {'final_loss':losses[-1],'min_loss':min(losses),'max_gain':max(gains),
            'min_D':min(denoms),'final_x_norm':float(np.linalg.norm(x)),
            'mean_horizon':float(np.mean(horizons)),
            'diverged':bool(not np.isfinite(losses[-1]) or losses[-1]>1e8)}

def mlp_run(kind, epochs=18):
    try:
        import torch
        from sklearn.datasets import load_digits
        from sklearn.model_selection import train_test_split
        torch.manual_seed(SEED)
        dev='cuda' if torch.cuda.is_available() else 'cpu'
        X,y=load_digits(return_X_y=True); X=X.astype('float32')/16.; y=y.astype('int64')
        Xtr,Xte,ytr,yte=train_test_split(X,y,test_size=.25,random_state=SEED,stratify=y)
        Xtr=torch.tensor(Xtr,device=dev); ytr=torch.tensor(ytr,device=dev); Xte=torch.tensor(Xte,device=dev); yte=torch.tensor(yte,device=dev)
        model=torch.nn.Sequential(torch.nn.Linear(64,32),torch.nn.Tanh(),torch.nn.Linear(32,10)).to(dev)
        opt=torch.optim.SGD(model.parameters(),lr=.08); losses=[]
        for _ in range(epochs):
            for ix in torch.randperm(len(Xtr),device=dev).split(128):
                opt.zero_grad(); loss=torch.nn.functional.cross_entropy(model(Xtr[ix]),ytr[ix]); loss.backward()
                if kind=='hh':
                    d,p,_=feedback(-1.,1.,1.,.5,1.,.2); mult=min(max(p/.5,.25),4.)
                    for par in model.parameters():
                        if par.grad is not None: par.grad.mul_(mult)
                opt.step()
            with torch.no_grad(): losses.append(float(torch.nn.functional.cross_entropy(model(Xtr),ytr).cpu()))
        with torch.no_grad(): acc=float((model(Xte).argmax(1)==yte).float().mean().cpu())
        return {'final_loss':losses[-1],'test_accuracy':acc,'device':dev,
                'feedback_multiplier':(feedback(-1,1,1,.5,1,.2)[1]/.5 if kind=='hh' else 1.)}
    except Exception as e:
        return {'error':type(e).__name__+': '+str(e)}

if __name__=='__main__':
    result={'math_check':math_check(),'quadratic':{},'digits_mlp':{}}
    for k in ('sgd','fixed_horizon','adaptive'): result['quadratic'][k]=quadratic_run(k)
    for k in ('sgd','hh'): result['digits_mlp'][k]=mlp_run(k)
    Path('results.json').write_text(json.dumps(result,indent=2)); print(json.dumps(result,indent=2))