"""Stasis 1 sixth tranche: 2-GPU DDP continuation from fifth recovery step 4,126.
Launch: torchrun --standalone --nproc_per_node=2 train_sixth.py --config RUN_CONFIG.json
Model code: frozen_stasis (unchanged sha pins). Never AutoModelForCausalLM."""
import argparse, json, math, os, sys, time
import numpy as np, torch, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
sys.path.insert(0,os.path.dirname(os.path.abspath(__file__)))
from frozen_stasis.config import ModelConfig
from frozen_stasis.model import StasisDecoder
import sixth_ckpt as C

W=2049; TARGETS=2048

def lr_at(u,c,N):
    if u<c['warmup_updates']: return c['peak_lr']*(u+1)/c['warmup_updates']
    if N is None: return c['peak_lr']
    p=min(1.0,(u-c['warmup_updates'])/max(1,N-c['warmup_updates']))
    return c['end_lr']+0.5*(c['peak_lr']-c['end_lr'])*(1+math.cos(math.pi*p))

def main():
    ap=argparse.ArgumentParser(); ap.add_argument('--config',required=True); a=ap.parse_args()
    c=json.load(open(a.config))
    backend=c.get('backend','nccl')
    dist.init_process_group(backend)
    rank,world=dist.get_rank(),dist.get_world_size()
    cuda=backend=='nccl'
    dev=torch.device(f'cuda:{rank}') if cuda else torch.device('cpu')
    if cuda: torch.cuda.set_device(dev)
    windows_per_update=c['windows_per_update']; micro=c['microbatch']
    C.need(windows_per_update%(world*micro)==0,'windows/update must split over ranks*microbatch')
    accum=windows_per_update//(world*micro)
    out=c['out_dir']; os.makedirs(out,exist_ok=True)
    log=open(os.path.join(out,'train_log.jsonl'),'a') if rank==0 else None
    def emit(**k):
        if rank==0: log.write(json.dumps(dict(t=round(time.time(),3),**k))+'\n'); log.flush()
    torch.backends.cuda.matmul.allow_tf32=False; torch.backends.cudnn.allow_tf32=False
    model=StasisDecoder(ModelConfig(**c.get('model',{}))).to(dev)
    C.need(sum(p.numel() for p in model.parameters())==109529856 or c.get('allow_test_model'),'parameter count')
    opt,names=C.make_optimizer(model,c['peak_lr'],fused=cuda and c.get('fused_adamw',True))
    scaler=torch.amp.GradScaler('cuda' if cuda else 'cpu',enabled=c.get('scaler',True))
    # ---- restore (strict) ----
    resume=c.get('resume_path')
    src=resume if resume else c['parent_path']
    payload=C.load_payload(src,expected_sha=None if resume else c.get('parent_sha256'),expected_bytes=None if resume else c.get('parent_bytes'))
    step,cursor=C.restore(payload,model,opt,names,scaler,rank=rank)
    if resume:
        C.need(payload['binding']['run_id']==c['run_id'],'resume run_id'); start=cursor['tranche_update']; N=payload['binding'].get('frozen_updates')
        base=payload['binding']['parent_cursor']
    else:
        C.need(step==c['parent_optimizer_step'] and cursor['optimizer_step']==step,'parent optimizer step'); start=0; N=None; base=cursor
    del payload
    for g in opt.param_groups: g['lr']=lr_at(start,c,N)
    ddp=DDP(model,device_ids=[rank] if cuda else None,gradient_as_bucket_view=True,broadcast_buffers=False)
    if c.get('compile'): ddp=torch.compile(ddp)
    data=np.memmap(c['train_path'],dtype='<u2',mode='r'); C.need(len(data)%W==0,'train.bin window multiple')
    data=data.reshape(-1,W); max_updates=min(len(data)//windows_per_update,c.get('max_updates',10**9))
    if N is not None: C.need(N<=max_updates,'frozen N exceeds data')
    stop_at=c['train_stop_epoch']; save_every=c['save_every_seconds']; last_save=time.time()
    binding=dict(schema='stasis.sixth_tranche.binding.v1',run_id=c['run_id'],parent_sha256=c.get('parent_sha256'),train_sha256=c.get('train_sha256'),
                 gates_sha256=c.get('gates_sha256'),schedule=dict(kind='sixth_rewarm_cosine',warmup_updates=c['warmup_updates'],peak_lr=c['peak_lr'],end_lr=c['end_lr']),
                 frozen_updates=N,parent_cursor=base,windows_per_update=windows_per_update,world_size=world,microbatch=micro)
    def cur(u):
        t=u*windows_per_update*TARGETS
        return dict(tranche_update=u,optimizer_step=base['optimizer_step']+u,cumulative_recovery_updates=base['cumulative_recovery_updates']+u,
                    tranche_targets=t,cumulative_recovery_targets=base['cumulative_recovery_targets']+t,lineage_targets=base['lineage_targets']+t,next_tranche_window=u*windows_per_update)
    def save(u,tag):
        if rank==0:
            binding['frozen_updates']=N
            p=C.capture(ddp.module if hasattr(ddp,'module') else model,opt,names,scaler,binding=binding,cursor=cur(u))
            path=os.path.join(out,f'recovery-{tag}.pt'); t0=time.time(); h=C.atomic_save(p,path)
            json.dump(dict(file=os.path.basename(path),sha256=h,bytes=os.path.getsize(path),cursor=cur(u)),open(path+'.manifest.json','w'))
            emit(event='save',tag=tag,update=u,sha256=h,seconds=round(time.time()-t0,1))
        dist.barrier()
    emit(event='start',rank_world=world,start_update=start,max_updates=max_updates,frozen_updates=N,accum=accum,micro=micro,parent_step=step)
    u=start; times=[]; consecutive_skips=0; t_update=time.time()
    raw=ddp.module if hasattr(ddp,'module') else model
    while True:
        if N is not None and u>=N: break
        if u>=max_updates: break
        # stop decision (rank0 time, broadcast)
        flag=torch.tensor([1 if time.time()>=stop_at else 0],device=dev); dist.broadcast(flag,0)
        if flag.item(): emit(event='time_stop',update=u); break
        lr=lr_at(u,c,N)
        for g in opt.param_groups: g['lr']=lr
        base_w=u*windows_per_update+rank*accum*micro
        tot=torch.zeros((),device=dev,dtype=torch.float64)
        for m in range(accum):
            x=torch.from_numpy(data[base_w+m*micro:base_w+(m+1)*micro].astype(np.int64)).to(dev,non_blocking=True)
            ctx=ddp.no_sync() if m<accum-1 else torch.enable_grad()
            with ctx:
                with torch.autocast('cuda' if cuda else 'cpu',dtype=torch.float16 if cuda else torch.bfloat16,enabled=c.get('amp',cuda)):
                    _,loss=ddp(x[:,:-1],x[:,1:])
                scaler.scale(loss/accum).backward()
            tot+=loss.detach().double()
        scaler.unscale_(opt)
        norm=torch.nn.utils.clip_grad_norm_(raw.parameters(),1.0)
        before=scaler.get_scale(); scaler.step(opt); scaler.update(); after=scaler.get_scale()
        opt.zero_grad(set_to_none=True)
        skipped=bool(after<before) if scaler.is_enabled() else False
        dist.all_reduce(tot); mean_loss=float(tot)/(accum*world)
        consecutive_skips=consecutive_skips+1 if skipped else 0
        u+=1; now=time.time(); dt=now-t_update; t_update=now; times.append(dt)
        emit(event='update',update=u,loss=round(mean_loss,6),lr=lr,grad_norm=float(norm),scale=after,skipped=skipped,seconds=round(dt,3))
        C.need(math.isfinite(mean_loss) or skipped,'nonfinite loss without skip')
        if consecutive_skips>=c.get('max_consecutive_skips',25):
            emit(event='abort_skips',update=u); save(u,'abort'); break
        # freeze N after calibration (warmup only so far, LR independent of N)
        if N is None and u==start+c['calibration_updates']:
            spu=torch.tensor([float(np.median(times[3:]))],device=dev); dist.broadcast(spu,0); spu=float(spu.item())
            remaining=stop_at-time.time()-c['reserve_seconds_for_saves']
            N=int(min(max_updates,u+math.floor(remaining/(spu*c['throughput_safety']))))
            C.need(N>c['warmup_updates'],'not enough time for schedule')
            binding['frozen_updates']=N
            if rank==0: json.dump(dict(frozen_updates=N,seconds_per_update_median=spu,targets=N*windows_per_update*TARGETS,frozen_at_update=u,
                                       remaining_seconds=remaining,safety=c['throughput_safety'],utc=time.time()),open(os.path.join(out,'SCHEDULE_FROZEN.json'),'w'),indent=1)
            emit(event='schedule_frozen',N=N,spu=spu)
        sflag=torch.tensor([1 if time.time()-last_save>=save_every else 0],device=dev); dist.broadcast(sflag,0)
        if sflag.item(): save(u,'latest'); last_save=time.time()
    save(u,'final')
    if rank==0:
        mo=dict(schema='stasis.sixth_tranche.model_only.v1',model={k:v.detach().cpu().clone() for k,v in raw.state_dict().items()},cursor=cur(u),binding=binding)
        h=C.atomic_save(mo,os.path.join(out,'endpoint.model_only.pt'))
        json.dump(dict(completed_updates=u,frozen_updates=N,schedule_complete=(N is not None and u>=N),endpoint_sha256=h,cursor=cur(u)),open(os.path.join(out,'TRAIN_STATUS.json'),'w'),indent=1)
    emit(event='done',update=u)
    dist.barrier(); dist.destroy_process_group()

if __name__=='__main__': main()