"""Original CPU teaching experiments. No pretrained weights or network access.
Architecture adapted from Karpathy/build-nanogpt; see NOTICE.md.
"""
import copy, io, json, math, platform
from dataclasses import dataclass, asdict
from pathlib import Path
import torch
from torch import nn
from torch.nn import functional as F

torch.set_num_threads(1)
torch.manual_seed(1337)
torch.use_deterministic_algorithms(True)

@dataclass
class Config:
    vocab_size: int = 17
    block_size: int = 12
    n_embd: int = 16
    n_head: int = 4
    n_layer: int = 2

class Attention(nn.Module):
    def __init__(self,c):
        super().__init__(); self.c=c
        self.qkv=nn.Linear(c.n_embd,3*c.n_embd)
        self.proj=nn.Linear(c.n_embd,c.n_embd)
        self.proj.residual=True
    def forward(self,x):
        B,T,C=x.shape; H=self.c.n_head
        q,k,v=[t.reshape(B,T,H,C//H).transpose(1,2) for t in self.qkv(x).chunk(3,dim=-1)]
        a=F.scaled_dot_product_attention(q,k,v,is_causal=True,dropout_p=0.0)
        return self.proj(a.transpose(1,2).contiguous().view(B,T,C))

class Block(nn.Module):
    def __init__(self,c):
        super().__init__(); self.ln1=nn.LayerNorm(c.n_embd); self.ln2=nn.LayerNorm(c.n_embd)
        self.attn=Attention(c)
        self.mlp=nn.Sequential(nn.Linear(c.n_embd,4*c.n_embd),nn.GELU(approximate='tanh'),nn.Linear(4*c.n_embd,c.n_embd))
        self.mlp[-1].residual=True
    def forward(self,x):
        x=x+self.attn(self.ln1(x)); return x+self.mlp(self.ln2(x))

class TinyGPT(nn.Module):
    def __init__(self,c=Config()):
        super().__init__(); self.c=c
        self.wte=nn.Embedding(c.vocab_size,c.n_embd); self.wpe=nn.Embedding(c.block_size,c.n_embd)
        self.blocks=nn.Sequential(*[Block(c) for _ in range(c.n_layer)])
        self.ln=nn.LayerNorm(c.n_embd); self.head=nn.Linear(c.n_embd,c.vocab_size,bias=False)
        self.head.weight=self.wte.weight
        self.apply(self._init)
    def _init(self,m):
        if isinstance(m,(nn.Linear,nn.Embedding)):
            std=.02*(2*self.c.n_layer)**-.5 if getattr(m,'residual',False) else .02
            nn.init.normal_(m.weight,0,std)
            if isinstance(m,nn.Linear) and m.bias is not None: nn.init.zeros_(m.bias)
    def forward(self,idx):
        B,T=idx.shape
        if T>self.c.block_size: raise ValueError('context overflow')
        x=self.wte(idx)+self.wpe(torch.arange(T,device=idx.device))
        return self.head(self.ln(self.blocks(x)))

def schedule(i,warm=4,end=20,peak=6e-4,floor=6e-5):
    if i<warm: return peak*(i+1)/warm
    if i>end: return floor
    return floor+.5*(1+math.cos(math.pi*(i-warm)/(end-warm)))*(peak-floor)

def flat_grad(m):
    return torch.cat([p.grad.reshape(-1) for p in m.parameters()])

def mean_loss(m,x,y):
    return F.cross_entropy(m(x).reshape(-1,m.c.vocab_size),y.reshape(-1))

def optimizer(m):
    ps=list(m.parameters())
    return torch.optim.AdamW([{'params':[p for p in ps if p.ndim>=2],'weight_decay':.1},{'params':[p for p in ps if p.ndim<2],'weight_decay':0.}],lr=.002,betas=(.9,.95),eps=1e-8)

def one_step(m,o,x,y):
    o.zero_grad(set_to_none=True); l=mean_loss(m,x,y); l.backward()
    norm=nn.utils.clip_grad_norm_(m.parameters(),1.); o.step()
    return l.item(),norm.item()

def tests():
    out={'python':platform.python_version(),'torch':torch.__version__,'device':'cpu','threads':torch.get_num_threads(),'seed':1337}
    m=TinyGPT().double(); c=m.c
    assert m.head.weight is m.wte.weight
    out['unique_parameters']=sum(p.numel() for p in m.parameters())
    # Same-valued independent tensors do not implement tying.
    independent=nn.Parameter(m.wte.weight.detach().clone())
    assert independent.data_ptr()!=m.wte.weight.data_ptr()
    out['tied_storage']=m.head.weight.data_ptr()==m.wte.weight.data_ptr()
    W=torch.randn(3,5,dtype=torch.float64); a=torch.randn(2,3,dtype=torch.float64)
    out['transpose_import_error']=(a@W-F.linear(a,W.T)).abs().max().item()
    q,k,v=[torch.randn(2,2,5,4,dtype=torch.float64,requires_grad=True) for _ in range(3)]
    mask=torch.ones(5,5,dtype=torch.bool).tril()
    manual=((q@k.transpose(-1,-2)/2).masked_fill(~mask,float('-inf')).softmax(-1))@v
    sdpa=F.scaled_dot_product_attention(q,k,v,is_causal=True)
    gm=torch.autograd.grad(manual.square().sum(),(q,k,v),retain_graph=True)
    gs=torch.autograd.grad(sdpa.square().sum(),(q,k,v))
    out['attention_forward_max_error']=(manual-sdpa).abs().max().item()
    out['attention_grad_max_error']=max((a-b).abs().max().item() for a,b in zip(gm,gs))
    assert out['attention_grad_max_error']<1e-12
    x=torch.randint(c.vocab_size,(5,6)); y=torch.randint(c.vocab_size,(5,6))
    models=[copy.deepcopy(m) for _ in range(3)]
    mean_loss(models[0],x,y).backward()
    for sl in [slice(0,2),slice(2,5)]:
        (mean_loss(models[1],x[sl],y[sl])*x[sl].shape[0]/5).backward()
        (mean_loss(models[2],x[sl],y[sl])/2).backward()
    g=flat_grad(models[0]); out['weighted_accumulation_error']=(g-flat_grad(models[1])).abs().max().item()
    out['wrong_equal_microbatch_error']=(g-flat_grad(models[2])).abs().max().item()
    assert out['weighted_accumulation_error']<1e-12 and out['wrong_equal_microbatch_error']>1e-5
    # DDP mathematical simulation only: two equal-size rank means averaged.
    ranks=[copy.deepcopy(m) for _ in range(2)]; full=copy.deepcopy(m)
    mean_loss(full,x[:4],y[:4]).backward()
    for r in range(2): mean_loss(ranks[r],x[2*r:2*r+2],y[2*r:2*r+2]).backward()
    out['ddp_average_math_error']=(flat_grad(full)-(flat_grad(ranks[0])+flat_grad(ranks[1]))/2).abs().max().item()
    assert out['ddp_average_math_error']<1e-12
    stream=torch.arange(30); B,T,R=2,3,2
    batches=[]
    for r in range(R):
        buf=stream[B*T*r:B*T*r+B*T+1]; bx=buf[:-1].view(B,T); by=buf[1:].view(B,T)
        assert torch.equal(by,bx+1); batches.append({'rank':r,'x':bx.tolist(),'y':by.tolist(),'next_start':B*T*r+B*T*R})
    out['batch_shifts']=batches
    out['schedule']={str(i):schedule(i) for i in [0,3,4,12,20,21]}
    assert schedule(3)==schedule(4)==6e-4 and schedule(20)==schedule(21)==6e-5
    # Completion-only scoring; row with lower sum is not always lower mean.
    losses=torch.tensor([[.7,.7,0],[.5,.5,.5]],dtype=torch.float64)
    mask=torch.tensor([[1,1,0],[1,1,1]],dtype=torch.float64)
    out['completion_score_sum_argmin']=int((losses*mask).sum(-1).argmin())
    out['completion_score_mean_argmin']=int(((losses*mask).sum(-1)/mask.sum(-1)).argmin())
    # Actual shifted completion mask: last prompt logit predicts first completion.
    tokens=torch.tensor([[1,2,3,4],[1,2,4,0]])
    cmask=torch.tensor([[0,0,1,1],[0,0,1,0]],dtype=torch.float64)
    logits=torch.zeros(2,4,5,dtype=torch.float64)
    logits[0,1,3]=5; logits[0,2,4]=5
    logits[1,1,4]=-5
    shifted=F.cross_entropy(logits[:,:-1].reshape(-1,5),tokens[:,1:].reshape(-1),reduction='none').reshape(2,3)
    good=(shifted*cmask[:,1:]).sum(-1)/cmask[:,1:].sum(-1)
    bad=(shifted*cmask[:,:-1]).sum(-1)/cmask[:,:-1].sum(-1)
    assert good.argmin().item()==0
    # Perturbing only prompt-target loss must not change completion-only score.
    changed_logits=logits.clone(); changed_logits[:,0,:]=torch.randn(2,5,dtype=torch.float64)*20
    shifted2=F.cross_entropy(changed_logits[:,:-1].reshape(-1,5),tokens[:,1:].reshape(-1),reduction='none').reshape(2,3)
    good2=(shifted2*cmask[:,1:]).sum(-1)/cmask[:,1:].sum(-1)
    assert torch.equal(good,good2)
    out['shifted_completion_scores']=good.tolist()
    out['wrong_unshifted_mask_scores']=bad.tolist()
    out['prompt_loss_invariance_error']=(good-good2).abs().max().item()
    # Causality: changing a future token cannot change earlier logits.
    changed=x.clone(); changed[:,4:]=(changed[:,4:]+1)%c.vocab_size
    out['causality_prefix_error']=(m(x)[:,:4]-m(changed)[:,:4]).abs().max().item()
    assert out['causality_prefix_error']==0
    # Exact CPU resume test, including optimizer and RNG, no hidden external state.
    train=TinyGPT().double(); opt=optimizer(train)
    one_step(train,opt,x,y)
    buf=io.BytesIO(); torch.save({'config':asdict(train.c),'model':train.state_dict(),'optimizer':opt.state_dict(),'rng':torch.get_rng_state(),'next_step':1,'loader_position':0},buf)
    nextx=torch.randint(c.vocab_size,(2,6)); nexty=torch.randint(c.vocab_size,(2,6))
    one_step(train,opt,nextx,nexty)
    buf.seek(0); ck=torch.load(buf,weights_only=True)
    restored=TinyGPT(Config(**ck['config'])).double(); restored.load_state_dict(ck['model']); ropt=optimizer(restored); ropt.load_state_dict(ck['optimizer']); torch.set_rng_state(ck['rng'])
    rx=torch.randint(c.vocab_size,(2,6)); ry=torch.randint(c.vocab_size,(2,6)); assert torch.equal(rx,nextx) and torch.equal(ry,nexty)
    one_step(restored,ropt,rx,ry)
    out['checkpoint_resume_max_error']=max((a-b).abs().max().item() for a,b in zip(train.parameters(),restored.parameters()))
    assert out['checkpoint_resume_max_error']==0 and restored.head.weight is restored.wte.weight
    # A tiny deterministic cyclic-language training run: not GPT-2 reproduction.
    torch.manual_seed(1337); toy=TinyGPT(); op=optimizer(toy)
    tx=(torch.arange(48).view(8,6)%c.vocab_size).long(); ty=(tx+1)%c.vocab_size
    initial=mean_loss(toy,tx,ty).item(); history=[]
    for step in range(100):
        loss,norm=one_step(toy,op,tx,ty)
        if step in [0,9,49,99]: history.append({'step':step,'pre_update_loss':loss,'pre_clip_norm':norm})
    final=mean_loss(toy,tx,ty).item(); assert final<initial
    out['toy_training']={'task':'next token is (current+1) modulo 17, fixed batch','initial_loss':initial,'final_post_update_loss':final,'history':history}
    # Supplemental notebook explorations, isolated RNG to preserve prior results.
    with torch.random.fork_rng():
        torch.manual_seed(1337)
        noise=torch.randn(100,768,dtype=torch.float64)
        unscaled=noise.sum(0); scaled=(noise/10).sum(0)
        out['residual_noise_std_unscaled']=unscaled.std().item()
        out['residual_noise_std_scaled']=scaled.std().item()
        assert torch.allclose(unscaled.std()/scaled.std(),torch.tensor(10.,dtype=torch.float64))
        net=nn.Sequential(nn.Linear(16,32),nn.GELU(),nn.Linear(32,1)).double()
        xx=torch.randn(4,16,dtype=torch.float64); yy=torch.randn(4,1,dtype=torch.float64)
        n2=copy.deepcopy(net)
        F.mse_loss(net(xx),yy).backward()
        for i in range(4): (F.mse_loss(n2(xx[i]),yy[i])/4).backward()
        out['notebook_mlp_equal_accumulation_error']=(flat_grad(net)-flat_grad(n2)).abs().max().item()
        assert out['notebook_mlp_equal_accumulation_error']<1e-12
    out['passed']=True
    return out

if __name__=='__main__':
    result=tests(); print(json.dumps(result,ensure_ascii=False,indent=2))