import json, math, random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

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

A = torch.tensor([[0.93, 0.10], [-0.08, 0.88]], dtype=torch.float32)
B = torch.tensor([[0.12], [0.22]], dtype=torch.float32)
K = torch.tensor([[0.72, 0.42]], dtype=torch.float32)

class Policy(nn.Module):
    def __init__(self, innovation=True):
        super().__init__(); self.innovation = innovation
        self.gru = nn.GRUCell(5, 24)
        self.z = nn.Linear(24, 2); self.oh = nn.Linear(2, 2)
        self.nom = nn.Linear(2, 1); self.L = nn.Linear(2, 2); self.lout = nn.Linear(2, 1)
        self.unc = nn.Linear(2, 1)
    def forward(self, obs, mask, uprev, h=None, c=2.0):
        if h is None: h = torch.zeros(obs.shape[0], 24, device=obs.device)
        h = self.gru(torch.cat([obs * mask, mask, uprev], 1), h)
        z = self.z(h); pred = self.oh(z); nu = obs - pred
        uncertainty = F.softplus(self.unc(z)) + 0.02
        alpha = (1.0 / (1.0 + c * uncertainty)).clamp(0, 1)
        nominal = self.nom(z)
        correction = self.lout(torch.tanh(self.L(z)) * nu)
        u = nominal - (alpha * correction if self.innovation else 0.0)
        return u.clamp(-1, 1), pred, nu, uncertainty, h

def batch_rollout(policy, batch=96, horizon=35, corrupt=False, train=True):
    Adev, Bdev = A.to(device), B.to(device)
    x = torch.randn(batch, 2, device=device) * 0.7
    h = None; up = torch.zeros(batch, 1, device=device)
    xs=[]; obs=[]; masks=[]; us=[]; preds=[]; nus=[]; uncs=[]
    for t in range(horizon):
        m = torch.ones(batch, 2, device=device)
        if corrupt and 12 <= t < 18: m.zero_()
        noise = torch.randn(batch,2,device=device)*0.10
        o = (x + noise) * m
        u,p,n,q,h = policy(o,m,up,h)
        # reference is zero; controller should dissipate state.
        xnext = x @ Adev.T + u @ Bdev.T + torch.randn(batch,2,device=device)*0.018
        xs.append(x); obs.append(o); masks.append(m); us.append(u); preds.append(p); nus.append(n); uncs.append(q)
        x=xnext; up=u
    return map(torch.stack, (xs,obs,masks,us,preds,nus,uncs))

def train(policy, steps=220):
    opt=torch.optim.Adam(policy.parameters(),lr=3e-3)
    policy.train()
    for step in range(steps):
        x,o,m,u,p,n,q = batch_rollout(policy, batch=64, horizon=28, train=True)
        # Regulation, action effort, and estimator prediction objectives.
        loss = x.pow(2).sum(-1).mean() + 0.035*u.pow(2).mean() + 0.18*(o-p).pow(2).mean()
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(policy.parameters(),1.0); opt.step()
    return float(loss.detach().cpu())

def evaluate(policy, corrupt=False, episodes=192):
    policy.eval()
    with torch.no_grad():
        x,o,m,u,p,n,q = batch_rollout(policy,episodes,45,corrupt)
        state_cost=x.pow(2).sum(-1)
        rms=float(torch.sqrt(state_cost.mean()).cpu())
        transient=float(torch.sqrt(state_cost[:,12:20].mean()).cpu())
        recovery=float(torch.sqrt(state_cost[:,20:].mean()).cpu())
        action_std=float(u.std().cpu())
        pred=float((o-p).pow(2).mean().sqrt().cpu())
    return {'rms_state_error':rms,'transient_error':transient,'post_corruption_error':recovery,'action_std':action_std,'innovation_rmse':pred}

def math_check():
    # The gate alpha=1/(1+c sigma) is monotone and gain-limiting; verify numerically.
    sig=torch.linspace(0,10,1001); c=2.0; alpha=1/(1+c*sig)
    monotone=bool(torch.all(alpha[1:] <= alpha[:-1]))
    bounds=bool(torch.all((alpha>=0)&(alpha<=1)))
    # For a scalar innovation correction, increasing uncertainty strictly shrinks it.
    correction=torch.tensor(1.7)*alpha
    shrink=bool(correction[0] > correction[-1] and correction[0] <= 1.7)
    return {'gate_monotone':monotone,'gate_bounds':bounds,'correction_shrinks':shrink,'alpha_sigma0':float(alpha[0]),'alpha_sigma10':float(alpha[-1])}

def main():
    check=math_check(); results={'device':str(device),'math_check':check}
    base=Policy(False).to(device); idea=Policy(True).to(device)
    results['train_loss_baseline']=train(base); results['train_loss_idea']=train(idea)
    results['clean_baseline']=evaluate(base,False); results['clean_idea']=evaluate(idea,False)
    results['corrupt_baseline']=evaluate(base,True); results['corrupt_idea']=evaluate(idea,True)
    with open('results.json','w') as f: json.dump(results,f,indent=2)
    print(json.dumps(results,indent=2))

if __name__=='__main__': main()