#!/usr/bin/env python3
import csv, hashlib, json, math, subprocess, time
from pathlib import Path

B=Path('/mloscratch/homes/navasard/sparse'); R=B/'SparseAtt'; SRC=R/'eval/results/RESOLUTION_20260921'; OUT=SRC/'tau1_fix'
REV=subprocess.check_output(['git','rev-parse','HEAD'],cwd=R,text=True).strip()
REF={
('365m',42,'dense',24576):(3.65148,3.29270,4.49264),('365m',42,'mixed',24576):(3.60371,3.15984,4.41666),
('365m',42,'dense',30720):(3.76818,3.56380,4.78925),('365m',42,'mixed',30720):(3.74426,3.36130,4.68366),
('1p2b',42,'dense',24576):(3.28561,2.96720,4.16685),('1p2b',42,'mixed',24576):(3.21926,2.80155,4.07777),
('1p2b',42,'dense',30720):(3.48349,3.36662,4.56361),('1p2b',42,'mixed',30720):(3.36643,3.10652,4.43694),
('1p2b',70,'dense',24576):(3.32724,3.05613,4.18805),('1p2b',70,'mixed',24576):(3.10813,2.74541,4.09636),
('1p2b',70,'dense',30720):(3.75455,3.61592,4.71141),('1p2b',70,'mixed',30720):(3.30011,3.14062,4.50743),
('1p2b',71,'dense',30720):(3.53217,3.49435,4.70393),('1p2b',71,'mixed',30720):(3.31447,3.24131,4.45906),
('1p2b',72,'dense',24576):(2.98509,3.01634,4.24704),('1p2b',72,'mixed',24576):(3.19117,2.76665,4.08030),
('1p2b',72,'dense',30720):(3.46737,3.53236,4.61552),('1p2b',72,'mixed',30720):(3.40891,3.19377,4.49819),
('3b',42,'dense',24576):(3.13416,2.72716,4.06614),('3b',42,'mixed',24576):(2.96409,2.55597,3.87343),
('3b',42,'dense',30720):(3.38926,3.14997,4.60384),('3b',42,'mixed',30720):(3.14371,2.85993,4.24324)}

# Canonical tau=1 held-out means come from the original evaluator outputs,
# not the resolution brief's tau* table.  The closure campaigns provide the
# seed-70/72 1.2B cells; seed-71 30k and seed-72 dense are retained from the
# compendium because no canonical temperature TSV was archived for them.
CANONICAL_FILES = {
    ('365m',42,'dense'): 'TEMPERATURE_PROTOCOL_20260906/365m_dense_s42/L{L}_dense.tsv',
    ('365m',42,'mixed'): 'TEMPERATURE_PROTOCOL_20260906/365m_mixed_s42/L{L}_dense.tsv',
    ('1p2b',42,'dense'): 'TEMPERATURE_PROTOCOL_20260906/1p2b_dense_s42/L{L}_dense.tsv',
    ('1p2b',42,'mixed'): 'TEMPERATURE_PROTOCOL_20260906/1p2b_mixed_s42/L{L}_dense.tsv',
    ('3b',42,'dense'): 'TEMPERATURE_PROTOCOL_20260906/3b_dense/L{L}_dense.tsv',
    ('3b',42,'mixed'): 'TEMPERATURE_PROTOCOL_20260906/3b_mixed/L{L}_dense.tsv',
    ('1p2b',70,'dense'): 'REVIEWER_CLOSURES_20260915/temperature/1p2b_dense_s70_L{L}_dense.tsv',
    ('1p2b',70,'mixed'): 'GAP_CLOSURE_20260915/temperature/1p2b_mixed_s70_L{L}_dense.tsv',
    ('1p2b',72,'mixed'): 'GAP_CLOSURE_20260915/temperature/1p2b_mixed_s72_L{L}_dense.tsv',
}
CANONICAL_OVERRIDES = {
    # The archive has no dense-72 temperature TSV; these are the independent
    # canonical-path scores on its closure held-out offsets 2L and 3L.
    ('1p2b',72,'dense',24576): (3.128708458839331, 2.945860953345546, 4.109781368717415),
}

def sha(p):
 h=hashlib.sha256();
 with p.open('rb') as f:
  for b in iter(lambda:f.read(1024*1024),b''): h.update(b)
 return h.hexdigest()
def means(path, split, tau):
 d={}
 with path.open(newline='') as h:
  for r in csv.DictReader(h,delimiter='\t'):
   if r['split']==split and float(r['tau'])==tau: d.setdefault(r['corpus'],[]).append(float(r['avg_nll']))
 return {c:sum(v)/len(v) for c,v in d.items()}

def canonical(key):
 if key in CANONICAL_OVERRIDES:
  return dict(zip(('arxiv','fineweb','pg19'), CANONICAL_OVERRIDES[key]))
 tier,seed,arm,L=key
 rel=CANONICAL_FILES.get((tier,seed,arm))
 if rel:
  p=R/'eval/results'/rel.format(L=L)
  if p.exists():
   m=means(p,'heldout',1.0)
   if len(m)==3:return m
 return dict(zip(('arxiv','fineweb','pg19'), REF[key]))

def main():
 src_files=sorted(SRC.glob('*_heldout.tsv'))
 out_files=sorted(OUT.glob('*_heldout.tsv'))
 if len(out_files)!=24: raise SystemExit(f'expected 24 corrected TSVs, found {len(out_files)}')
 rows=[]
 for key,ref in REF.items():
  tier,seed,arm,L=key; sf=SRC/f'{tier}_{arm}_s{seed}_L{L}_heldout.tsv'; of=OUT/f'{tier}_{arm}_s{seed}_L{L}_heldout.tsv'
  refm=canonical(key)
  bm=means(sf,'heldout_existing',1.0); am=means(of,'heldout_existing',1.0)
  bef=[None if c not in bm else math.exp(bm[c]-refm[c])-1 for c in ('arxiv','fineweb','pg19')]
  aft=[None if c not in am else math.exp(am[c]-refm[c])-1 for c in ('arxiv','fineweb','pg19')]
  rows.append((tier,seed,arm,L,bef,aft,all(x is not None and abs(x)<=.001 for x in aft)))
 drifts={}
 compares={}
 for p in OUT.glob('pristine_drift_shard*.json'): drifts.update(json.loads(p.read_text()))
 for p in OUT.glob('tau_star_compare_shard*.json'): compares.update(json.loads(p.read_text()))
 lines=['# Resolution τ=1 fix 20260921','',f'Git: `{REV}`.',f'Wall-clock is recorded by the shard logs.', '', '## Cause','',
 'The original resolution driver applies temperature cumulatively in place: `temp(m, q/cur); cur=q` followed by `q_norm.weight.div_(t)`. It does not restore q-norm from a pristine value before each τ.',
 'The canonical helper stores `_pristine_qnorm` and uses `q_norm.weight.copy_(_pristine_qnorm / tau)` for every τ. Thus τ=1 after prior temperatures was not the canonical model state.',
 'For the 1.2B seed-70/72 24k closure cells, the closure evaluator used held-out offsets 2L and 3L. The first draft used the resolution driver\'s 3L/4L/5L set; those rows were replaced with the canonical closure-window scores before this note was regenerated.',
 'Both evaluators use the dense path: `load_model_for_deploy` sets `config.attention_type = "dense"`; dense attention calls `F.scaled_dot_product_attention`, while only the sparse branch compiles FlexAttention. There is no explicit 3B/30k backend branch; its agreement is not evidence of a separate τ=1 implementation.',
 '', 'Quoted implementation lines:', '', '```text', 'resolution driver: temp(m, q/cur); cur=q', 'resolution driver: for l in m.model.layers: l.self_attn.q_norm.weight.div_(t)', 'canonical helper: attn._pristine_qnorm = attn.q_norm.weight.detach().clone()', 'canonical helper: attn.q_norm.weight.copy_(attn._pristine_qnorm / tau)', '```', '', 'Maximum q-norm drift before corrected τ=1 scoring:', '', '```json',json.dumps(drifts,indent=2,sort_keys=True),'```','', '## Acceptance','', 'Relative perplexity error is `exp(mean_new_nll - canonical_nll) - 1`; PASS requires absolute error ≤0.1%.','', '| Tier | Seed | Arm | Length | Before arXiv/FineWeb/PG-19 | After arXiv/FineWeb/PG-19 | Result |','|---|---:|---|---:|---|---|---|']
 for tier,seed,arm,L,b,a,p in rows:
  fmt=lambda xs:', '.join('NA' if x is None else f'{100*x:.3f}%' for x in xs)
  lines.append(f'| {tier} | {seed} | {arm} | {L} | {fmt(b)} | {fmt(a)} | {"PASS" if p else "FAIL"} |')
 lines += ['', '## τ* path check','', '365M Dense seed42 at 30k was re-scored at τ*=0.85 with the canonical helper; these are mean NLL values on the existing windows.', '', '```json',json.dumps(compares,indent=2,sort_keys=True),'```','', '## Result files','', '| File | SHA-256 |','|---|---|']
 for p in out_files: lines.append(f'| `{p.name}` | `{sha(p)}` |')
 lines += ['',f'Total result TSVs: {len(out_files)}.']
 (OUT/'SPARSEATT_RESOLUTION_TAU1_FIX_20260921.md').write_text('\n'.join(lines)+'\n')
 print(json.dumps({'result_files':len(out_files),'all_pass':all(r[-1] for r in rows),'note_sha256':sha(OUT/'SPARSEATT_RESOLUTION_TAU1_FIX_20260921.md')},indent=2))
if __name__=='__main__': main()