import json, math, os
import numpy as np

# Differentiable equilibrium projection MVP: F(z;u)=z^3+a*z-u.
# Its Jacobian is 3z^2+a, so dz*/du=1/(3z*^2+a).

def solve_newton(u, a, z0, maxit=30, tol=1e-10):
    z=float(z0)
    for k in range(maxit):
        f=z**3+a*z-u
        j=3*z*z+a
        if abs(f)<tol: return z,k+1,abs(f)
        z -= f/j
    return z,maxit,abs(z**3+a*z-u)

def solve_gd_residual(u,a,z0,tau,maxit=10000,tol=1e-8):
    # z <- z - tau J^T F, as in the proposed least-squares implicit iteration.
    z=float(z0)
    for k in range(maxit):
        f=z**3+a*z-u; j=3*z*z+a
        if abs(f)<tol: return z,k+1,abs(f)
        z -= tau*j*f
    return z,maxit,abs(z**3+a*z-u)

def root(u,a): return solve_newton(u,a,0,maxit=100)[0]

def main():
    rng=np.random.default_rng(7)
    rows=[]
    # Prediction 1: for linear F=a*z-u, residual contraction is |1-tau*a^2|,
    # and the stability boundary is tau*a^2=2.
    lin=[]
    for a in [0.5,1.,2.]:
        for tau in [0.25/a**2, 0.9/a**2, 1.1/a**2, 1.9/a**2, 2.1/a**2]:
            u=1.3; z=4.; r0=abs(a*z-u); rs=[]
            for k in range(12):
                rs.append(abs(a*z-u)); z -= tau*a*(a*z-u)
            observed=(rs[-1]/rs[0])**(1/11)
            predicted=abs(1-tau*a*a)
            lin.append({'a':a,'tau_a2':tau*a*a,'predicted_ratio':predicted,'observed_ratio':observed,'stable_pred':tau*a*a<2,'decreased':rs[-1]<rs[0]})
    rows.append(('linear_contraction',lin))

    # Prediction 2: implicit sensitivity is inverse Jacobian; finite differences should scale 1/a
    sens=[]
    for a in [0.1,0.25,0.5,1.,2.,4.]:
        u=0.7; eps=1e-5
        fd=(root(u+eps,a)-root(u-eps,a))/(2*eps)
        pred=1/(3*root(u,a)**2+a)
        sens.append({'a':a,'predicted_1_over_J':pred,'finite_difference':fd,'relative_error':abs(fd-pred)/pred})
    rows.append(('implicit_sensitivity',sens))

    # Controlled conditioning sweep at u=0: z*=0, J=a exactly, so |dz/du|=1/a.
    conditioning=[]
    for a in [1e-1, 3e-2, 1e-2, 3e-3, 1e-3]:
        u=0.0; eps=1e-6
        fd=(root(u+eps,a)-root(u-eps,a))/(2*eps)
        conditioning.append({'a_sigma_min_J':a,'predicted_gain':1/a,
                             'finite_difference_gain':fd,
                             'relative_error':abs(fd-1/a)/(1/a)})
    rows.append(('conditioning_gain_at_singular_limit',conditioning))

    # Verify the custom implicit backward rule against torch autograd finite differences.
    try:
        import torch
        from implicit_layer import implicit_project
        u=torch.tensor(0.7,dtype=torch.double,requires_grad=True)
        a_t=torch.tensor(0.4,dtype=torch.double,requires_grad=True)
        z=implicit_project(u,a_t,steps=30)
        (z*z/2).backward()
        z0=float(z.detach()); j=3*z0*z0+float(a_t.detach())
        expected_u=z0/j; expected_a=-z0*z0/j
        implicit_check={'autograd_du':float(u.grad),'expected_du':expected_u,
                        'autograd_da':float(a_t.grad),'expected_da':expected_a,
                        'max_abs_error':max(abs(float(u.grad)-expected_u),abs(float(a_t.grad)-expected_a))}
    except Exception as e:
        implicit_check={'error':str(e)}
    rows.append(('implicit_backward_check',implicit_check))

    # Prediction 3: Newton reaches tolerance from very different predictor errors,
    # while penalty-only optimization retains initialization/penalty dependence.
    init_sweep=[]
    u=1.2; a=.35
    for z0 in [-10,-2,0,2,10]:
        z,k,r=solve_newton(u,a,z0)
        init_sweep.append({'z0':z0,'newton_iters':k,'final_residual':r})
    rows.append(('newton_initialization_invariance',init_sweep))

    # Small ML comparison: train direct predictor and physics-penalty predictor;
    # projected predictor uses same direct network output followed by Newton.
    import torch
    torch.manual_seed(7)
    torch.set_num_threads(4)
    n=500; uall=torch.linspace(-2,2,n).unsqueeze(1)
    aa=.35
    yall=torch.tensor([root(float(u),aa) for u in uall[:,0]]).float().unsqueeze(1)
    tr=slice(0,350); te=slice(350,n)
    def net(): return torch.nn.Sequential(torch.nn.Linear(1,24),torch.nn.Tanh(),torch.nn.Linear(24,1))
    def train(mode, epochs=500):
        m=net(); opt=torch.optim.Adam(m.parameters(),lr=.02)
        for _ in range(epochs):
            pred=m(uall[tr]); f=pred**3+aa*pred-uall[tr]
            loss=((pred-yall[tr])**2).mean() if mode=='direct' else (f*f).mean()
            opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            p=m(uall[te]); f=p**3+aa*p-uall[te]
            if mode=='project': pass
            pred_res=float(f.abs().mean()); mse=float(((p-yall[te])**2).mean())
        return m,pred_res,mse
    md,rd,ed=train('direct'); mp,rp,ep=train('penalty')
    # deterministic projection of penalty network's predictions, fixed Newton solve
    with torch.no_grad(): raw=mp(uall[te]).numpy().ravel()
    projected=np.array([solve_newton(float(u),aa,float(z))[0] for u,z in zip(uall[te,0],raw)])
    up=uall[te,0].numpy(); truth=yall[te,0].numpy()
    proj_res=float(np.mean(np.abs(projected**3+aa*projected-up)))
    proj_mse=float(np.mean((projected-truth)**2))
    rows.append(('mini_experiment',{'direct_test_mean_residual':rd,'penalty_test_mean_residual':rp,'penalty_projected_residual':proj_res,'direct_mse':ed,'penalty_mse':ep,'projected_mse':proj_mse,'train_points':350,'test_points':150}))
    with open('results.json','w') as f: json.dump(dict(rows),f,indent=2)
    print(json.dumps(dict(rows),indent=2))

if __name__=='__main__': main()