"""Saved-data scientific analysis; no model invocation or fitting."""
import csv,json
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from .core import concentration


def stats(x):
    x=np.asarray(x,dtype=np.float64);valid=x[np.isfinite(x)]
    if not valid.size:return {'count':0,'p10':None,'median':None,'p90':None,'mean':None}
    q=np.quantile(valid,[.1,.5,.9]);return {'count':int(valid.size),'p10':float(q[0]),'median':float(q[1]),'p90':float(q[2]),'mean':float(valid.mean())}


def csvwrite(p,rows):
    with p.open('w') as f:
        w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)


def main(root):
    root=Path(root);figdir=root/'figures';figdir.mkdir(exist_ok=True);status=json.loads((root/'status.json').read_text());assert status['state']=='completed'
    rows=[];joint_rows=[];pool={};curves={};joint_pool={};obs={};winner=[];joint_by_layer={}
    for entry in status['completed_links']:
        sample=entry['sample'];source=entry['source'];data=dict(np.load(root/'raw'/f'{sample}-link-{source:02d}.npz'))
        if 'single_JS' in data:
            data['single_sqrtJS']=np.sqrt(data['single_JS'])
            for label in ('weak_half','strong_half','random_half','all_heads'):
                data['joint_'+label+'_sqrtJS']=np.sqrt(data['joint_'+label+'_JS'])
                masks=data.get('mask_'+label,np.ones_like(data['mask_weak_half']))
                data['individual_sum_'+label+'_sqrtJS']=np.sum(data['single_sqrtJS'].transpose(2,1,0)*masks,-1).T
        for key in ('Q','K','V','QKV','JS','sqrtJS'):
            if 'single_'+key not in data:continue
            d=data['single_'+key][:,7:,:].transpose(1,2,0);c=concentration(d);valid=c['valid'];h=d.shape[-1]
            metric={k:v.reshape(-1) for k,v in c.items() if k!='valid'}
            row={'sample':sample,'source':source,'recipient':source+1,'metric':key,'source_heads':h,'recipient_attention_type':'full_attention' if entry['full_attention_recipient'] else 'linear_attention','observations':d.shape[0]*d.shape[1],'undefined_zero_total':int((~valid).sum())}
            for k,v in metric.items():
                pool.setdefault(key,{}).setdefault(k,[]).extend(v)
                for s,z in stats(v).items():row[k+'_'+s]=z
            rows.append(row)
            # Raw per-observation statistics for independently reproducing figures.
            for k,v in metric.items():obs.setdefault(key+'_'+k,[]).append(v.astype(np.float32))
            n=d.shape[0]*d.shape[1]
            for k,v in [('source',np.full(n,source)),('sample_index',np.full(n,sorted({e['sample'] for e in status['completed_links']}).index(sample))),('token',np.repeat(np.arange(7,32),d.shape[1])),('recipient_head',np.tile(np.arange(d.shape[1]),25))]:obs.setdefault(key+'_'+k,[]).append(v.astype(np.int16))
            p=np.divide(d,d.sum(-1,keepdims=True),out=np.zeros_like(d),where=d.sum(-1,keepdims=True)>0);cum=np.concatenate([np.zeros((*d.shape[:2],1)),np.cumsum(np.sort(p,-1)[...,::-1],-1)],-1)
            grid=np.linspace(0,1,101);cc=np.array([np.interp(grid,np.linspace(0,1,h+1),v) for v in cum[valid]])
            curves.setdefault(key,[]).extend(cc)
            win=np.argmax(d,-1);counts=np.bincount(win[valid],minlength=h)
            winner.append({'sample':sample,'source':source,'metric':key,'most_frequent_winner_fraction':float(counts.max()/max(counts.sum(),1)),'distinct_winning_heads':int((counts>0).sum()),'head_counts':json.dumps(counts.tolist())})
            for label in ('weak_half','strong_half','random_half','all_heads'):
                values=data['joint_'+label+'_'+key][7:];summed=data['individual_sum_'+label+'_'+key][7:]
                jr={'sample':sample,'source':source,'recipient':source+1,'metric':key,'removed':label}
                for k,x in [('joint_effect',values),('sum_individual',summed)]:
                    for s,z in stats(x).items():jr[k+'_'+s]=z
                joint_rows.append(jr);joint_pool.setdefault(key,{}).setdefault(label,[]).extend(values.ravel());joint_by_layer.setdefault((key,label,source),[]).extend(values.ravel())
    csvwrite(root/'concentration-by-link.csv',rows);csvwrite(root/'joint-by-link.csv',joint_rows);csvwrite(root/'winner-identities.csv',winner)
    np.savez_compressed(root/'observations.npz',**{k:np.concatenate(v) for k,v in obs.items()})
    pooled={k:{m:stats(v) for m,v in metrics.items()} for k,metrics in pool.items()};joint_stats={k:{m:stats(v) for m,v in metrics.items()} for k,metrics in joint_pool.items()}
    paired={}
    for key,labels in joint_pool.items():
        weak=np.asarray(labels['weak_half']);strong=np.asarray(labels['strong_half']);random=np.asarray(labels['random_half']);valid=weak>0
        paired[key]={'weak_less_than_strong_fraction':float(np.mean(weak<strong)),'weak_less_than_random_fraction':float(np.mean(weak<random)),'zero_weak_effect_count':int((~valid).sum()),'strong_over_weak':stats(strong[valid]/weak[valid]),'random_over_weak':stats(random[valid]/weak[valid])}
    full_sources=sorted({x['source'] for x in status['completed_links'] if x['full_attention_recipient']})
    by_type={}
    for key in pool:
        src=np.concatenate(obs[key+'_source']);by_type[key]={}
        for label,choose in [('full_attention',np.isin(src,full_sources)),('linear_attention',~np.isin(src,full_sources))]:
            if not choose.any():continue
            by_type[key][label]={m:stats(np.concatenate(obs[key+'_'+m])[choose]) for m in pool[key]}
    summary={'scale_sensitivity':'sqrt(JS) added after reviewing completed primary summaries; primary JS and all raw interventions unchanged; squared divergences cannot be equated to probability mass or linear-scale effects','by_recipient_type':by_type,'full_attention_source_indices':full_sources,'paired_joint':paired,'pooling':'query7..31, per text/token/recipient-query-head; descriptive dependent observations, not independent samples','concentration':pooled,'joint':joint_stats,'counts':{'links':len(status['completed_links']),'texts':len(json.loads((root/'samples.json').read_text())),'layer_transitions':23},'sample_index_key':sorted({e['sample'] for e in status['completed_links']}),'winner_identity_file':'winner-identities.csv'}
    (root/'analysis-summary.json').write_text(json.dumps(summary,indent=2,allow_nan=False)+'\n')
    def save(fig,name):
        for ext in ('png','pdf','svg'):fig.savefig(figdir/(name+'.'+ext),dpi=180)
        plt.close(fig)
    fig,axs=plt.subplots(1,2,figsize=(14,4.8),layout='constrained')
    colors={'Q':'C0','K':'C1','V':'C2','QKV':'C3'}
    for key in colors:
        for ax,metric in zip(axs,['effective_fraction','top1']):
            med=[];lo=[];hi=[]
            source= np.concatenate(obs[key+'_source']);vals=np.concatenate(obs[key+'_'+metric])
            for layer in range(23):
                z=stats(vals[source==layer]);lo.append(z['p10']);med.append(z['median']);hi.append(z['p90'])
            ax.plot(range(1,24),med,label=key,color=colors[key]);ax.fill_between(range(1,24),lo,hi,color=colors[key],alpha=.10)
    axs[0].set(title='Continuous effective-head fraction (median, p10–p90)',ylabel='N_eff / number of donor heads')
    axs[1].set(title='Largest donor-head influence share',ylabel='Top-1 normalized ablation effect')
    for ax in axs:
        for j,source in enumerate(full_sources):ax.axvspan(source+.6,source+1.4,color='gray',alpha=.08,label='Full-attention receiver' if j==0 else None)
        ax.set(xlabel='Recipient layer index',ylim=(0,1.02));ax.legend();ax.grid(alpha=.2)
    fig.suptitle('All 23 adjacent transitions; four texts, query positions 7–31; no training')
    save(fig,'continuous_dominance')
    fig,axs=plt.subplots(1,2,figsize=(13,4.8),layout='constrained');curve_rows=[]
    for key in colors:
        a=np.asarray(curves[key]);lo,med,hi=np.quantile(a,[.1,.5,.9],axis=0);x=np.linspace(0,1,101)
        axs[0].plot(x,med,label=key,color=colors[key]);axs[0].fill_between(x,lo,hi,color=colors[key],alpha=.1)
        for j,xx in enumerate(x):curve_rows.append({'metric':key,'fraction_ranked_heads':xx,'p10':lo[j],'median':med[j],'p90':hi[j]})
        v=np.sort(np.asarray(pool[key]['effective_fraction']));axs[1].plot(v,np.arange(1,len(v)+1)/len(v),label=key,color=colors[key])
    axs[0].plot([0,1],[0,1],':',color='gray',label='Uniform');axs[0].set(title='Influence concentration curve',xlabel='Fraction of donor heads, largest first',ylabel='Cumulative normalized single-head effects')
    axs[1].set(title='Continuous effective-head fraction: empirical CDF',xlabel='N_eff / donor heads',ylabel='Fraction of observed cases')
    for ax in axs:ax.set(xlim=(0,1),ylim=(0,1.02));ax.legend();ax.grid(alpha=.2)
    fig.suptitle('Observation-wise ranking; dominant identities can differ; shaded p10–p90')
    csvwrite(root/'concentration-curve.csv',curve_rows);save(fig,'ranked_influence_distribution')
    fig,axs=plt.subplots(1,2,figsize=(14,4.8),layout='constrained')
    for label,color in zip(['weak_half','strong_half','random_half','all_heads'],['C0','C3','C1','gray']):
        for ax,key in zip(axs,['QKV','JS']):
            selected=[r for r in joint_rows if r['metric']==key and r['removed']==label]
            xs=sorted(set(r['recipient'] for r in selected));ys=[np.median(joint_by_layer[(key,label,x-1)]) for x in xs]
            ax.plot(xs,np.maximum(ys,1e-12),'o-',label=label,color=color,markersize=4)
    axs[0].set(title='Joint removal: balanced QKV relative effect',ylabel='Median token/head relative L2 score')
    axs[1].set(title='Full-attention recipients: retrieval change',ylabel='Median token/head JS divergence (nats)')
    for ax in axs:ax.set(xlabel='Recipient layer index',yscale='log');ax.legend();ax.grid(alpha=.2)
    fig.suptitle('Equal-count weak/strong/random removal; token-wise masks; log display floor 1e-12')
    save(fig,'joint_removal_effects')
    fig,axs=plt.subplots(1,2,figsize=(13,4.8),layout='constrained')
    for key,label,color in [('JS','JS divergence (squared scale)','C4'),('sqrtJS','sqrt(JS) distance','C5')]:
        a=np.asarray(curves[key]);lo,med,hi=np.quantile(a,[.1,.5,.9],axis=0);x=np.linspace(0,1,101)
        axs[0].plot(x,med,label=label,color=color);axs[0].fill_between(x,lo,hi,color=color,alpha=.12)
        for j,xx in enumerate(x):curve_rows.append({'metric':key,'fraction_ranked_heads':xx,'p10':lo[j],'median':med[j],'p90':hi[j]})
        v=np.sort(np.asarray(pool[key]['effective_fraction']));axs[1].plot(v,np.arange(1,len(v)+1)/len(v),label=label,color=color)
    axs[0].plot([0,1],[0,1],':',color='gray',label='Uniform');axs[0].set(title='Retrieval influence rank curve',xlabel='Fraction of donor heads, largest first',ylabel='Cumulative normalized intervention effects')
    axs[1].set(title='Sensitivity to divergence versus distance scale',xlabel='Effective donor-head fraction',ylabel='Fraction of observations')
    for ax in axs:ax.set(xlim=(0,1),ylim=(0,1.02));ax.legend();ax.grid(alpha=.2)
    fig.suptitle('Full-attention recipients; sqrt(JS) is post-summary scale sensitivity, not a new model run')
    csvwrite(root/'concentration-curve.csv',curve_rows);save(fig,'retrieval_scale_sensitivity')
    lines=['# 上游注意力头对下一层 Q/K/V 的影响集中度','', '实测全部23个相邻层连接、4条既有中英文文本各32 token。主统计取query7–31，按文本/token/接收query头分别计算，再汇总。不是独立样本的显著性检验。','', '影响是移除上游一个头后，下一层Q/K/V向量的相对L2变化；重新计算中间MLP/残差/归一化。各头影响归一化仅构成消融影响分布，不是天然可相加的贡献概率。','', '| 对象 | 有效头比例中位数 [P10,P90] | 最大头份额中位数 | 前25%头份额中位数 | 最小25%头份额中位数 | 90%影响所需头数中位数 |','|---|---:|---:|---:|---:|---:|']
    for k in ('Q','K','V','QKV','JS','sqrtJS'):
        if k not in pooled:continue
        z=pooled[k];a=z['effective_fraction'];lines.append(f"| {k} | {a['median']:.3f} [{a['p10']:.3f},{a['p90']:.3f}] | {z['top1']['median']:.1%} | {z['top25']['median']:.1%} | {z['bottom25']['median']:.1%} | {z['heads90']['median']:.1f} |")
    lines+=['','sqrtJS为JS距离尺度敏感性分析：在初次完整汇总后增加，原始消融和原JS主结果不变。JS散度近似平方尺度，归一化后会放大强弱差距；JS影响份额不是注意力概率质量份额。','', '按接收层类型分层（避免总体分布混合不同机制）：','', '| 接收层类型 | Q有效头比例中位数 | K | V | QKV |','|---|---:|---:|---:|---:|']
    for label in ('linear_attention','full_attention'):
        lines.append('| '+label+' | '+' | '.join(format(by_type[k][label]['effective_fraction']['median'],'.3f') for k in ('Q','K','V','QKV'))+' |')
    lines+=['','QKV是三个相对变化平方的均值再开根号，避免混加不同Q/K/V尺度；JS仅普通注意力接收层有。普通注意力GQA的K/V复制到query头索引用于配对，不是额外独立KV头。','', '| 联合移除 | QKV相对影响中位数 [P10,P90] | 检索JS中位数 [P10,P90] |','|---|---:|---:|']
    for label in ('weak_half','strong_half','random_half','all_heads'):
        a=joint_stats['QKV'][label];b=joint_stats['JS'][label];lines.append(f"| {label} | {a['median']:.5f} [{a['p10']:.5f},{a['p90']:.5f}] | {b['median']:.6f} [{b['p10']:.6f},{b['p90']:.6f}] |")
    lines+=['', '配对观测：QKV弱半数的实际联合影响小于强半数的比例为'+format(paired['QKV']['weak_less_than_strong_fraction'],'.1%')+'；强/弱影响比中位数为'+format(paired['QKV']['strong_over_weak']['median'],'.3f')+'。这是逐观测比值的中位数，不是两个总体中位数之比。弱效应为零的观测另计，不做除法。', '', '联合实验的弱/强/随机半数头，在各token、各接收头上分别选择，形成完整序列再重算。排名依据单头QKV影响，同样用于JS；它是输入自适应诊断，不是免费可用的剪枝策略。线性注意力的卷积和普通注意力的QK可能传播跨token影响。所有上游头移除仍保留残差背景，不能把它称为删除全部上游信息。', '', '连续指标：N_eff=1/Σp²；有效头比例=N_eff/H；C=(Σp²−1/H)/(1−1/H)。全零总影响时集中度未定义，原始数组保留NaN并在统计时排除、另计数量。单头相对L2影响在做头间占比归一化之前的总和及最大值均保留，避免把微小效应中的集中误认作重要。字段max_absolute指最大无符号相对效应，并非隐藏坐标中的未归一化L2；total是单头效应之和，真正全头联合效果另测。', '', '联合实验中保存的单头影响之和仅为按query掩码加权的参考。单头消融移除整段序列的一个头，联合掩码则可跨位置变化；二者的差不能单独解释成纯非线性交互，因为时间支持也不同。应直接比较实际弱/强/随机联合效果。', '', '## 文件与限制','', '`raw/*.npz`保留每次单头影响、联合影响、相应单头影响之和、删除掩码及空操作对照；`observations.npz`保留逐观测连续指标及索引；CSV和JSON为全部图表源表。原模型权重和大基线缓存继续保留在服务器。','', '只测短文本上的局部下一层作用；不代表全LM质量、可直接删头、所有输入普遍集中，或相似头可互相替代。所有头贡献可能是广泛分布的；报告不以发现稀疏为成功条件。','',f"CPU测量耗时{status['elapsed_seconds']:.1f}秒，峰值RSS{status['peak_rss_kib']/1024:.1f} MiB；无训练、GPU或清理。"]
    (root/'REPORT.md').write_text('\n'.join(lines)+'\n')
    print((root/'REPORT.md').read_text())
if __name__=='__main__':
    import sys
    main(sys.argv[1])