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

SEED = 2765
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)


def tube_rollout(Js, r0, ds):
    r = np.asarray(r0, dtype=float).copy()
    out = [r.copy()]
    for J, d in zip(Js, ds):
        r = np.abs(np.asarray(J)) @ r + np.asarray(d)
        out.append(r.copy())
    return np.asarray(out)


def toy_verification():
    # For J=gamma*lambda, the theory predicts q=1 boundary, slope log(q) above it,
    # and stationary radius d/(1-q) below it.
    lam = 0.8
    gamma_values = np.array([0.50, 0.90, 1.10, 1.30, 1.50])
    d = 0.01; H = 400
    rows = []
    for gamma in gamma_values:
        q = gamma * lam
        rs = tube_rollout([[[q]]]*H, [0.0], [[d]]*H)[:, 0]
        if q < 1:
            predicted_plateau = d/(1-q)
            observed_plateau = float(np.mean(rs[-10:]))
            relerr = abs(observed_plateau-predicted_plateau)/predicted_plateau
            slope_pred = 0.0
            slope_obs = 0.0
        else:
            predicted_plateau = float('inf')
            # Fit log radius after the transient; asymptotic slope should log(q).
            ix = np.arange(150, H+1)
            slope_obs = float(np.polyfit(ix, np.log(rs[150:]), 1)[0])
            slope_pred = math.log(q)
            relerr = abs(slope_obs-slope_pred)/abs(slope_pred)
        rows.append(dict(gamma=float(gamma), q=float(q), stable_pred=bool(q<1),
                         stable_observed=(bool(abs(slope_obs) <= 1e-10) if q >= 1 else bool(abs(slope_obs) <= 1e-10)), predicted_plateau=predicted_plateau,
                         observed_plateau=observed_plateau if q<1 else None,
                         predicted_log_slope=slope_pred, observed_log_slope=slope_obs,
                         relative_error=float(relerr)))
    # Disturbance scaling prediction at q=.72: equilibrium is linear in d.
    q=.72; ds=np.array([.002,.005,.01,.02]); final=[]
    for di in ds:
        rr=tube_rollout([[[q]]]*120,[0],[[di]]*120)[-1,0]
        final.append(float(rr))
    ratio=np.asarray(final)/ds
    scaling_error=float(np.max(np.abs(ratio-1/(1-q))/(1/(1-q))))
    # Locate transition empirically among a fine sweep: first q >= 1.
    fine=np.linspace(.70,1.50,161); empirical_boundary=float(fine[np.argmax(fine*lam>=1)])
    return {'gain_sweep':rows, 'predicted_boundary_gamma':1/lam,
            'observed_boundary_gamma':empirical_boundary, 'disturbances':ds.tolist(),
            'final_radii':final, 'radius_over_d':ratio.tolist(),
            'predicted_radius_over_d':1/(1-q), 'scaling_max_relative_error':scaling_error}


class Transition(nn.Module):
    def __init__(self, dim=2):
        super().__init__()
        self.net=nn.Sequential(nn.Linear(dim+1,24),nn.Tanh(),nn.Linear(24,dim))
    def forward(self,z,u): return z + 0.18*self.net(torch.cat([z,u],-1))


def make_data(n=900, horizon=12):
    zs=[]; us=[]; ys=[]
    for _ in range(n):
        z=np.random.uniform(-1,1,2).astype('float32')
        for k in range(horizon):
            u=np.random.uniform(-.7,.7,1).astype('float32')
            # nonlinear stable plant with transition noise
            y=np.array([.88*z[0]+.12*np.tanh(z[1])+ .08*u[0],
                        .80*z[1]+.10*np.sin(z[0])- .04*u[0]],dtype='float32')
            y += np.random.normal(0,.008,2).astype('float32')
            zs.append(z); us.append(u); ys.append(y); z=y
    return torch.tensor(np.asarray(zs)),torch.tensor(np.asarray(us)),torch.tensor(np.asarray(ys))


def neural_comparison():
    z,u,y=make_data(); n=len(z); split=int(.8*n)
    # Same initialization and training budget; tube penalty uses exact autograd Jacobian.
    def train(tube):
        torch.manual_seed(SEED); m=Transition(); opt=torch.optim.Adam(m.parameters(),lr=3e-3)
        for epoch in range(28):
            ix=torch.randperm(split)[:128]; zb,ub,yb=z[ix],u[ix],y[ix]
            pred=m(zb,ub); loss=((pred-yb)**2).mean()
            if tube:
                # Conservative residual radius: held-out-like fixed upper quantile estimate.
                r=torch.full((len(ix),2),.025)
                zz=zb.detach().clone().requires_grad_(True)
                pp=m(zz,ub)
                J=[]
                for j in range(2):
                    grad=torch.autograd.grad(pp[:,j].sum(),zz,create_graph=False,retain_graph=True)[0]
                    J.append(grad)
                J=torch.stack(J,1).abs()
                rnext=torch.bmm(J,r.unsqueeze(-1)).squeeze(-1)+.012
                # nominal state constraint |z_i| <= 1, robust backoff.
                robust_violation=torch.relu(torch.abs(pred)+rnext-1.0)
                loss=loss + .015*rnext.mean() + .20*robust_violation.pow(2).mean()
            opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            one=((m(z[split:],u[split:])-y[split:])**2).mean().item()
        # Free-running 20-step rollout on test starts; compare state constraint violations.
        viol=[]; mse=[]
        for start in range(split,min(split+100, n-20),20):
            zz=z[start:start+1].clone(); truth=z[start:start+20].clone();
            for k in range(20):
                uu=u[start+k:start+k+1]
                zz=m(zz,uu); mse.append(float(((zz-truth[k:k+1])**2).mean()))
                viol.append(float((torch.abs(zz)>1).any()))
        return {'one_step_mse':one,'rollout_mse':float(np.mean(mse)),
                'violation_rate':float(np.mean(viol))}
    return {'baseline':train(False),'tube_regularized':train(True)}


def main():
    out={'seed':SEED,'toy_verification':toy_verification()}
    try:
        out['neural_comparison']=neural_comparison()
        out['device']='cpu'
    except Exception as e:
        out['neural_error']=repr(e)
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()