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

SEED=1469
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
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')

# Exact stochastic scalar map: x[t+1]=exp(mu+sigma*eps[t])*x[t].
# Its FTLE is exactly Normal(mu, sigma^2/T), making this a direct sanity check.
def toy_check():
    rng=np.random.default_rng(SEED)
    mu=-0.08; z=1.645; n=200000
    rows=[]
    for T in [4,16,64]:
        for sigma in [0.05,0.15,0.30]:
            lam=mu + sigma*rng.standard_normal((n,))/math.sqrt(T)
            obs_mu=float(lam.mean()); obs_sd=float(lam.std(ddof=1))
            p=float((lam>0).mean())
            pred_p=0.5*math.erfc((-mu)/(math.sqrt(2)*sigma/math.sqrt(T)))
            rows.append({'T':T,'sigma':sigma,'mu_obs':obs_mu,'mu_pred':mu,
                         'sd_obs':obs_sd,'sd_pred':sigma/math.sqrt(T),
                         'p_obs':p,'p_pred':pred_p})
    # Boundary sweep at T=16: UCB changes sign at sigma=-mu*sqrt(T)/z.
    T=16; boundary=-mu*math.sqrt(T)/z
    boundary_rows=[]
    for sigma in [0.12,0.18,0.195,0.23,0.30]:
        lam=mu + sigma*rng.standard_normal((n,))/math.sqrt(T)
        ucb=float(lam.mean()+z*lam.std(ddof=1))
        boundary_rows.append({'sigma':sigma,'ucb_obs':ucb,
                              'p_obs':float((lam>0).mean()),
                              'ucb_pred':mu+z*sigma/math.sqrt(T)})
    return {'predictions':{
        'variance_scaling':'sd(lambda_T)=sigma/sqrt(T)',
        'tail':'p_+=Phi(mu*sqrt(T)/sigma)',
        'boundary':f'ucb=0 at sigma={boundary:.5f} for mu={mu}, T={T}, z={z}'},
        'scaling_rows':rows,'boundary_rows':boundary_rows}

class NoisyRNN(nn.Module):
    def __init__(self,h=12):
        super().__init__(); self.h=h
        self.W=nn.Parameter(torch.randn(h,h)*0.45)
        self.U=nn.Parameter(torch.randn(h,2)*0.35)
        self.b=nn.Parameter(torch.zeros(h)); self.out=nn.Linear(h,1)
    def run(self,x, noise_sigma=0.0, return_ftle=False, K=4):
        B,T,_=x.shape; h=torch.zeros(B,self.h,device=x.device)
        if return_ftle:
            qs=[torch.randn(B,self.h,device=x.device) for _ in range(K)]
            qs=[q/(q.norm(dim=1,keepdim=True)+1e-8) for q in qs]
            sums=[torch.zeros(B,device=x.device) for _ in range(K)]
        for t in range(T):
            eps=torch.randn(K,self.h,self.h,device=x.device) if noise_sigma else None
            Wt=self.W + (noise_sigma*eps[0] if eps is not None else 0)
            pre=h@Wt.T + x[:,t]@self.U.T + self.b
            h=torch.tanh(pre)
            if return_ftle:
                deriv=1-torch.tanh(pre)**2
                new_qs=[]; new_sums=[]
                for k in range(K):
                    wk=self.W + (noise_sigma*eps[k] if eps is not None else 0)
                    q=(qs[k]@wk.T)*deriv
                    norm=q.norm(dim=1)+1e-8
                    new_qs.append(q/norm[:,None])
                    new_sums.append(sums[k]+torch.log(norm))
                qs, sums=new_qs, new_sums
        y=self.out(h).squeeze(-1)
        if return_ftle:
            return y, torch.stack(sums,dim=0)/T
        return y

def make_batch(B=64,T=20):
    x=torch.rand(B,T,2,device=device)
    # Two marker positions, target is sum of the marked values in channel 0.
    p1=torch.randint(0,T//2,(B,),device=device); p2=torch.randint(T//2,T,(B,),device=device)
    x[:,:,1]=0
    x[torch.arange(B,device=device),p1,1]=1
    x[torch.arange(B,device=device),p2,1]=1
    y=x[:,:,0].gather(1,p1[:,None]).squeeze(1)+x[:,:,0].gather(1,p2[:,None]).squeeze(1)
    return x,y

def train_variant(kind, steps=220, T=20, sigma=.08, rho=.35):
    torch.manual_seed(SEED+{'base':0,'mean':1,'ucb':2}[kind])
    model=NoisyRNN().to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
    z=1.645; history=[]
    for step in range(steps):
        x,y=make_batch(64,T)
        opt.zero_grad(); pred,lams=model.run(x,sigma,True,K=4)
        task=((pred-y)**2).mean(); mu=lams.mean(); sd=lams.std(unbiased=True)
        risk=torch.relu(mu + (z*sd if kind=='ucb' else 0.0))**2 if kind!='base' else mu*0
        # mean variant penalizes only positive estimated mean FTLE.
        loss=task + rho*risk
        loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5.0); opt.step()
        history.append((float(task.detach()),float(mu.detach()),float(sd.detach()),float(loss.detach())))
    # Clean task and independent noisy FTLE evaluation.
    with torch.no_grad():
        vals=[]; ys=[]; ts=[]
        for _ in range(12):
            x,y=make_batch(64,T); pred,lams=model.run(x,sigma,True,K=8)
            vals.append(float(((pred-y)**2).mean())); ts.append(lams.cpu().numpy().ravel())
        arr=np.concatenate(ts)
    return {'final_train':history[-1], 'cleanish_mse':float(np.mean(vals)),
            'ftle_mean':float(arr.mean()),'ftle_sd':float(arr.std(ddof=1)),
            'positive_fraction':float((arr>0).mean()),
            'ucb':float(arr.mean()+z*arr.std(ddof=1))}

def main():
    global device
    out={'device':str(device),'toy':toy_check(),'rnn':{}}
    for k in ['base','mean','ucb']:
        try: out['rnn'][k]=train_variant(k)
        except Exception as e:
            if device.type=='cuda':
                device=torch.device('cpu')
                out['rnn'][k]={'error':str(e),'fallback':'cpu'}
            else: raise
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()