"""Fixed source-only D/L/F acquisition pilot. Never samples native answers."""
import argparse
import hashlib
import json
import random
import time
from pathlib import Path

import torch
from torch.nn import functional as F

from audit_latent_source_boundary import split_trace
from extract_latent_teacher import sha
from latent_deliberation import LatentDeliberation, feature_loss
from run_baseline import FastModel, QWEN_CONFIG_06_B, Qwen3Tokenizer

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


def prepare(path, count, tok):
    rows = [json.loads(line) for line in path.read_text().splitlines()][:count]
    assert len(rows) == count
    return [dict(row, split=split_trace(row, tok)) for row in rows]


def components(model, row, teacher):
    device = next(model.parameters()).device
    prompt = torch.tensor([row['prompt']], device=device)
    boundary = torch.tensor([row['split']['separator']], device=device)
    target = torch.tensor([row['split']['canonical_answer']], device=device)
    ce, states, logits = model.answer_loss(prompt, boundary, target, features=True)
    alignment = feature_loss(states[:, 1:], teacher[None].to(device))
    # Diagnostics do not add a second differentiable CE path to the objective.
    with torch.no_grad():
        per_token = F.cross_entropy(logits.flatten(0,1), target.flatten(), reduction='none')
        body = torch.tensor(row['split']['canonical_content_mask'], device=device)
        metrics = {'answer_ce': float(ce), 'feature_loss': float(alignment),
                   'body_overlap_ce': float(per_token[body].mean()),
                   'syntax_eos_ce': float(per_token[~body].mean())}
    return ce, alignment, metrics


def text_loss(model, row):
    device = next(model.parameters()).device
    return model.text_loss(torch.tensor([row['prompt']], device=device),
                           torch.tensor([row['completion']], device=device))


@torch.no_grad()
def audit(model, rows, teacher):
    model.eval()
    records = []
    for row in rows:
        _, _, metrics = components(model, row, teacher[row['index']])
        metrics['text_ce'] = float(text_loss(model, row))
        metrics['index'] = row['index']
        assert all(__import__('math').isfinite(value) for key,value in metrics.items() if key != 'index')
        records.append(metrics)
    means = {key: sum(row[key] for row in records)/len(records)
             for key in records[0] if key != 'index'}
    return {'means': means, 'rows': records}


def work_for(row, slots):
    p = len(row['prompt']); d = len(row['split']['separator'])
    a = len(row['split']['canonical_answer']); t = len(row['completion'])
    # Attention uses dense Q-by-K tensors even where the causal mask disallows pairs.
    answer_positions = p+slots+d+a-1
    answer_pairs = p*p + sum(p+j+1 for j in range(slots)) + d*(p+slots+d)
    if a > 1:
        answer_pairs += (a-1)*(p+slots+d+a-1)
    text_positions = p+t-1
    return {'answer_trunk_positions':answer_positions, 'answer_head_positions':a,
            'answer_attention_pairs':answer_pairs, 'explicit_trunk_positions':text_positions,
            'explicit_head_positions':t, 'explicit_attention_pairs':text_positions**2,
            'latent_positions':slots}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--contract',required=True)
    ap.add_argument('--arm',choices=['D','L','F'],required=True)
    args = ap.parse_args()
    cp=ROOT/args.contract; c=json.loads(cp.read_text())
    assert c['status']=='FROZEN_SOURCE_ONLY_LATENT_ACQUISITION_PILOT'
    for path,digest in c['input_sha256'].items():
        assert sha(ROOT/path)==digest,path
    assert c['fit_questions']==128 and c['hold_questions']==64 and c['epochs']==2 and c['batch']==4
    torch.set_num_threads(4);torch.manual_seed(c['seed'])
    torch.backends.cuda.matmul.allow_tf32=False;torch.backends.cudnn.allow_tf32=False
    start=time.monotonic()
    out=ROOT/c['output']/args.arm;out.mkdir(parents=True,exist_ok=False)
    tok=Qwen3Tokenizer(tokenizer_file_path=ROOT/'code/data/tokenizer-base.json')
    fit=prepare(ROOT/c['fitting_source'],128,tok)
    hold=prepare(ROOT/c['holdout_source'],64,tok)
    assert not {r['index'] for r in fit}&{r['index'] for r in hold}
    saved=torch.load(ROOT/c['teacher_features'],map_location='cpu',weights_only=True)
    teacher={index:features for index,features in zip(saved['indices'],saved['features'])}
    assert all(r['index'] in teacher for r in fit+hold)
    base=FastModel(dict(QWEN_CONFIG_06_B,context_length=4096,dtype=torch.float32)).cuda()
    base.load_state_dict(torch.load(ROOT/c['parent'],map_location='cpu',weights_only=True))
    slots=0 if args.arm=='D' else c['slots']
    model=LatentDeliberation(base,slots,checkpoint_blocks=True)
    assert sum(p.numel() for p in model.parameters() if p.requires_grad)==751632384
    before={'fit':audit(model,fit,teacher),'hold':audit(model,hold,teacher)}
    (out/'before.json').write_text(json.dumps(before,indent=2)+'\n')
    opt=torch.optim.AdamW(model.parameters(),lr=c['learning_rate'],weight_decay=c['weight_decay'],foreach=False)
    work={'training':{},'audit':{},'training_rows':0,'optimizer_updates':0,
          'checkpointing':'Each differentiable transformer block recomputes in backward; report separately.'}
    sequence=[]
    for epoch in range(2):
        order=list(range(128));random.Random(c['seed']+epoch).shuffle(order)
        sequence.extend(order[j:j+4] for j in range(0,128,4))
    torch.cuda.reset_peak_memory_stats();model.train()
    with (out/'updates.jsonl').open('x') as log:
        for update,batch in enumerate(sequence,1):
            if time.monotonic()-start>2700:
                raise TimeoutError('Pilot timing bound exceeded; retain logs, no success assertion')
            opt.zero_grad(set_to_none=True); records=[]
            for j in batch:
                row=fit[j]
                ce,alignment,metrics=components(model,row,teacher[row['index']])
                objective=ce+c['alignment_weight']*alignment if args.arm=='F' else ce
                assert torch.isfinite(objective)
                (objective/4).backward()
                del ce,alignment,objective
                explicit=text_loss(model,row)
                assert torch.isfinite(explicit)
                (explicit/4).backward()
                metrics.update(index=row['index'],text_ce=float(explicit.detach()))
                records.append(metrics);del explicit
                for key,value in work_for(row,slots).items():
                    work['training'][key]=work['training'].get(key,0)+value
                work['training_rows']+=1
            norm=torch.nn.utils.clip_grad_norm_(model.parameters(),c['clip'],error_if_nonfinite=True)
            opt.step();work['optimizer_updates']+=1
            record={'update':update,'rows':records,'preclip_norm':float(norm),'seconds':time.monotonic()-start}
            log.write(json.dumps(record)+'\n');log.flush()
            if update%8==0:
                print(json.dumps({'arm':args.arm,'update':update,'seconds':time.monotonic()-start}),flush=True)
    opt.zero_grad(set_to_none=True)
    assert all(torch.isfinite(p).all() for p in model.parameters())
    after={'fit':audit(model,fit,teacher),'hold':audit(model,hold,teacher)}
    for row in fit+hold:
        for key,value in work_for(row,slots).items():
            work['audit'][key]=work['audit'].get(key,0)+2*value
    # Final model only. No optimizer history is needed after a completed fixed pilot.
    torch.save({'model':{name:p.detach().cpu() for name,p in base.state_dict().items()},
                'contract_sha256':sha(cp),'arm':args.arm,'updates':64},out/'final.pt')
    result={'status':'COMPLETE','scope':'source-only acquisition, no native sample or grade',
            'arm':args.arm,'slots':slots,'contract_sha256':sha(cp),'before':before,'after':after,
            'work':work,'checkpoint_sha256':sha(out/'final.pt'),
            'peak_allocated_bytes':torch.cuda.max_memory_allocated(),
            'seconds_including_setup_audits_save':time.monotonic()-start}
    (out/'summary.json').write_text(json.dumps(result,indent=2)+'\n')
    print(json.dumps({key:value for key,value in result.items() if key not in ('before','after')},indent=2))


if __name__=='__main__':
    main()
