"""CPU-only, read-only parameter cloud census with paired uniform references."""
import argparse, hashlib, json, math, os, time
from pathlib import Path
from types import SimpleNamespace
import numpy as np
from pilot.tensors import load_tensors
from finetune_compare.lora import install
from model_shell.groups import enumerate_groups

MAX_GRID=np.array([.5,.75,.9,.95,.99,1.,1.05,1.1,1.25])
RMS_GRID=np.linspace(.5,1.5,41)

def seed(key):return int.from_bytes(hashlib.sha256(key.encode()).digest()[:8],'little')

def guard(start):
    mem=dict((s.split(':')[0],int(s.split()[1])) for s in Path('/proc/meminfo').read_text().splitlines() if s.startswith('MemAvailable:'))
    rss=int(Path('/proc/self/statm').read_text().split()[1])*os.sysconf('SC_PAGE_SIZE')
    if rss>512*1024**2 or mem['MemAvailable']<128*1024 or time.monotonic()-start>7200:
        raise RuntimeError(f'Resource guard: RSS={rss}, available_kB={mem}, elapsed={time.monotonic()-start:.0f}')

def sample(t,g):
    rng=np.random.default_rng(seed(g['id']+'/indices'))
    n=min(512,g['population']);d=min(8192,g['dimension'])
    points=np.sort(rng.choice(g['population'],n,replace=False))
    features=np.sort(rng.choice(g['dimension'],d,replace=False)) if d<g['dimension'] else np.arange(d)
    width=math.prod(t.shape[1:]);out=np.empty((n,d),np.float32)
    if g['orientation']=='rows':
        for j,i in enumerate(points):
            row=i+g['start'];out[j]=t.read(slice(row*width,(row+1)*width))[features]
            if j%32==0:t.evict()
    else:
        for j,i in enumerate(features):
            out[:,j]=t.read(slice(i*width,(i+1)*width))[points]
            if j%32==0:t.evict()
    t.evict();out=out[rng.permutation(n)]
    digest=hashlib.sha256(points.tobytes()+features.tobytes()).hexdigest()
    return out,digest

def participation(y):
    gram=y@y.T;energy=float(np.trace(gram));den=float(np.square(gram).sum())
    return energy**2/den if den else 0.

def geometry(x,ncal,center,raw=False):
    # Float64 arithmetic; bounded sample matrix, no model-wide conversion.
    z=x.astype(np.float64);mu=z[:ncal].mean(0) if center=='mean' else np.zeros(z.shape[1])
    y=z-mu;norm=np.linalg.norm(y,axis=1);cal=norm[:ncal];r=norm[ncal:]
    maximum=float(cal.max());rms=float(np.sqrt(np.mean(cal**2)))
    valid=maximum>0 and rms>0
    a=r/maximum if maximum else np.zeros_like(r);b=r/rms if rms else np.zeros_like(r)
    coordvar=z[:ncal].var(0,ddof=1);cv=float(coordvar.sum()**2/np.square(coordvar).sum()) if coordvar.any() else 0.
    e=np.square(y[ncal:]);den=np.square(e).sum(1)
    pe=np.divide(e.sum(1)**2,den,out=np.zeros(len(e)),where=den>0)
    out=dict(valid=bool(valid),ncal=ncal,neval=len(r),d=z.shape[1],rmax=maximum,rms=rms,
        median_max=float(np.median(a)),median_rms=float(np.median(b)),
        p10_rms=float(np.quantile(b,.1)),p90_rms=float(np.quantile(b,.9)),
        outside=float(np.mean(a>1)),cdf_max=[float(np.mean(a<=v)) for v in MAX_GRID],
        cdf_rms=[float(np.mean(b<=v)) for v in RMS_GRID],
        covariance_participation=participation(y[:ncal]),rank_ceiling=ncal-(center=='mean'),
        coordinate_variance_participation=cv,point_energy_participation_median=float(np.median(pe)),
        zero_fraction=float(np.mean(z==0)))
    if raw:out.update(cal_norm=cal.tolist(),eval_norm=r.tolist())
    return out

BALL_CACHE={}
def references(x,ncal,key):
    d=x.shape[1];cachekey=(d,ncal)
    if cachekey not in BALL_CACHE:
        rng=np.random.default_rng(seed(f'ball/{d}/{ncal}'))
        y=rng.standard_normal((ncal+512,d),dtype=np.float32)
        y*=((rng.random(len(y))**(1/d))/np.linalg.norm(y,axis=1))[:,None]
        BALL_CACHE[cachekey]={c:geometry(y,ncal,c) for c in ('origin','mean')}
    cal=x[:ncal].astype(np.float64);lo=cal.min(0);hi=cal.max(0);mu=cal.mean(0);sd=cal.std(0,ddof=1)
    rng=np.random.default_rng(seed(key+'/uniform'))
    u=rng.random((ncal+512,d),dtype=np.float32)
    refs={'ball':BALL_CACHE[cachekey]}
    for name,left,right in [('range_box',lo,hi),('moment_box',mu-np.sqrt(3)*sd,mu+np.sqrt(3)*sd)]:
        y=(left+(right-left)*u).astype(np.float32)
        refs[name]={c:geometry(y,ncal,c) for c in ('origin','mean')}
    return refs,dict(heldout_outside_coordinate_range=float(np.mean((x[ncal:]<lo)|(x[ncal:]>hi))),
                     range_box_expected_energy=float(np.sum((lo**2+lo*hi+hi**2)/3)),
                     observed_cal_energy=float(np.mean(np.sum(cal**2,axis=1))))

def analyze(x,key):
    ncal=len(x)//2;refs,extra=references(x,ncal,key);out={}
    for center in ('origin','mean'):
        p=geometry(x,ncal,center,True);p['references']=refs; p.update(extra)
        p['comparisons']={}
        for name,bycenter in refs.items():
            ref=bycenter[center];diff=np.array(p['cdf_rms'])-ref['cdf_rms']
            p['comparisons'][name]=dict(interior90_gap=p['cdf_max'][2]-ref['cdf_max'][2],
                interior95_gap=p['cdf_max'][3]-ref['cdf_max'][3],
                rms_cdf_positive_gap=float(max(0,diff.max())),rms_cdf_negative_gap=float(max(0,-diff.min())),
                median_rms_delta=p['median_rms']-ref['median_rms'],
                participation_ratio=p['covariance_participation']/ref['covariance_participation'] if ref['covariance_participation'] else None)
        p['ball_origin_log_expected_cdf90']=float(x.shape[1]*np.log(.9)+np.log(ncal/(ncal+1)))
        out[center]=p
    return out

def main():
    ap=argparse.ArgumentParser();ap.add_argument('--root',default='runs');ap.add_argument('--out',default='runs/model-shell');args=ap.parse_args()
    root=Path(args.root);out=Path(args.out);out.mkdir(exist_ok=True,parents=True);(out/'checkpoints').mkdir(exist_ok=True)
    cfg=json.loads((root/'cpu-qwen35-2b/model_config.json').read_text())['text_config']
    base=load_tensors(root/'cpu-qwen35-2b/weights');ft=SimpleNamespace(tensors=dict(base),config=cfg)
    meta=json.loads((root/'finetune-head-comparison/adapter_meta.json').read_text())
    pairs,scale=install(ft,root/'finetune-head-comparison',meta)
    groups,skips=enumerate_groups(base,cfg)
    manifest=dict(groups=groups,skips=skips,total_tensors=len(base),parameters=sum(t.numel for t in base.values()),changed_tensors=sorted(pairs),adapter_meta=meta,
                  max_grid=MAX_GRID.tolist(),rms_grid=RMS_GRID.tolist(),null_eval=512,point_cap=512,feature_cap=8192)
    (out/'inventory.json').write_text(json.dumps(manifest,indent=2));start=time.monotonic()
    for i,g in enumerate(groups):
        target=out/'checkpoints'/f'{i:04d}.json'
        if target.exists():continue
        guard(start);x,digest=sample(base[g['tensor']],g);p=analyze(x,g['id']);del x
        changed=g['tensor'] in pairs
        if changed:
            guard(start);x,digest2=sample(ft.tensors[g['tensor']],g);assert digest==digest2
            q=analyze(x,g['id']);del x
        else:q=None
        record=dict(group=g,sampling_sha256=digest,changed=changed,original=p,finetuned=q,finetuned_reuses_original=not changed)
        temp=target.with_suffix('.tmp');temp.write_text(json.dumps(record,allow_nan=False));temp.replace(target)
        if i%10==0:print(f'{i+1}/{len(groups)} elapsed={time.monotonic()-start:.1f}s',flush=True)
    print('COMPLETE',len(groups),flush=True)

if __name__=='__main__':main()