from __future__ import annotations
from pathlib import Path
import sys, json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch

ROOT=Path(__file__).resolve().parents[1]
sys.path.insert(0,str(ROOT/'scripts'))
from audit_core import load_panel, fit_local, residual_from_difference, COORDS

PDF_METADATA={'Title':'Quality and Quantity supplementary analysis','Author':'Anonymous','Creator':'Anonymised reproducibility archive'}

plt.rcParams.update({
    'font.size': 10,
    'axes.titlesize': 11,
    'axes.labelsize': 10,
    'legend.fontsize': 8.5,
    'xtick.labelsize': 9,
    'ytick.labelsize': 9,
    'lines.linewidth': 1.5,
})

OUT=ROOT/'outputs/qaq_extension'
(OUT/'tables').mkdir(parents=True,exist_ok=True)
(OUT/'figures').mkdir(parents=True,exist_ok=True)

panel=load_panel(ROOT)
X=panel[COORDS].to_numpy(float)


def _chart(indices):
    a=np.asarray(indices,dtype=float)
    return {'indices':list(map(int,indices)),'centre':float(a.mean()),'scale':float(a.std())}

def common_window_decomp(boundary:int):
    left=np.arange(boundary-8,boundary)
    right=np.arange(boundary,boundary+8)
    audit=np.arange(boundary-3,boundary+3)
    lf=fit_local(X,_chart(left),degree=2,ridge=1.0)
    rf=fit_local(X,_chart(right),degree=2,ridge=1.0)
    diff=(np.column_stack([((audit-rf.centre)/rf.scale)**k for k in range(rf.degree+1)])@rf.coefficients
          -np.column_stack([((audit-lf.centre)/lf.scale)**k for k in range(lf.degree+1)])@lf.coefficients)
    res=residual_from_difference(diff,0.15)
    L=X[left]; R=X[right]
    return {
        'boundary_index':boundary,
        'boundary_quarter':panel.iloc[boundary].quarter,
        'translation_norm':float(res['translation_norm']),
        'shape_residual':float(res['after_gauge']),
        'raw_rms':float(res['raw_rms']),
        'mean_shift_norm':float(np.linalg.norm(R.mean(axis=0)-L.mean(axis=0))),
        'median_shift_norm':float(np.linalg.norm(np.median(R,axis=0)-np.median(L,axis=0))),
        'covariance_shift_frobenius':float(np.linalg.norm(np.cov(R,rowvar=False)-np.cov(L,rowvar=False),'fro')),
    }

controls=list(range(8,36))+list(range(68,73))
actual={'pre--shock':43,'shock--transition':45,'transition--post':60}
rows=[]
for b in controls:
    rows.append({'kind':'within-regime control','label':panel.iloc[b].quarter,**common_window_decomp(b)})
for label,b in actual.items():
    rows.append({'kind':'actual boundary','label':label,**common_window_decomp(b)})
all_df=pd.DataFrame(rows)
all_df.to_csv(OUT/'tables'/'level_shape_boundary_diagnostics.csv',index=False)
ctrl=all_df[all_df.kind=='within-regime control']
rank_rows=[]
for label,b in actual.items():
    r=all_df[(all_df.kind=='actual boundary')&(all_df.label==label)].iloc[0]
    d={'boundary':label,'quarter':r.boundary_quarter}
    for col in ['translation_norm','shape_residual','raw_rms','mean_shift_norm','median_shift_norm','covariance_shift_frobenius']:
        d[col]=float(r[col])
        d[col+'_strict_rank_percent']=100*float((ctrl[col]<r[col]).sum())/len(ctrl)
    rank_rows.append(d)
pd.DataFrame(rank_rows).to_csv(OUT/'tables'/'boundary_benchmark_ranks.csv',index=False)

# Level-shape map.
fig,ax=plt.subplots(figsize=(7.8,5.7))
ax.scatter(ctrl['translation_norm'],ctrl['shape_residual'],s=34,alpha=.65,label='within-regime controls')
marks=['o','s','^']
for (label,_),mk in zip(actual.items(),marks):
    r=all_df[(all_df.kind=='actual boundary')&(all_df.label==label)].iloc[0]
    ax.scatter([r.translation_norm],[r.shape_residual],s=105,marker=mk,label=label)
    ax.annotate(label,(r.translation_norm,r.shape_residual),xytext=(6,6),textcoords='offset points',fontsize=9)
ax.set_xlabel('Registered level-translation norm')
ax.set_ylabel('Post-registration shape residual')
ax.set_title('Level conversion and residual shape disagreement')
ax.legend(loc='best',frameon=True)
ax.margins(x=.08,y=.10)
fig.tight_layout()
fig.savefig(OUT/'figures'/'Fig3_level_shape_map.pdf',bbox_inches='tight',metadata=PDF_METADATA)
fig.savefig(OUT/'figures'/'Fig3_level_shape_map.png',dpi=260,bbox_inches='tight')
plt.close(fig)

# Monte Carlo discrimination experiment.
rng=np.random.default_rng(20260722)
t=np.linspace(-1,1,8)
p=3
base_intercept=np.array([.35,.25,.20])
base_slope=np.array([.08,-.04,.05])
level=np.array([.30,-.20,.25])
shape=np.array([.18,-.12,.16])
scenarios={
    'no break':(np.zeros(p),np.zeros(p)),
    'level only':(level,np.zeros(p)),
    'shape only':(np.zeros(p),shape),
    'level and shape':(level,shape),
}
mc=[]
for sigma in [.02,.05,.10]:
    for scenario,(delta,kappa) in scenarios.items():
        for rep in range(1000):
            left=base_intercept+np.outer(t,base_slope)+rng.normal(0,sigma,size=(8,p))
            right=base_intercept+delta+np.outer(t,base_slope+kappa)+rng.normal(0,sigma,size=(8,p))
            diff=right-left
            audit=residual_from_difference(diff,.15)
            mc.append({'sigma':sigma,'scenario':scenario,'replication':rep+1,
                       'raw_rms':audit['raw_rms'],'translation_norm':audit['translation_norm'],
                       'shape_residual':audit['after_gauge'],
                       'mean_shift_norm':float(np.linalg.norm(right.mean(0)-left.mean(0)))})
mcdf=pd.DataFrame(mc)
mcdf.to_csv(OUT/'tables'/'simulation_draws.csv',index=False)
summary=[]
for (sigma,scenario),g in mcdf.groupby(['sigma','scenario'],sort=False):
    row={'sigma':sigma,'scenario':scenario}
    for col in ['raw_rms','translation_norm','shape_residual','mean_shift_norm']:
        row[col+'_median']=float(g[col].median())
        row[col+'_q90']=float(g[col].quantile(.90))
    row['shape_residual_gt_0_12_share']=float((g.shape_residual>.12).mean())
    summary.append(row)
pd.DataFrame(summary).to_csv(OUT/'tables'/'simulation_summary.csv',index=False)

# Boxplot at sigma .05.
sub=mcdf[mcdf.sigma==.05]
order=list(scenarios.keys())
fig,ax=plt.subplots(figsize=(8.0,5.0))
data=[sub[sub.scenario==s].shape_residual.to_numpy() for s in order]
ax.boxplot(data,tick_labels=order,showfliers=False)
ax.axhline(.12,linestyle='--',linewidth=1.0,label='illustrative lower tolerance')
ax.set_ylabel('Post-registration shape residual')
ax.set_title(r'Discrimination between level and shape breaks ($\sigma=0.05$)')
ax.legend(loc='upper left',frameon=True)
fig.tight_layout()
fig.savefig(OUT/'figures'/'Fig4_simulation_discrimination.pdf',bbox_inches='tight',metadata=PDF_METADATA)
fig.savefig(OUT/'figures'/'Fig4_simulation_discrimination.png',dpi=260,bbox_inches='tight')
plt.close(fig)

# Two-row workflow figure with publication-size text.
fig,ax=plt.subplots(figsize=(9.6,5.0))
ax.set_xlim(0,1); ax.set_ylim(0,1); ax.axis('off')
steps=[
    (0.03,0.59,0.27,0.25,'1  Register records','Regimes, variables, source versions,\nwindows and authorised corrections'),
    (0.365,0.59,0.27,0.25,'2  Register levels','Estimate the declared level crosswalk\nand report its magnitude'),
    (0.70,0.59,0.27,0.25,'3  Measure shape','Compute the disagreement remaining\nafter registered operations'),
    (0.20,0.14,0.27,0.25,'4  Test coherence','Check graph-wide crosswalk consistency\nand full-record feasibility'),
    (0.535,0.14,0.27,0.25,'5  Document action','Report evidence, margins, provenance\nand the accountable disposition'),
]
for x,y,w,h,title,body in steps:
    patch=FancyBboxPatch((x,y),w,h,boxstyle='round,pad=0.012,rounding_size=0.018',
                         linewidth=1.2,facecolor='white',edgecolor='black')
    ax.add_patch(patch)
    ax.text(x+w/2,y+h*.68,title,ha='center',va='center',fontsize=10.5,fontweight='bold')
    ax.text(x+w/2,y+h*.33,body,ha='center',va='center',fontsize=9.2,linespacing=1.25)

def arrow(x1,y1,x2,y2):
    ax.annotate('',xy=(x2,y2),xytext=(x1,y1),arrowprops=dict(arrowstyle='->',lw=1.25))
arrow(.30,.715,.365,.715)
arrow(.635,.715,.70,.715)
arrow(.835,.59,.47,.39)
arrow(.47,.265,.535,.265)
ax.set_title('Registered compatibility audit: from regime records to an accountable disposition',pad=12,fontweight='bold')
fig.tight_layout()
fig.savefig(OUT/'figures'/'Fig1_audit_architecture.pdf',bbox_inches='tight',metadata=PDF_METADATA)
fig.savefig(OUT/'figures'/'Fig1_audit_architecture.png',dpi=260,bbox_inches='tight')
plt.close(fig)

# Residual-only threshold sensitivity across the registered tolerance grid.
thresholds=pd.read_csv(ROOT/'outputs/tables'/'threshold_sensitivity.csv')
disposition_order=['Conditional pooling','Review before pooling','Keep separate']
boundary_order=['pre--shock','shock--transition','transition--post']
fig,ax=plt.subplots(figsize=(7.8,4.8))
x=np.arange(len(boundary_order)); width=.23
for j,disposition in enumerate(disposition_order):
    vals=[int((thresholds[b]==disposition).sum()) for b in boundary_order]
    bars=ax.bar(x+(j-1)*width,vals,width,label=disposition)
    for bar,val in zip(bars,vals):
        if val:
            ax.text(bar.get_x()+bar.get_width()/2,val+.25,str(val),ha='center',va='bottom',fontsize=8.5)
ax.set_xticks(x,boundary_order)
ax.set_ylabel('Registered threshold pairs (out of 16)')
ax.set_ylim(0,17.2)
ax.set_title('Residual-only threshold sensitivity')
ax.legend(frameon=True)
fig.tight_layout()
fig.savefig(OUT/'figures'/'Fig9_threshold_sensitivity.pdf',bbox_inches='tight',metadata=PDF_METADATA)
fig.savefig(OUT/'figures'/'Fig9_threshold_sensitivity.png',dpi=260,bbox_inches='tight')
plt.close(fig)

summary_obj={'actual_benchmarks':rank_rows,'simulation_rows':len(mcdf),'controls':len(controls),'seed':20260722}
(OUT/'extension_summary.json').write_text(json.dumps(summary_obj,indent=2),encoding='utf-8')
print(json.dumps(summary_obj,indent=2))
