"""Shared implementation for a future frozen broad learning comparison."""
import collections
import json
import random
from pathlib import Path
import torch
from batched_latent_acquisition import answer_loss,collate,bounded_head_ce
from compressed_checkpoint import save_compressed,load_compressed
from extract_latent_teacher import sha
from latent_deliberation import LatentDeliberation,feature_loss
from microbatch_latent_acquisition import partition_rows
from run_baseline import FastModel,QWEN_CONFIG_06_B,Qwen3Tokenizer
from train_latent_acquisition_pilot import prepare

ROOT=Path(__file__).resolve().parents[1]


def validate(c):
    assert c['status']=='FROZEN_BROAD_OWN_FEATURE_LEARNING'
    for path,digest in c['input_sha256'].items():assert sha(ROOT/path)==digest,path
    assert c['trainable_parameters']==751632384 and c['slots']==32


def initialize(seed):
    torch.set_num_threads(4);torch.manual_seed(seed)
    torch.backends.cuda.matmul.allow_tf32=False;torch.backends.cudnn.allow_tf32=False


def make_model(c,arm):
    base=FastModel(dict(QWEN_CONFIG_06_B,context_length=4096,dtype=torch.float32)).cuda()
    model=LatentDeliberation(base,c['slots'] if arm in ('L','F') else 0,True)
    assert sum(p.numel() for p in model.parameters() if p.requires_grad)==c['trainable_parameters']
    return model


def optimizer(model,c,stage):
    return torch.optim.AdamW(model.parameters(),lr=c[stage+'_learning_rate'],
                            weight_decay=c['weight_decay'],foreach=False)


def data(c):
    tok=Qwen3Tokenizer(tokenizer_file_path=ROOT/'code/data/tokenizer-base.json')
    fitting=prepare(ROOT/c['fitting_source'],c['fit_questions'],tok)
    hold=prepare(ROOT/c['holdout_source'],c['hold_questions'],tok)
    saved=torch.load(ROOT/c['teacher_features'],map_location='cpu',weights_only=True)
    teacher={int(i):f for i,f in zip(saved['indices'],saved['features'])}
    assert not {r['index'] for r in fitting}&{r['index'] for r in hold}
    assert all(r['index'] in teacher for r in fitting+hold)
    assert all(len(r['prompt'])<=960 and len(r['completion'])<=1024 and
               len(r['split']['canonical_answer'])<=64 for r in fitting+hold)
    boundaries={tuple(r['split']['separator']) for r in fitting+hold}
    assert len(boundaries)==1 and list(next(iter(boundaries)))==c['boundary']
    return tok,fitting,hold,teacher


def sequence(n,epochs,batch,seed):
    result=[]
    for epoch in range(epochs):
        order=list(range(n));random.Random(seed+epoch).shuffle(order)
        result.extend(order[j:j+batch] for j in range(0,n,batch))
    return result


def group_work(rows,slots,mode):
    n=len(rows);p=max(len(r['prompt']) for r in rows);t=max(len(r['completion']) for r in rows)
    work={'rows':n,'explicit_trunk_positions':n*(p+t-1),'explicit_head_positions':n*t,
          'explicit_valid_targets':sum(len(r['completion']) for r in rows),
          'explicit_dense_attention_pairs':n*(p+t-1)**2}
    if mode!='O':
        a=max(len(r['split']['canonical_answer']) for r in rows);d=len(rows[0]['split']['separator'])
        pairs=p*p+sum(p+j+1 for j in range(slots))+d*(p+slots+d)
        if a>1:pairs+=(a-1)*(p+slots+d+a-1)
        work.update(answer_trunk_positions=n*(p+slots+d+a-1),answer_head_positions=n*a,
                    answer_valid_targets=sum(len(r['split']['canonical_answer']) for r in rows),
                    answer_dense_attention_pairs=n*pairs,latent_positions=n*slots)
    return work


@torch.no_grad()
def audit_source(model,rows,teacher,*,pad,chunk):
    model.eval();records=[];work=collections.Counter()
    for start in range(0,len(rows),4):
        for group in partition_rows(rows[start:start+4]):
            prompts=[r['prompt'] for r in group]
            targets=[r['split']['canonical_answer'] for r in group]
            _,features,tokens,valid=answer_loss(model,prompts,targets,group[0]['split']['separator'],pad=pad,chunk=chunk)
            target=torch.stack([teacher[r['index']] for r in group]).to(features.device)
            scale=target.std(-1,correction=0,keepdim=True).clamp_min(1e-6)
            alignment=(torch.nn.functional.smooth_l1_loss(features[:,1:],target,reduction='none')/scale).mean((1,2))
            p,mask,t,tvalid=collate(prompts,[r['completion'] for r in group],pad,features.device)
            ids=torch.cat((p,t),1)[:,:-1]
            mask=torch.cat((mask,torch.ones_like(t[:,:-1],dtype=torch.bool)),1)
            h,_=model.hidden(ids=ids,mask=mask)
            _,text_tokens=bounded_head_ce(model,h[:,p.shape[1]-1:],t,tvalid,chunk)
            for i,row in enumerate(group):
                token=tokens[i,:len(targets[i])]
                body=torch.tensor(row['split']['canonical_content_mask'],dtype=torch.bool,device=token.device)
                metrics={'index':row['index'],'answer_ce':float(token.mean()),
                         'body_overlap_ce':float(token[body].mean()),'syntax_eos_ce':float(token[~body].mean()),
                         'text_ce':float(text_tokens[i].sum()/tvalid[i].sum()),'feature_loss':float(alignment[i])}
                assert all(__import__('math').isfinite(v) for k,v in metrics.items() if k!='index')
                records.append(metrics)
            work.update(group_work(group,model.slots,'F'))
    assert [r['index'] for r in records]==[r['index'] for r in rows]
    means={k:sum(r[k] for r in records)/len(records) for k in records[0] if k!='index'}
    return {'rows':records,'means':means,'forward_work':dict(work)}


def workspace_bytes():
    # Do not traverse shared-data symlinks. Count allocated files in our workspace.
    import os
    total=0
    for directory,dirs,files in os.walk(ROOT,followlinks=False):
        dirs[:]=[d for d in dirs if not (Path(directory)/d).is_symlink()]
        for name in files:
            p=Path(directory)/name
            if not p.is_symlink():total+=p.stat().st_size
    return total


def save_state(payload,path,*,final):
    # Conservative bounds allow slightly more than even an uncompressed archive.
    bound=3_100_000_000 if final else 9_100_000_000
    if workspace_bytes()+bound>49_800_000_000:
        raise RuntimeError('Checkpoint peak would exceed the50GB workspace budget')
    return save_compressed(payload,path)


def rotate_restart(folder,previous,summary):
    if previous is None:return
    old=Path(previous['checkpoint']['path'])
    assert old.is_relative_to(folder) and old.name.startswith('restart_') and old.exists()
    with (folder/'retention.jsonl').open('a') as log:
        log.write(json.dumps({'removed':str(old),'old_checkpoint':previous['checkpoint'],
                              'successor_chunk':summary['chunk'],'successor_checkpoint':summary['checkpoint'],
                              'reason':'Verified successor preserves resume or final endpoint; intermediate optimizer state disposable.'})+'\n')
    old.unlink()
