from __future__ import annotations

import argparse
import datetime as dt
import json
import sys
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

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,
})

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "scripts"))
from audit_core import (COORDS, OVERLAP_ORDER, audit_edges, baseline_charts,
                        circular_block_bootstrap, common_window_score,
                        fit_local, load_panel, mechanical_class, predict,
                        residual_from_difference, run_specification)


PDF_METADATA = {
    "Title": "Tunisia policy-data audit generated figure",
    "Author": "Anonymous",
    "Creator": "Anonymised reproducibility archive",
    "CreationDate": dt.datetime(2026, 7, 22, 0, 0, 0),
    "ModDate": dt.datetime(2026, 7, 22, 0, 0, 0),
}


def save_figure(fig, stem: str):
    out = ROOT / "outputs/figures"
    fig.savefig(out / f"{stem}.pdf", bbox_inches="tight", metadata=PDF_METADATA)
    fig.savefig(out / f"{stem}.png", dpi=200, bbox_inches="tight")
    plt.close(fig)


def main(verify: bool=False):
    panel = load_panel(ROOT)
    values = panel[COORDS].to_numpy(float)
    out_tables = ROOT / "outputs/tables"
    out_logs = ROOT / "outputs/logs"
    out_tables.mkdir(parents=True, exist_ok=True)
    out_logs.mkdir(parents=True, exist_ok=True)

    # Local model coefficients and baseline edge audit.
    charts = baseline_charts()
    fits = {name: fit_local(values, chart, degree=2, ridge=1.0) for name,chart in charts.items()}
    coef_rows = []
    for chart_name, fit in fits.items():
        for coord_idx, coord in enumerate(COORDS):
            row = {"chart": chart_name, "coordinate": coord, "intercept": fit.coefficients[0,coord_idx], "linear": fit.coefficients[1,coord_idx]}
            row["quadratic"] = fit.coefficients[2,coord_idx]
            coef_rows.append(row)
    pd.DataFrame(coef_rows).to_csv(out_tables / "local_model_coefficients.csv", index=False)

    baseline = run_specification(values)
    baseline_rows=[]
    detail_rows=[]
    for overlap,left,right,window in audit_edges():
        d=baseline[overlap]
        baseline_rows.append({
            "overlap":overlap,"window_start":panel.iloc[window[0]].quarter,"window_end":panel.iloc[window[-1]].quarter,"n":len(window),
            "raw_rms":d["raw_rms"],"after_translation":d["after_translation"],"after_gauge":d["after_gauge"],"translation_norm":d["translation_norm"],
            **{f"translation_{c}":d["translation"][i] for i,c in enumerate(COORDS)},
            **{f"gauge_correction_{c}":d["correction"][i] for i,c in enumerate(COORDS)},
            **{f"corrected_component_rms_{c}":d["corrected_component_rms"][i] for i,c in enumerate(COORDS)},
        })
        for row_i,q in enumerate(window):
            for coord_i,c in enumerate(COORDS):
                detail_rows.append({"overlap":overlap,"quarter":panel.iloc[q].quarter,"coordinate":c,
                    "untranslated_difference":d["difference"][row_i,coord_i],
                    "translated_difference":d["translated_difference"][row_i,coord_i],
                    "corrected_difference":d["corrected_difference"][row_i,coord_i]})
    baseline_df=pd.DataFrame(baseline_rows)
    baseline_df.to_csv(out_tables / "baseline_overlap_audit.csv",index=False)
    pd.DataFrame(detail_rows).to_csv(out_tables / "baseline_difference_detail.csv",index=False)

    # Circular moving-block resampling. A single RNG stream is used sequentially across overlaps.
    rng=np.random.default_rng(20260714)
    boot_rows=[]
    band_rows=[]
    for overlap in OVERLAP_ORDER:
        draws=circular_block_bootstrap(baseline[overlap]["difference"],rng,2000,2,0.15)
        for i,v in enumerate(draws): boot_rows.append({"overlap":overlap,"replication":i+1,"corrected_residual":v})
        qlo,qhi=np.quantile(draws,[0.025,0.975])
        band_rows.append({"overlap":overlap,"lower_2_5":qlo,"upper_97_5":qhi})
    pd.DataFrame(boot_rows).to_csv(out_tables / "block_resampling_draws.csv",index=False)
    bands=pd.DataFrame(band_rows)
    bands.to_csv(out_tables / "block_resampling_bands.csv",index=False)

    # Leave-one-quarter-out sensitivity with fixed local chart fits.
    deletion_rows=[]
    deletion_summary=[]
    for overlap in OVERLAP_ORDER:
        r=baseline[overlap]["difference"]
        scores=[]
        _,_,_,window=next(e for e in audit_edges() if e[0]==overlap)
        for i,q in enumerate(window):
            score=residual_from_difference(np.delete(r,i,axis=0),0.15)["after_gauge"]
            scores.append(score)
            deletion_rows.append({"overlap":overlap,"omitted_quarter":panel.iloc[q].quarter,"corrected_residual":score})
        deletion_summary.append({"overlap":overlap,"minimum":min(scores),"maximum":max(scores)})
    pd.DataFrame(deletion_rows).to_csv(out_tables / "leave_one_quarter_out.csv",index=False)
    pd.DataFrame(deletion_summary).to_csv(out_tables / "leave_one_quarter_out_summary.csv",index=False)

    # Gauge sensitivity.
    gauge_rows=[]
    for rho in [0,.05,.10,.15,.20,.25]:
        spec=run_specification(values,gauge_fraction=rho)
        for overlap in OVERLAP_ORDER:
            gauge_rows.append({"gauge_fraction":rho,"gauge_percent":100*rho,"overlap":overlap,"corrected_residual":spec[overlap]["after_gauge"]})
    gauge_df=pd.DataFrame(gauge_rows)
    gauge_df.to_csv(out_tables / "gauge_sensitivity.csv",index=False)

    # Named model/window sensitivity.
    named_specs=[
        ("baseline quadratic",2,1.0,False),
        ("extended shock, linear",1,1.0,True),
        ("extended shock, quadratic",2,1.0,True),
        ("linear local trend",1,1.0,False),
        ("quadratic, strong ridge",2,10.0,False),
        ("quadratic, weak ridge",2,0.1,False),
    ]
    model_rows=[]
    for label,degree,ridge,extended in named_specs:
        spec=run_specification(values,degree,ridge,.15,extended)
        scores=[spec[o]["after_gauge"] for o in OVERLAP_ORDER]
        for o,s in zip(OVERLAP_ORDER,scores): model_rows.append({"specification":label,"degree":degree,"ridge":ridge,"extended_shock":extended,"overlap":o,"corrected_residual":s})
    model_df=pd.DataFrame(model_rows)
    model_df.to_csv(out_tables / "local_model_sensitivity.csv",index=False)

    # Registered 36 specification grid.
    grid_rows=[]
    for degree in [1,2]:
        for ridge in [.1,1.0,10.0]:
            for rho in [0,.05,.10,.15,.20,.25]:
                spec=run_specification(values,degree,ridge,rho,False)
                scores=[spec[o]["after_gauge"] for o in OVERLAP_ORDER]
                largest=OVERLAP_ORDER[int(np.argmax(scores))]
                grid_rows.append({"degree":degree,"ridge":ridge,"gauge_fraction":rho,
                    **{o:s for o,s in zip(OVERLAP_ORDER,scores)},"largest_boundary":largest})
    grid_df=pd.DataFrame(grid_rows)
    grid_df.to_csv(out_tables / "registered_specification_grid_36.csv",index=False)
    counts=grid_df["largest_boundary"].value_counts().reindex(OVERLAP_ORDER,fill_value=0)
    pd.DataFrame({"overlap":counts.index,"count":counts.values,"share_percent":100*counts.values/len(grid_df)}).to_csv(out_tables / "largest_boundary_counts.csv",index=False)

    # Rolling same-design controls and actual boundary scores.
    controls=[]
    control_boundaries=list(range(8,36)) + list(range(68,73))
    for b in control_boundaries:
        controls.append({"boundary_index":b,"boundary_quarter":panel.iloc[b].quarter,"stable_region":"pre" if b<43 else "post","corrected_residual":common_window_score(values,b)})
    controls_df=pd.DataFrame(controls)
    controls_df.to_csv(out_tables / "rolling_within_regime_controls.csv",index=False)
    actual_map={"pre--shock":43,"shock--transition":45,"transition--post":60}
    actual_rows=[]
    ctrl=controls_df.corrected_residual.to_numpy()
    for overlap,b in actual_map.items():
        s=common_window_score(values,b)
        rank=100*np.sum(ctrl<s)/len(ctrl)
        actual_rows.append({"overlap":overlap,"boundary_quarter":panel.iloc[b].quarter,"common_window_score":s,"strict_rank_percent":rank,"n_controls":len(ctrl)})
    actual_df=pd.DataFrame(actual_rows)
    actual_df.to_csv(out_tables / "rolling_actual_boundary_scores.csv",index=False)
    quantiles=pd.DataFrame([{"q25":np.quantile(ctrl,.25),"median":np.quantile(ctrl,.5),"q90":np.quantile(ctrl,.9),"q95":np.quantile(ctrl,.95)}])
    quantiles.to_csv(out_tables / "rolling_control_quantiles.csv",index=False)

    # Threshold sensitivity.
    threshold_rows=[]
    for tau_l in [.08,.10,.12,.14]:
        for tau_h in [.16,.18,.20,.22]:
            if tau_h<=tau_l: continue
            row={"tau_L":tau_l,"tau_H":tau_h}
            for overlap in OVERLAP_ORDER:
                row[overlap]=mechanical_class(baseline[overlap]["after_gauge"],tau_l,tau_h)
            threshold_rows.append(row)
    pd.DataFrame(threshold_rows).to_csv(out_tables / "threshold_sensitivity.csv",index=False)

    # Evidence-integrated action rule under the illustrative registry.
    decisions=[]
    actual_lookup=actual_df.set_index("overlap")
    for overlap in OVERLAP_ORDER:
        d=baseline[overlap]["after_gauge"]
        p=float(actual_lookup.loc[overlap,"strict_rank_percent"])
        mech=mechanical_class(d,.12,.18)
        if d>.18 or p>97.5:
            action="Keep separate"
        elif d>=.12 or p>=90:
            action="Review before pooling"
        else:
            action="Conditional pooling"
        decisions.append({"overlap":overlap,"baseline_residual":d,"control_percentile":p,"translation_norm":baseline[overlap]["translation_norm"],"mechanical_class":mech,"evidence_integrated_action":action,"metadata_flag":0,"translation_flag":"not classified"})
    decisions_df=pd.DataFrame(decisions)
    decisions_df.to_csv(out_tables / "evidence_integrated_decisions.csv",index=False)

    # Figures 2-7.
    fig,ax=plt.subplots(figsize=(8.2,4.8))
    ax.plot(panel.year,panel.trust_obs,label="Trust/governance")
    ax.plot(panel.year,panel.shock_obs,label="Crisis stress")
    ax.plot(panel.year,panel.policy_proxy,label="Visible policy response")
    for x in [2010.75,2011.25,2015.0]: ax.axvline(x,linestyle=":",linewidth=.8)
    ax.set_xlabel("Year"); ax.set_ylabel("Aligned planning coordinate [0,1]"); ax.set_title("Source-aligned Tunisia planning panel"); ax.legend(); ax.set_ylim(-.03,1.05)
    save_figure(fig,"Figure_2_Tunisia_aligned_panel")

    merged=baseline_df.merge(bands,on="overlap").set_index("overlap").loc[OVERLAP_ORDER].reset_index()
    fig,ax=plt.subplots(figsize=(7.8,4.8))
    y=merged.after_gauge.to_numpy(); low=y-merged.lower_2_5.to_numpy(); high=merged.upper_97_5.to_numpy()-y
    ax.bar(np.arange(3),y,yerr=np.vstack([low,high]),capsize=3)
    ax.set_xticks(np.arange(3),OVERLAP_ORDER,rotation=15); ax.set_ylabel("Corrected RMS residual"); ax.set_title("Baseline residuals and resampling sensitivity bands")
    save_figure(fig,"Figure_3_baseline_residuals")

    fig,ax=plt.subplots(figsize=(7.8,4.8))
    ax.hist(ctrl,bins=12)
    for _,row in actual_df.iterrows(): ax.axvline(row.common_window_score,linestyle="--",label=f"{row.overlap} ({row.strict_rank_percent:.1f}th pct.)")
    ax.set_xlabel("Common-window pseudo-boundary residual"); ax.set_ylabel("Count"); ax.set_title("Rolling within-regime controls (n=33)"); ax.legend(fontsize=8)
    save_figure(fig,"Figure_4_rolling_controls")

    fig,ax=plt.subplots(figsize=(7.8,4.8))
    for o in OVERLAP_ORDER:
        sub=gauge_df[gauge_df.overlap==o]
        ax.plot(sub.gauge_percent,sub.corrected_residual,marker="o",label=o)
    ax.set_xlabel("Gauge bound (percent of translated component RMS)"); ax.set_ylabel("Corrected RMS residual"); ax.set_title("Gauge-bound sensitivity"); ax.legend()
    save_figure(fig,"Figure_5_gauge_sensitivity")

    pivot=model_df.pivot(index="specification",columns="overlap",values="corrected_residual").loc[[s[0] for s in named_specs],OVERLAP_ORDER]
    fig,ax=plt.subplots(figsize=(9.2,5.2))
    x=np.arange(len(pivot)); width=.25
    for i,o in enumerate(OVERLAP_ORDER): ax.bar(x+(i-1)*width,pivot[o],width,label=o)
    short_labels=["baseline quadratic","extended shock, linear","extended shock, quadratic","linear local trend","quadratic, strong ridge","quadratic, weak ridge"]
    ax.set_xticks(x,short_labels,rotation=24,ha="right"); ax.set_ylabel("Corrected RMS residual"); ax.set_title("Local-model and shock-window sensitivity"); ax.legend()
    save_figure(fig,"Figure_6_model_sensitivity")

    fig,ax=plt.subplots(figsize=(7.8,4.8))
    ax.axhspan(0,.12,alpha=.15,label="conditional pooling band"); ax.axhspan(.12,.18,alpha=.15,label="review band"); ax.axhspan(.18,.25,alpha=.15,label="separate band")
    ax.scatter(np.arange(3),[baseline[o]["after_gauge"] for o in OVERLAP_ORDER])
    ax.set_xticks(np.arange(3),OVERLAP_ORDER,rotation=15); ax.set_ylim(0,.25); ax.set_ylabel("Baseline corrected residual"); ax.set_title("Illustrative agency tolerance registry"); ax.legend(fontsize=8)
    save_figure(fig,"Figure_7_tolerance_registry")

    summary={
        "baseline_after_gauge":{o:baseline[o]["after_gauge"] for o in OVERLAP_ORDER},
        "bootstrap_bands":bands.set_index("overlap").to_dict(orient="index"),
        "leave_one_out":pd.DataFrame(deletion_summary).set_index("overlap").to_dict(orient="index"),
        "rolling_actual":actual_df.set_index("overlap").to_dict(orient="index"),
        "control_quantiles":quantiles.iloc[0].to_dict(),
        "largest_boundary_counts":counts.to_dict(),
        "decisions":decisions_df.set_index("overlap")["evidence_integrated_action"].to_dict(),
    }
    (out_logs / "reproduction_summary.json").write_text(json.dumps(summary,indent=2),encoding="utf-8")

    if verify:
        expected={
            "baseline_round3":{"pre--shock":.102,"shock--transition":.209,"transition--post":.147},
            "raw_round3":{"pre--shock":.316,"shock--transition":.353,"transition--post":.230},
            "translation_round3":{"pre--shock":.102,"shock--transition":.210,"transition--post":.149},
            "translation_norm_round3":{"pre--shock":.518,"shock--transition":.497,"transition--post":.332},
            "bands_round3":{"pre--shock":[.030,.146],"shock--transition":[.076,.294],"transition--post":[.049,.198]},
            "loqo_round3":{"pre--shock":[.084,.111],"shock--transition":[.170,.225],"transition--post":[.117,.160]},
            "common_round3":{"pre--shock":.251,"shock--transition":.438,"transition--post":.372},
            "rank_round1":{"pre--shock":93.9,"shock--transition":100.0,"transition--post":97.0},
        }
        failures=[]
        bdf=baseline_df.set_index("overlap")
        for o,v in expected["baseline_round3"].items():
            if round(float(bdf.loc[o,"after_gauge"]),3)!=v: failures.append(f"baseline {o}")
        for o,v in expected["raw_round3"].items():
            if round(float(bdf.loc[o,"raw_rms"]),3)!=v: failures.append(f"raw {o}")
        for o,v in expected["translation_round3"].items():
            if round(float(bdf.loc[o,"after_translation"]),3)!=v: failures.append(f"translation {o}")
        for o,v in expected["translation_norm_round3"].items():
            if round(float(bdf.loc[o,"translation_norm"]),3)!=v: failures.append(f"translation norm {o}")
        bnd=bands.set_index("overlap")
        for o,(lo,hi) in expected["bands_round3"].items():
            if [round(float(bnd.loc[o,"lower_2_5"]),3),round(float(bnd.loc[o,"upper_97_5"]),3)] != [lo,hi]: failures.append(f"bootstrap {o}")
        lodf=pd.DataFrame(deletion_summary).set_index("overlap")
        for o,(lo,hi) in expected["loqo_round3"].items():
            if [round(float(lodf.loc[o,"minimum"]),3),round(float(lodf.loc[o,"maximum"]),3)] != [lo,hi]: failures.append(f"LOQO {o}")
        act=actual_df.set_index("overlap")
        for o,v in expected["common_round3"].items():
            if round(float(act.loc[o,"common_window_score"]),3)!=v: failures.append(f"common {o}")
        for o,v in expected["rank_round1"].items():
            if round(float(act.loc[o,"strict_rank_percent"]),1)!=v: failures.append(f"rank {o}")
        if counts.get("shock--transition",0)!=30 or counts.get("transition--post",0)!=6 or counts.get("pre--shock",0)!=0:
            failures.append("36-specification ranking counts")
        report={"status":"PASS" if not failures else "FAIL","failures":failures,"checked_against":"main manuscript Tables 4-10 and Figures 2-7"}
        (out_logs / "verification_report.json").write_text(json.dumps(report,indent=2),encoding="utf-8")
        if failures:
            raise SystemExit("Verification failed: "+", ".join(failures))
        print("Verification PASS: all reported main-text values reproduce to displayed precision.")
    else:
        print("Reproduction complete. Run with --verify to check reported values.")

if __name__=="__main__":
    parser=argparse.ArgumentParser()
    parser.add_argument("--verify",action="store_true")
    args=parser.parse_args()
    main(args.verify)
