import json, math, random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
SEED=1669
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
try:
    device="cuda" if torch.cuda.is_available() else "cpu"
    if device=="cuda": torch.zeros(1,device="cuda")
except Exception:
    device="cpu"

class Schedule:
    def __init__(self,n,lam=0.,rho=0.,s=1):
        self.n,self.lam,self.rho,self.s=n,lam,rho,s
        assert lam+rho<=1
    def sample(self,batch,generator=None):
        u=torch.rand(batch,generator=generator); branch=torch.zeros(batch,dtype=torch.long)
        low=(u>=1-self.rho-self.lam)&(u<1-self.rho); full=u>=1-self.rho
        branch[low]=1; branch[full]=2
        vis=torch.full((batch,),self.n-1,dtype=torch.long)
        if low.any(): vis[low]=torch.randint(0,self.s+1,(int(low.sum()),),generator=generator)
        vis[full]=0
        rank=torch.rand(batch,self.n,generator=generator).argsort(1)
        visible=rank < vis[:,None]
        return ~visible,branch,vis
    def analytic(self):
        return {"v_min":0 if self.lam+self.rho>0 else self.n-1,
                "pi_s":self.rho+self.lam if self.s>=1 else self.rho+self.lam/(self.s+1),
                "full_atom":self.rho}

def schedule_check():
    rows=[]
    for lam,rho in [(0,0),(.1,0),(.1,.01),(.25,.01)]:
        sc=Schedule(8,lam,rho,1); g=torch.Generator().manual_seed(SEED+7)
        _,b,v=sc.sample(100000,g)
        rows.append({"lambda":lam,"rho":rho,"pred_pi1":sc.analytic()["pi_s"],"obs_pi1":float((v<=1).float().mean()),"pred_full":rho,"obs_full":float((b==2).float().mean()),"v_min":int(v.min())})
    return rows

def anchor_check():
    p=.8; eta=1.; steps=100; rows=[]
    for rho in [.01,.03,.1,.2]:
        b=0.
        for _ in range(steps): b-=eta*rho*(1/(1+math.exp(-b))-p)
        q=1/(1+math.exp(-b)); obs=abs(q-p); pred=.3*(1-eta*rho*p*(1-p))**steps
        rows.append({"rho":rho,"observed_error":obs,"linear_prediction":pred,"ratio":obs/pred})
    return rows

class TinyMaskedPredictor(nn.Module):
    def __init__(self,n=8):
        super().__init__(); self.n=n; self.emb=nn.Embedding(3,12); self.pos=nn.Embedding(n,12)
        self.net=nn.Sequential(nn.Linear(n*12,40),nn.Tanh(),nn.Linear(40,n*2))
    def forward(self,x):
        z=self.emb(x.long())+self.pos(torch.arange(self.n,device=x.device))[None]
        return self.net(z.reshape(x.shape[0],-1)).reshape(x.shape[0],self.n,2)

def make_data(num,n=8,p=.8):
    mode=(torch.rand(num)<p).long(); alt=torch.tensor([j%2 for j in range(n)]).long()
    x=torch.where(mode[:,None].bool(),torch.ones(num,n,dtype=torch.long),alt[None].expand(num,n)).clone()
    noise=torch.rand(num,n)<.04; return torch.where(noise,1-x,x),mode

def train_eval(lam,rho,steps=250):
    torch.manual_seed(SEED); n=8; x,_=make_data(1200,n); xv,_=make_data(500,n)
    model=TinyMaskedPredictor(n).to(device); opt=torch.optim.Adam(model.parameters(),lr=.005); g=torch.Generator().manual_seed(SEED+99); sc=Schedule(n,lam,rho,1)
    for _ in range(steps):
        ix=torch.randint(0,len(x),(64,),generator=g); xb=x[ix].to(device); masks,_,_=sc.sample(len(ix),g); masks=masks.to(device)
        inp=xb.clone(); inp[masks]=2; loss=F.cross_entropy(model(inp)[masks],xb[masks])
        opt.zero_grad(); loss.backward(); opt.step()
    model.eval()
    with torch.no_grad():
        xvd=xv.to(device); inp=xvd.clone(); mask=torch.zeros_like(inp,dtype=torch.bool); mask[:,0]=True; inp[mask]=2
        cond=float((model(inp).argmax(-1)[mask]==xvd[mask]).float().mean().cpu())
        empty=torch.full((len(xv),n),2,dtype=torch.long,device=device); q=float(model(empty).softmax(-1)[:,0,1].mean().cpu())
    return {"lambda":lam,"rho":rho,"conditional_accuracy":cond,"empty_mode_prob":q,"mode_abs_error":abs(q-.8)}

def main():
    out={"device":device,"schedule_checks":schedule_check(),"anchor_checks":anchor_check(),"runs":[train_eval(*z) for z in [(0,0),(.1,0),(.1,.01),(.25,.01)]]}
    open("results.json","w").write(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
if __name__=="__main__": main()