"""Shared production operations; CPU-heavy validation/export is invoked on mini."""
import hashlib,json,os,random
from pathlib import Path
import numpy as np
import torch
from dense import Config,DenseLM,prepare_layout

ROOT=Path(__file__).resolve().parent
OUT=Path('/mnt/experiments/DENSE_VALUE_DIRECT_20260911')
DATA=Path('/home/youthk/Documents/Codex/2026-08-29/new-chat-6/DST_CYCLIC_ENRICHED_100M_20260908/data/cycle')
TOKENIZER=Path('/home/youthk/デスクトップ/訓練用_共通ファイル/tokenizers/babylm_2026_strict_small_16k_bpe_98dfab9e/tokenizer.json')
ARMS=('mha_gated_ffn_router',)
SOURCE_NAMES=('dense.py','configuration_dense.py','modeling_dense.py','runtime.py','train.py','exclusive.py','test_cpu.py','gpu_acceptance.py','resume_acceptance.py','parent_dense.py')
def source_identity():return {n:sha(ROOT/n) for n in SOURCE_NAMES}
def sha(p):
    h=hashlib.sha256()
    with Path(p).open('rb') as f:
        for block in iter(lambda:f.read(8<<20),b''):h.update(block)
    return h.hexdigest()
def write_json(p,value):
    p=Path(p);p.parent.mkdir(parents=True,exist_ok=True)
    temp=p.with_name(p.name+'.partial');temp.write_text(json.dumps(value,indent=2)+'\n');os.replace(temp,p)
def save_checkpoint(p,value):
    p=Path(p);p.parent.mkdir(parents=True,exist_ok=True)
    temp=p.with_name(p.name+'.partial')
    with temp.open('wb') as f:torch.save(value,f);f.flush();os.fsync(f.fileno())
    os.replace(temp,p)
    write_json(str(p)+'.integrity.json',dict(bytes=p.stat().st_size,sha256=sha(p)))
def setup_gpu():
    torch.set_num_threads(2)
    torch.manual_seed(20260907);random.seed(20260907);np.random.seed(20260907)
    torch.cuda.manual_seed_all(20260907)
    torch.backends.cuda.matmul.allow_tf32=True
    torch.set_float32_matmul_precision('high')
    torch.use_deterministic_algorithms(True)
    torch.cuda.set_per_process_memory_fraction(21.5*2**30/torch.cuda.get_device_properties(0).total_memory)
def build(arm,compiled=False,value_variance=True,tensor_fusion=True):
    model=DenseLM(Config(arm=arm,backend='flash',value_variance=value_variance)).cuda()
    opt=torch.optim.AdamW([{'params':[p for p in model.parameters() if p.ndim>=2],'weight_decay':.1},{'params':[p for p in model.parameters() if p.ndim<2],'weight_decay':0.}],lr=8e-4,betas=(.9,.95),eps=1e-8,fused=True)
    if compiled:
        torch._dynamo.config.recompile_limit=64
        # Real batches vary in document count. Compile stable tensor-only regions;
        # keep varlen FlashAttention and its CPU-built metadata outside Dynamo.
        for block in model.blocks:
            block.ffn.forward=torch.compile(block.ffn.forward,fullgraph=True,dynamic=False,mode='default')
            if tensor_fusion:
                block.project_attention=torch.compile(block.project_attention,fullgraph=True,dynamic=False,mode='default')
                block.finish_attention=torch.compile(block.finish_attention,fullgraph=True,dynamic=False,mode='default')
            else:
                block.attn_gate.forward=torch.compile(block.attn_gate.forward,fullgraph=True,dynamic=False,mode='default')
    return model,opt
class Epoch:
    def __init__(self,number,root=DATA):
        self.row=json.loads((root/'dataset_manifest.json').read_text())['epochs'][number]
        self.length=self.row['sequence_length'];self.rows=self.row['rows']
        path=root/f"epoch_{self.row['epoch']:02d}_len_{self.length}"
        self.tokens=np.memmap(path/'tokens.bin',mode='r',dtype='<u2',shape=(self.rows,self.length))
        self.segments=np.memmap(path/'segments.bin',mode='r',dtype='<u2',shape=(self.rows,self.length))
        self.lengths=np.memmap(path/'lengths.bin',mode='r',dtype='<u2')
        self.words=np.memmap(path/'words.bin',mode='r',dtype='<u4')
    def host(self,start,stop):
        ids=torch.from_numpy(np.array(self.tokens[start:stop],dtype=np.int64))
        seg=torch.from_numpy(np.array(self.segments[start:stop],dtype=np.int32))
        valid=torch.arange(self.length)[None,:]<torch.from_numpy(np.array(self.lengths[start:stop],dtype=np.int64))[:,None]
        seg=torch.where(valid,seg+1,0)
        labels=torch.full_like(ids,-100)
        same=(seg[:,:-1]>0)&(seg[:,:-1]==seg[:,1:])
        labels[:,:-1]=torch.where(same,ids[:,1:],-100)
        return ids,seg,labels,int(same.sum()),int(self.words[start:stop].sum())
def update(model,opt,data,start,stop,micro,words):
    from dataclasses import replace
    probability=0. if words<30_000_000 else .02 if words<60_000_000 else .05
    for block in model.blocks:
        if block.ffn.c.dropout!=probability:
            block.ffn.c=replace(block.ffn.c,dropout=probability)
    batches=[data.host(i,min(i+micro,stop)) for i in range(start,stop,micro)]
    count=sum(b[3] for b in batches);exposure=sum(b[4] for b in batches)
    assert count>0
    opt.zero_grad(set_to_none=True);total=torch.zeros((),device='cuda')
    for ids,seg,labels,n,_ in batches:
        layout=prepare_layout(seg,'cuda','flash')
        with torch.autocast('cuda',dtype=torch.bfloat16):loss=model(ids.cuda(),layout,labels.cuda())*(n/count)
        loss.backward();total+=loss.detach()
    norm=torch.nn.utils.clip_grad_norm_(model.parameters(),1.,error_if_nonfinite=True)
    base_lr=8e-4 if words<30_000_000 else 7e-4 if words<60_000_000 else 6e-4
    lr=base_lr*min(1.,max((words+exposure)/1_600_000,1e-4))
    for group in opt.param_groups:group['lr']=lr
    opt.step()
    assert torch.isfinite(total)
    return dict(loss=float(total),grad_norm=float(norm),lr=lr,words=exposure,targets=count,ffn_dropout=probability)
def rng():return dict(torch=torch.get_rng_state(),cuda=torch.cuda.get_rng_state_all(),numpy=np.random.get_state(),python=random.getstate())
def restore_rng(r):
    torch.set_rng_state(r['torch']);torch.cuda.set_rng_state_all(r['cuda']);np.random.set_state(r['numpy']);random.setstate(r['python'])

class TailAverage:
    """Arithmetic mean of raw post-update weights, separate from training weights."""
    def __init__(self):self.count=0;self.mean={};self.steps=[]
    @torch.no_grad()
    def add(self,model,step):
        self.count+=1;self.steps.append(step)
        for name,value in model.state_dict().items():
            if self.count==1:self.mean[name]=value.detach().clone()
            elif value.is_floating_point():self.mean[name].add_((value-self.mean[name])/self.count)
            else:self.mean[name].copy_(value)
    def state_dict(self):return dict(count=self.count,mean=self.mean,steps=list(self.steps))
    def load_state_dict(self,state,device):
        self.count=state['count'];self.steps=list(state['steps']);self.mean={k:v.to(device).clone() for k,v in state['mean'].items()}

def validate_acceptance(contract):
    for name in ('CPU_ACCEPTANCE.json','GPU_ACCEPTANCE.json'):
        path=ROOT/name
        assert sha(path)==contract['acceptance_sha256'][name],name
        report=json.loads(path.read_text())
        assert report['status']=='PASS',name
        assert report['source_sha256']==contract['source_sha256'],name