"""Source-pinned, CPU-only secondary analysis. No inference or model training.

Run: python code/checkpoint_analysis.py --download --raw-dir raw_sources
Then: python code/checkpoint_analysis.py --raw-dir raw_sources
Only JSON and safetensors are downloaded; no executable model code is loaded.
"""
from __future__ import annotations
import argparse, concurrent.futures, hashlib, json, math, struct, time, urllib.request
from pathlib import Path
import numpy as np

REPO='ModelOrganismsForEM/Qwen2.5-14B-Instruct_R1_0_1_0_extended_train'
REV='1551d514e5977637d5a6fd5198edef4d1d6b80e4'
STEPS=list(range(1,11))+list(range(15,791,5))+[792]
COLS=['step','A_norm','B_norm','effective_norm','A_cos_prev','B_cos_prev','effective_cos_prev','effective_distance_prev','effective_cos_final','B_cos_final','B_corner_k5','effective_corner_k5','B_corner_k10','effective_corner_k10','B_corner_k15','effective_corner_k15']

def sha(b: bytes) -> str: return hashlib.sha256(b).hexdigest()

def download(raw: Path) -> dict:
    paths=[f'checkpoints/checkpoint-{s}/{f}' for s in STEPS for f in ['adapter_model.safetensors','adapter_config.json']]
    paths += ['checkpoints/checkpoint-300/trainer_state.json','checkpoints/checkpoint-792/trainer_state.json']
    def get(path: str) -> dict:
        target=raw/path; target.parent.mkdir(parents=True,exist_ok=True)
        for attempt in range(3):
            try:
                if target.exists(): b=target.read_bytes()
                else:
                    req=urllib.request.Request(f'https://huggingface.co/{REPO}/resolve/{REV}/{path}',headers={'User-Agent':'declaration-study-secondary-analysis/1.0'})
                    with urllib.request.urlopen(req,timeout=40) as r: b=r.read()
                    target.write_bytes(b)
                return {'path':path,'bytes':len(b),'sha256':sha(b)}
            except Exception as ex:
                if attempt==2: return {'path':path,'error':repr(ex)}
                time.sleep(attempt+1)
        raise RuntimeError('unreachable')
    with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool: rows=list(pool.map(get,paths))
    m={'repo':REPO,'revision':REV,'steps':STEPS,'files':rows}
    (raw/'download_manifest.json').write_text(json.dumps(m,indent=2))
    bad=[r for r in rows if 'error' in r]
    if bad: raise RuntimeError(json.dumps(bad))
    return m

def read_ab(path: Path) -> tuple[np.ndarray,np.ndarray]:
    b=path.read_bytes()
    if len(b)<8: raise ValueError(f'Truncated file: {path}')
    n=struct.unpack('<Q',b[:8])[0]
    if n>1000000 or n+8>len(b): raise ValueError(f'Invalid header: {path}')
    header=json.loads(b[8:8+n]); arrays={}
    for key,m in header.items():
        if key=='__metadata__': continue
        if m['dtype']!='F32': raise ValueError(f'Unexpected dtype: {m}')
        start,end=m['data_offsets']; shape=m['shape']
        if start<0 or end<start or end>len(b)-8-n: raise ValueError('Invalid data range')
        if end-start!=math.prod(shape)*4: raise ValueError('Shape-byte mismatch')
        a=np.frombuffer(b[8+n+start:8+n+end],dtype='<f4').copy().astype(np.float64)
        if not np.isfinite(a).all(): raise ValueError('Nonfinite tensor')
        arrays[key]=a
    ak=[k for k in arrays if '.lora_A.' in k];bk=[k for k in arrays if '.lora_B.' in k]
    if len(arrays)!=2 or len(ak)!=1 or len(bk)!=1: raise ValueError('Expected one rank-one adapter pair')
    if '.layers.21.' not in ak[0] or '.layers.21.' not in bk[0]: raise ValueError('Unexpected layer')
    if arrays[ak[0]].size!=13824 or arrays[bk[0]].size!=5120: raise ValueError('Unexpected dimension')
    return arrays[ak[0]],arrays[bk[0]]

def cosine(G: np.ndarray,i: int,j: int) -> float:
    den=np.sqrt(G[i,i]*G[j,j])
    return float(np.clip(G[i,j]/den,-1,1)) if den>0 else float('nan')

def corner(G: np.ndarray,p: int,i: int,n: int) -> float:
    uv=G[p,n]-G[p,i]-G[i,n]+G[i,i]
    u2=G[p,p]+G[i,i]-2*G[p,i];v2=G[n,n]+G[i,i]-2*G[n,i]
    return float(np.clip(uv/np.sqrt(u2*v2),-1,1)) if u2>0 and v2>0 else float('nan')

def lowrank_inner(U,V):
    return float(sum(np.dot(b,c)*np.dot(a,d) for b,a in U for c,d in V))

def analyze(raw: Path,out: Path) -> dict:
    out.mkdir(parents=True,exist_ok=True)
    cfgs=[json.loads((raw/f'checkpoints/checkpoint-{s}/adapter_config.json').read_text()) for s in STEPS]
    if any(c!=cfgs[0] for c in cfgs): raise ValueError('Configurations differ')
    c=cfgs[0]
    if c['r']!=1 or c['lora_alpha']!=64 or c['use_dora'] or c['fan_in_fan_out'] or c['alpha_pattern'] or c['rank_pattern'] or c['layers_to_transform']!=[21]:raise ValueError('Unsupported configuration')
    ab=[read_ab(raw/f'checkpoints/checkpoint-{s}/adapter_model.safetensors') for s in STEPS]
    A=np.stack([p[0] for p in ab]);B=np.stack([p[1] for p in ab]);scale=64.0
    GA=A@A.T;GB=B@B.T;GW=scale**2*GA*GB
    na=np.sqrt(np.diag(GA));nb=np.sqrt(np.diag(GB));nw=np.sqrt(np.diag(GW));idx={s:i for i,s in enumerate(STEPS)}
    rows=[]
    for i,s in enumerate(STEPS):
        row=[s,na[i],nb[i],nw[i]]
        if i:row += [cosine(GA,i-1,i),cosine(GB,i-1,i),cosine(GW,i-1,i),float(np.sqrt(max(0,GW[i,i]+GW[i-1,i-1]-2*GW[i,i-1])))]
        else:row += [float('nan')]*4
        row += [cosine(GW,i,len(STEPS)-1),cosine(GB,i,len(STEPS)-1)]
        for k in [5,10,15]:
            row += [corner(GB,idx[s-k],i,idx[s+k]),corner(GW,idx[s-k],i,idx[s+k])] if s%5==0 and s-k in idx and s+k in idx else [float('nan')]*2
        rows.append(row)
    g=np.asarray(rows,dtype='<f8'); np.save(out/'geometry.npy',g,allow_pickle=False)
    np.savetxt(out/'geometry.csv',g,delimiter=',',header=','.join(COLS),comments='',fmt='%.17g')
    regular=[idx[s] for s in STEPS if s%5==0];Br=B[regular];Bc=Br-Br.mean(axis=0);Gp=Bc@Bc.T
    values,vectors=np.linalg.eigh(Gp);order=np.argsort(values)[::-1];values=np.maximum(values[order],0);vectors=vectors[:,order]
    pc=vectors[:,:2]*np.sqrt(values[:2])
    for j in range(2):
        if pc[np.argmax(np.abs(pc[:,j])),j]<0:pc[:,j]*=-1
    np.save(out/'pca.npy',np.column_stack([[STEPS[i] for i in regular],pc]),allow_pickle=False)
    logs={s:json.loads((raw/f'checkpoints/checkpoint-{s}/trainer_state.json').read_text()) for s in [300,792]}
    def field(log,name):return np.asarray([[v['step'],v[name]] for v in log['log_history'] if name in v],dtype='<f8')
    grad=field(logs[792],'grad_norm');ev=field(logs[792],'eval_loss');oldg=field(logs[300],'grad_norm');olde=field(logs[300],'eval_loss')
    assert grad.shape==(792,2) and ev.shape==(79,2)
    assert np.array_equal(grad[:300],oldg) and np.array_equal(ev[:30],olde)
    np.save(out/'gradient.npy',grad,allow_pickle=False);np.save(out/'eval.npy',ev,allow_pickle=False)
    sensitivity=[]
    for q,k in enumerate([5,10,15]):
        for threshold in [0.,.0001,.0005,.001,.002,.005]:
            selected=[]
            for i,s in enumerate(STEPS):
                if s%5 or s-k not in idx or s+k not in idx:continue
                p,n=idx[s-k],idx[s+k]
                ampA=max(np.linalg.norm(A[p]-A[i]),np.linalg.norm(A[n]-A[i]));ampB=max(np.linalg.norm(B[p]-B[i]),np.linalg.norm(B[n]-B[i]))
                if ampA<threshold and ampB<threshold:continue
                selected.append((s,g[i,10+2*q],g[i,11+2*q]))
            b=max(selected,key=lambda r:r[1]);w=max(selected,key=lambda r:r[2])
            sensitivity.append({'lag':k,'threshold':threshold,'eligible':len(selected),'B_peak_step':b[0],'B_peak':b[1],'effective_peak_step':w[0],'effective_peak':w[2]})
    def difference(j,i):return [(B[j]-B[i],A[j]),(B[i],A[j]-A[i])]
    errors=[]
    for i,s in enumerate(STEPS):
        for q,k in enumerate([5,10,15]):
            if s%5 or s-k not in idx or s+k not in idx:continue
            U,V=difference(idx[s-k],i),difference(idx[s+k],i)
            v=lowrank_inner(U,V)/np.sqrt(lowrank_inner(U,U)*lowrank_inner(V,V))
            errors.append(abs(v-g[i,11+2*q]))
    summary={'repo':REPO,'revision':REV,'config':c,'n_checkpoints':len(STEPS),'n_gradient':len(grad),'n_eval':len(ev),'columns':COLS,'pca_variance_ratios':(values[:5]/values.sum()).tolist(),'zero_effective_update_steps':[STEPS[i] for i in range(len(STEPS)) if nw[i]==0],'effective_norm_monotone':bool(np.all(np.diff(nw)>0)),'B_norm_monotone':bool(np.all(np.diff(nb)>0)),'amplitude_sensitivity':sensitivity,'source_prefix_match':True,'prior_gradient_binary_sha256':sha(oldg.tobytes()),'prior_eval_binary_sha256':sha(olde.tobytes()),'max_stable_vs_gram_corner_error':max(errors),'stable_corner_checks':len(errors),'source_log_sha256':{str(s):sha((raw/f'checkpoints/checkpoint-{s}/trainer_state.json').read_bytes()) for s in [300,792]}}
    (out/'analysis_summary.json').write_text(json.dumps(summary,indent=2,allow_nan=False))
    return summary

if __name__=='__main__':
    p=argparse.ArgumentParser(description=__doc__);p.add_argument('--raw-dir',type=Path,default=Path('raw_sources'));p.add_argument('--output-dir',type=Path,default=Path('recomputed'));p.add_argument('--download',action='store_true');a=p.parse_args()
    if a.download:download(a.raw_dir)
    print(json.dumps(analyze(a.raw_dir,a.output_dir),indent=2))