"""Validate complete checkpoints; write compact scientific data, figures and report."""
import argparse,csv,gzip,hashlib,json,platform,sys
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from model_shell.analyze import MAX_GRID,RMS_GRID,sample,seed
from pilot.tensors import load_tensors

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 main():
    ap=argparse.ArgumentParser();ap.add_argument('--input',default='runs/model-shell');ap.add_argument('--out',default='results/model-shell');a=ap.parse_args()
    source=Path(a.input);out=Path(a.out);out.mkdir(parents=True,exist_ok=True);figdir=out/'figures';figdir.mkdir(exist_ok=True)
    inventory=json.loads((source/'inventory.json').read_text());groups=inventory['groups'];records=[];rows=[];radii={}
    for i,g in enumerate(groups):
        r=json.loads((source/'checkpoints'/f'{i:04d}.json').read_text());assert r['group']==g
        assert r['changed']==(g['tensor'] in inventory['changed_tensors'])
        for model in ('original','finetuned'):
            profiles=r[model] if r[model] is not None else r['original']
            for center,p in profiles.items():
                if center=='origin':
                    p['ball_log_expected_max_cdf']=[float(p['d']*np.log(z)+np.log(p['ncal']/(p['ncal']+1))) for z in MAX_GRID if z<=1]
                assert p['valid'] and p['d']==min(8192,g['dimension'])
                assert p['covariance_participation']<=p['rank_ceiling']+1e-6
                v=np.asarray(p['eval_norm']);cal=np.asarray(p['cal_norm'])
                assert abs(p['rmax']-cal.max())<1e-12
                assert abs(p['outside']-np.mean(v>p['rmax']))<1e-12
                np.testing.assert_allclose(p['cdf_max'],[np.mean(v<=z*p['rmax']) for z in MAX_GRID])
                np.testing.assert_allclose(p['cdf_rms'],[np.mean(v<=z*p['rms']) for z in RMS_GRID])
                key=f'{i}_{model}_{center}';radii[key+'_cal']=cal;radii[key+'_eval']=v
                row=dict(group_index=i,group=g['id'],model=model,center=center,kind=g['kind'],stratum=g['stratum'],changed=r['changed'],projected=g['dimension']>8192,
                         dimension=g['dimension'],analyzed_dimension=p['d'],population=g['population'],ncal=p['ncal'],neval=p['neval'],sampling_sha256=r['sampling_sha256'])
                for k in ('median_max','median_rms','p10_rms','p90_rms','outside','rmax','rms','covariance_participation','rank_ceiling','coordinate_variance_participation','point_energy_participation_median','zero_fraction','heldout_outside_coordinate_range'):
                    row[k]=p[k]
                for name,c in p['comparisons'].items():
                    for k,vv in c.items():row[name+'_'+k]=vv
                rows.append(row)
        # Compact profiles keep CDF and controls; raw measured radii live in NPZ.
        for model in ('original','finetuned'):
            if r[model] is not None:
                for p in r[model].values():p.pop('cal_norm');p.pop('eval_norm')
        records.append(r)
    with gzip.open(out/'profiles.json.gz','wt',encoding='utf-8') as f:json.dump(records,f,allow_nan=False)
    np.savez_compressed(out/'measured_radii.npz',**radii)
    with (out/'metrics.csv').open('w',newline='') as f:
        w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)
    (out/'inventory.json').write_text(json.dumps(inventory,indent=2))
    def subset(model='original',center='origin',kind=None,changed=None):
        return [r for r in rows if r['model']==model and r['center']==center and r['stratum']=='primary' and not r['projected'] and (kind is None or r['kind']==kind) and (changed is None or r['changed']==changed)]
    summary={}
    for center in ('origin','mean'):
        for kind in ('matrix','head_rows'):
            rs=subset(center=center,kind=kind);s={'n':len(rs),
                'median_zero_fraction':float(np.median([r['zero_fraction'] for r in rs])),
                'median_coordinate_energy_participation_fraction':float(np.median([r['point_energy_participation_median']/r['analyzed_dimension'] for r in rs])),
                'median_coordinate_variance_participation_fraction':float(np.median([r['coordinate_variance_participation']/r['analyzed_dimension'] for r in rs]))}
            for ref in ('ball','range_box','moment_box'):
                gap=np.array([r[ref+'_interior90_gap'] for r in rs]);pos=np.array([r[ref+'_rms_cdf_positive_gap'] for r in rs]);neg=np.array([r[ref+'_rms_cdf_negative_gap'] for r in rs])
                s[ref]=dict(median_interior90_gap=float(np.median(gap)),fraction_gap_gt10pp=float(np.mean(gap>.1)),median_positive_rms_gap=float(np.median(pos)),median_negative_rms_gap=float(np.median(neg)),
                            fraction_two_sided_rms_gt10pp=float(np.mean((pos>.1)&(neg>.1))),median_participation_ratio=float(np.median([r[ref+'_participation_ratio'] for r in rs])))
            summary[center+'/'+kind]=s
    paired={}
    for center in ('origin','mean'):
        b=subset(center=center,changed=True);t=subset(model='finetuned',center=center,changed=True);assert [r['group'] for r in b]==[r['group'] for r in t]
        paired[center]={'n':len(b)}
        for field in ('median_max','median_rms','covariance_participation','moment_box_interior90_gap'):
            delta=np.array([q[field]-p[field] for p,q in zip(b,t)])
            paired[center][field]=dict(median=float(np.median(delta)),p10=float(np.quantile(delta,.1)),p90=float(np.quantile(delta,.9)),max_abs=float(np.max(abs(delta))))
    secondary=[]
    for model in ('original','finetuned'):
        for center in ('origin','mean'):
            for label in ('small_sample','projected'):
                ss=[r for r in rows if r['model']==model and r['center']==center and (r['stratum']=='small_sample' if label=='small_sample' else r['projected'])]
                secondary.append(dict(model=model,center=center,stratum=label,n=len(ss),median_max=float(np.median([r['median_max'] for r in ss])),moment_interior90_gap=float(np.median([r['moment_box_interior90_gap'] for r in ss]))))
    (out/'summary.json').write_text(json.dumps(dict(strata=summary,paired_changed=paired,secondary=secondary),indent=2))
    plt.rcParams.update({'font.size':10,'axes.spines.top':False,'axes.spines.right':False})
    def save(fig,name):
        fig.tight_layout()
        for ext in ('png','pdf','svg'):fig.savefig(figdir/(name+'.'+ext),dpi=180,bbox_inches='tight')
        plt.close(fig)
    fig,axs=plt.subplots(2,2,figsize=(11,7),sharex=True,sharey=True)
    for i,center in enumerate(('origin','mean')):
        for j,kind in enumerate(('matrix','head_rows')):
            rs=subset(center=center,kind=kind);ax=axs[i,j]
            for ref,label in [('ball','Uniform ball'),('range_box','Range-matched box'),('moment_box','Moment-matched box')]:
                v=np.sort([r[ref+'_interior90_gap'] for r in rs]);ax.plot(v,np.arange(1,len(v)+1)/len(v),label=label)
            ax.axvline(0,color='gray',lw=.8);ax.set_title(f'{center} center / {kind}, n={len(rs)}');ax.set_xlabel('F(data r/Rcal <= 0.9) - F(null)');ax.set_ylabel('Fraction of sampled groups')
    axs[0,0].legend(fontsize=8);fig.suptitle('Actual Qwen3.5-2B: interior mass depends on the reference');save(fig,'interior_reference_sensitivity')
    fig,axs=plt.subplots(1,2,figsize=(11,4.3))
    for ax,center in zip(axs,('origin','mean')):
        for kind,marker in [('matrix','o'),('head_rows','.')]:
            rs=subset(center=center,kind=kind);ax.scatter([r['moment_box_rms_cdf_positive_gap'] for r in rs],[r['moment_box_rms_cdf_negative_gap'] for r in rs],s=10,alpha=.25,label=kind,marker=marker)
        ax.axvline(.1,color='gray',lw=.7);ax.axhline(.1,color='gray',lw=.7);ax.set(xlabel='Largest positive CDF gap (more inner tail)',ylabel='Largest negative CDF gap (more outer tail)',title=f'{center} / moment-matched box',xlim=(-.02,1),ylim=(-.02,1));ax.legend()
    fig.suptitle('RMS-normalized radial shape: both tails can grow together');save(fig,'radial_shape_two_tails')
    fig,axs=plt.subplots(1,2,figsize=(11,4.3))
    for ax,center in zip(axs,('origin','mean')):
        b=subset(center=center,changed=True);t=subset(model='finetuned',center=center,changed=True)
        x=[r['median_max'] for r in b];y=[r['median_max'] for r in t];ax.scatter(x,y,s=10,alpha=.3);low=min(x+y);high=max(x+y);ax.plot([low,high],[low,high],color='gray',lw=1)
        ax.set(xlabel='Original median r/Rcal',ylabel='GRPO-LoRA median r/Rcal',title=f'{center}, changed primary groups n={len(b)}')
    fig.suptitle('Paired same-index comparison; unchanged groups excluded');save(fig,'finetune_paired_radii')
    fig,axs=plt.subplots(1,2,figsize=(11,4.3))
    for ax,center in zip(axs,('origin','mean')):
        for kind in ('matrix','head_rows'):
            rs=subset(center=center,kind=kind);v=np.sort([r['moment_box_participation_ratio'] for r in rs]);ax.plot(v,np.arange(1,len(v)+1)/len(v),label=kind)
        ax.set(xlabel='Empirical participation / equal-sample moment-box reference',ylabel='Fraction of groups',title=center);ax.axvline(1,color='gray',lw=.8);ax.legend()
    fig.suptitle('Finite-sample participation is not the true intrinsic dimension');save(fig,'participation_reference')
    examples=[('model.language_model.layers.0.linear_attn.in_proj_qkv.weight::Q_head0','L0 GDN Q head0'),
              ('model.language_model.layers.3.self_attn.q_proj.weight::Q_head0','L3 full-attention Q head0'),
              ('model.language_model.layers.0.mlp.up_proj.weight::rows','L0 MLP up rows')]
    fig,axs=plt.subplots(2,3,figsize=(13,7),sharey=True)
    byid={r['group']['id']:r for r in records}
    indices={g['id']:i for i,g in enumerate(groups)}
    base=load_tensors(Path('runs/cpu-qwen35-2b/weights'));example_radii={}
    for j,(gid,label) in enumerate(examples):
        p=byid[gid]['original']['origin']
        g=byid[gid]['group'];x,digest=sample(base[g['tensor']],g);assert digest==byid[gid]['sampling_sha256']
        ncal=len(x)//2;d=x.shape[1];cal=x[:ncal].astype(np.float64)
        rng=np.random.default_rng(seed(gid+'/uniform'));u=rng.random((ncal+512,d),dtype=np.float32)
        mu=cal.mean(0);sd=cal.std(0,ddof=1);lo=mu-np.sqrt(3)*sd;hi=mu+np.sqrt(3)*sd
        moment=(lo+(hi-lo)*u).astype(np.float32)
        rng=np.random.default_rng(seed(f'ball/{d}/{ncal}'));ball=rng.standard_normal((ncal+512,d),dtype=np.float32)
        ball*=((rng.random(len(ball))**(1/d))/np.linalg.norm(ball,axis=1))[:,None]
        arrays={'Actual weights':(radii[f'{indices[gid]}_original_origin_cal'],radii[f'{indices[gid]}_original_origin_eval'])}
        for name,ref,raw in [('Uniform ball','ball',ball),('Moment box','moment_box',moment)]:
            nr=np.linalg.norm(raw.astype(np.float64),axis=1);arrays[name]=(nr[:ncal],nr[ncal:])
            np.testing.assert_allclose([np.mean(nr[ncal:]/nr[:ncal].max()<=z) for z in MAX_GRID],p['references'][ref]['origin']['cdf_max'])
            np.testing.assert_allclose([np.mean(nr[ncal:]/np.sqrt(np.mean(nr[:ncal]**2))<=z) for z in RMS_GRID],p['references'][ref]['origin']['cdf_rms'])
        for name,(nc,ne) in arrays.items():example_radii[f'example{j}_{name.replace(" ","_")}_cal']=nc;example_radii[f'example{j}_{name.replace(" ","_")}_eval']=ne
        for i,(field,grid) in enumerate([('cdf_max',MAX_GRID),('cdf_rms',RMS_GRID)]):
            ax=axs[i,j]
            for name,(nc,ne) in arrays.items():
                scale=nc.max() if i==0 else np.sqrt(np.mean(nc**2));v=np.sort(ne/scale)
                ax.step(np.r_[0,v, max(1.5,float(v[-1]))],np.r_[0,np.arange(1,len(v)+1)/len(v),1],where='post',label=name,lw=1.8,ls='-' if name=='Actual weights' else '--')
            ax.set(title=label,xlabel='r / calibration '+('maximum' if i==0 else 'RMS'),ylabel='Held-out CDF',ylim=(-.03,1.03),xlim=(0,1.5));ax.legend(fontsize=8)
    np.savez_compressed(out/'example_radii.npz',**example_radii)
    fig.suptitle('Fixed role examples: maximum-radius and RMS normalization');save(fig,'fixed_role_radial_cdfs')
    lines=['# 两个实际2B检查点：参数云的球壳检验','',f'完整覆盖：{len(groups)}组/模型，{len(rows)}个模型×中心组合。632个参数张量，2,274,069,824个参数（包含语言、视觉和MTP）。',
           '','## 分布计数与范围','', '固定规则给出554个矩阵行／列视图、1,004个语言/MTP头视图；视图互相重叠，不能相加成参数节省量。1,522组有至少64个可用点，36组只有16个点；1组高维嵌入列视图使用8,192个抽样坐标。主表只汇总样本充足且没有投影的1,521组。319个一维张量和18个低维卷积核不构成此规则下的高维样本云。',
           '','## 原模型结果（逐组等权，非独立观测）','', '|中心|视图|组数|相对均匀球：90%半径内质量差中位数|相对矩匹配均匀盒：质量差中位数|矩匹配且RMS归一后，内外两尾差均>10个百分点的组占比|','|---|---|---:|---:|---:|---:|']
    for key,s in summary.items():
        c,k=key.split('/');lines.append(f'|{c}|{k}|{s["n"]}|{s["ball"]["median_interior90_gap"]:.3f}|{s["moment_box"]["median_interior90_gap"]:.3f}|{s["moment_box"]["fraction_two_sided_rms_gt10pp"]:.1%}|')
    lines+=['','这些是描述性质量差，不是显著性检验。正差表示所指定半径以内有更多样本；不能将其自动解释为整条径向分布向内平移。若RMS归一后内外尾均变多，更准确的解释是径向分布更宽。',
            '', '|中心/视图|逐组零值率中位数|点坐标能量参与率÷维度，中位数|坐标方差参与率÷维度，中位数|','|---|---:|---:|---:|']
    for key,s in summary.items():lines.append(f'|{key}|{s["median_zero_fraction"]:.6g}|{s["median_coordinate_energy_participation_fraction"]:.4f}|{s["median_coordinate_variance_participation_fraction"]:.4f}|')
    lines+=['','能量参与率衡量坐标能量集中程度，依赖坐标基；有限样本Gram参与率反映样本间几何，另见图。两者均不是精确非零坐标计数，也不是模型功能的有效维度。',
            '','## 微调前后','', '原模型是官方Qwen3.5-2B；第二个模型是在同一revision上进行GSM8K GRPO的已有纯LoRA检查点，不是另一个完整模型。72个投影矩阵使用只读W+2BA展开，其余参数共享。972个视图受影响；其中936个属于主分析。未改变视图复用完全相同的分析结果。','', '|中心|指标|配对变化中位数|P10|P90|最大绝对变化|','|---|---|---:|---:|---:|---:|']
    for c,p in paired.items():
        for k,v in p.items():
            if k!='n':lines.append(f'|{c}|{k}|{v["median"]:.6g}|{v["p10"]:.6g}|{v["p90"]:.6g}|{v["max_abs"]:.6g}|')
    lines+=['','## 单独描述的小样本与坐标抽样组','', '|模型|中心|类别|组数|中位r/Rcal的组间中位数|相对矩匹配盒的90%半径质量差中位数|','|---|---|---|---:|---:|---:|']
    for s in secondary:lines.append(f'|{s["model"]}|{s["center"]}|{s["stratum"]}|{s["n"]}|{s["median_max"]:.4f}|{s["moment_interior90_gap"]:.4f}|')
    lines+=['','## 如何读图','', '1. `interior_reference_sensitivity`：不同均匀参照下，90%校准半径内质量差的组间CDF。横轴向右表示更多内侧质量。','2. `radial_shape_two_tails`：匹配矩且按各自校准RMS归一后的最大正／负CDF差。右上角同时增加两尾，并非单纯向内。','3. `finetune_paired_radii`：同一视图的微调前后中位半径；对角线表示不变。','4. `participation_reference`：Gram矩阵参与率相对同样本量参照的比值。受样本秩上限影响，不能当作真实内在维度。',
            '5. `fixed_role_radial_cdfs`：预先按角色固定的GDN Q、全注意力Q和MLP上投影示例，上排按校准最大半径、下排按RMS归一化。不是按结果挑选最明显的组。',
            '','## 方法与限制','', '每组最多512个不同参数向量，一半校准、一半留出；模型固定，留出指未用于拟合中心/半径的参数向量，不是独立训练模型。对照每组512个验证向量，最坏单项二项标准误约2.21个百分点；模型留出样本的误差还取决于样本数，且参数向量不保证独立。10个百分点仅是描述性筛选尺度，没有多重检验或总体显著性结论。',
            '','均匀球、坐标范围均匀盒、匹配均值和逐坐标方差的均匀盒都按同一规则拟合自身中心、最大半径及RMS。均匀球原点中心的精确拟合半径CDF为a^d*ncal/(ncal+1)，保留其对数以避免把模拟零命中误报成零概率。盒对照不是协方差匹配。校准最大半径不是已知总体支撑边界；留出越界率已记录。',
            '','`origin`保留坐标均值贡献；`mean`使用校准向量均值。没有使用所有数据来选边界、训练或推理；没有删除模型。坐标稀疏性、有限样本参与率与径向内侧质量分别测量，不能相互代替。该实验不能证明更好的模型更靠里，也不能证明功能重复、可剪枝或可复用。',
            '','复现：`OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 venv/bin/python -m model_shell.analyze`，完成后 `... python -m model_shell.report`。只在已授权CPU服务器运行。分组协议见 `model_shell/PROTOCOL.md`；全量指标在metrics.csv，完整CDF/参照在profiles.json.gz，实测范数在measured_radii.npz。']
    (out/'REPORT.md').write_text('\n'.join(lines)+'\n')
    sources={'base':Path('runs/cpu-qwen35-2b/weights/model.safetensors-00001-of-00001.safetensors'),'adapter':Path('runs/finetune-head-comparison/adapter.safetensors'),
             'config':Path('runs/cpu-qwen35-2b/model_config.json'),'adapter_meta':Path('runs/finetune-head-comparison/adapter_meta.json')}
    sourcehash={k:dict(path=str(p),bytes=p.stat().st_size,sha256=sha(p)) for k,p in sources.items()}
    (out/'environment.json').write_text(json.dumps(dict(python=sys.version,platform=platform.platform(),numpy=np.__version__,matplotlib=matplotlib.__version__,
        source_code_sha256={str(p):sha(p) for p in sorted(Path('model_shell').glob('*.py'))},protocol_sha256=sha(Path('model_shell/PROTOCOL.md')),seed='first 8 SHA256 bytes, little endian; group id plus purpose suffix',
        thread_limit=1,server='8.211.154.134',training=False,inference=False,gpu=False),indent=2))
    assert sourcehash['base']['sha256']=='aa33250c4fc64891ddfaba3a314fd9542ea371843c387178b425fbcc5ed680b1'
    assert sourcehash['adapter']['sha256']=='93a3f43a4356f84da0d13ecc622fee13e41ba15d237cdaad6094f243899c62d6'
    audit=dict(complete_groups=len(groups),profile_rows=len(rows),all_finite=True,cdf_recomputed=True,sample_rank_checked=True,sources=sourcehash,files={str(p.relative_to(out)):sha(p) for p in sorted(out.rglob('*')) if p.is_file() and p.name!='audit.json'})
    (out/'audit.json').write_text(json.dumps(audit,indent=2));print(json.dumps(dict(summary=summary,paired=paired),indent=2))

if __name__=='__main__':main()