"""Offline rerun: python analysis.py --input . --output rerun
Requires existing pandas, numpy, matplotlib. No installation or network calls.
Source: CDC/NCHS xbxb-epbu, national all-sex/all-race rows, 1999–2016.
"""
from pathlib import Path
import argparse, json, hashlib, sys
import pandas as pd
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

def analyze(input_dir=".", output_dir="rerun"):
    src, out = Path(input_dir), Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)
    raw = (src/"source_national.json").read_bytes()
    manifest = json.loads((src/"source_manifest.json").read_text(encoding="utf-8"))
    assert hashlib.sha256(raw).hexdigest() == manifest["sha256"]["source_national.json"]
    df = pd.DataFrame(json.loads(raw))
    numeric = ["year","deaths","population","crude_death_rate","age_adjusted_rate",
        "lower_confidence_limit_for_crude_rate","upper_confidence_limit_for_crude_rate",
        "lower_confidence_limit_for_age_adjusted_rate","upper_confidence_limit_for_age_adjusted_rate"]
    for c in numeric:
        df[c] = pd.to_numeric(df[c], errors="coerce")
    assert df.state.eq("United States").all()
    assert df.sex.eq("Both Sexes").all()
    assert df.race_and_hispanic_origin.eq("All Races-All Origins").all()
    assert len(df)==162 and not df.duplicated(["year","age_group"]).any()
    assert set(df.year)==set(range(1999,2017))
    assert not df[["deaths","population","crude_death_rate"]].isna().any().any()
    assert (df.population>0).all()
    assert abs(df.crude_death_rate-df.deaths/df.population*100000).max()<0.000051
    a = df[df.age_group.eq("All Ages")].set_index("year").sort_index()
    g = df[~df.age_group.eq("All Ages")]
    bands = ["0–14","15–24","25–34","35–44","45–54","55–64","65–74","75+"]
    pop = g.pivot(index="year",columns="age_group",values="population")[bands]
    deaths = g.pivot(index="year",columns="age_group",values="deaths")[bands]
    assert not pop.isna().any().any() and not deaths.isna().any().any()
    assert (pop.sum(axis=1)==a.population).all()
    residual_deaths = a.deaths-deaths.sum(axis=1)
    assert residual_deaths.between(0,16).all()
    p = pop.div(pop.sum(axis=1),axis=0)
    r = deaths/pop*100000
    residual_rate = residual_deaths/a.population*100000
    exact = a.deaths/a.population*100000
    composition = (p.loc[2016]-p.loc[1999])*(r.loc[2016]+r.loc[1999])/2
    within = (r.loc[2016]-r.loc[1999])*(p.loc[2016]+p.loc[1999])/2
    residual_change = residual_rate.loc[2016]-residual_rate.loc[1999]
    change = exact.loc[2016]-exact.loc[1999]
    assert abs(composition.sum()+within.sum()+residual_change-change)<1e-10
    annual = a[[c for c in numeric if c!="year"]].copy()
    assert not annual.isna().any().any()
    annual["exact_crude_rate"] = exact
    annual["adjusted_minus_crude"] = a.age_adjusted_rate-a.crude_death_rate
    annual["adjusted_over_crude"] = a.age_adjusted_rate/a.crude_death_rate
    annual["crude_change_pct"] = a.crude_death_rate.pct_change()*100
    annual["adjusted_change_pct"] = a.age_adjusted_rate.pct_change()*100
    annual["unallocated_age_deaths"] = residual_deaths
    annual["unallocated_age_rate"] = residual_rate
    annual["fixed_1999_composition_known_age_rate"] = r.mul(p.loc[1999],axis=1).sum(axis=1)
    annual["fixed_2016_composition_known_age_rate"] = r.mul(p.loc[2016],axis=1).sum(axis=1)
    decomp = pd.DataFrame({"population_share_1999":p.loc[1999],
        "population_share_2016":p.loc[2016],"rate_1999":r.loc[1999],"rate_2016":r.loc[2016],
        "composition_contribution":composition,"within_age_rate_contribution":within})
    df.to_csv(out/"national_extract.csv", index=False)
    annual.to_csv(out/"annual_comparison.csv")
    decomp.to_csv(out/"age_decomposition.csv")
    summary = {
        "crude_increase_pct":float((a.crude_death_rate.iloc[-1]/a.crude_death_rate.iloc[0]-1)*100),
        "adjusted_increase_pct":float((a.age_adjusted_rate.iloc[-1]/a.age_adjusted_rate.iloc[0]-1)*100),
        "maximum_absolute_relative_rate_gap_pct":float(abs(annual.adjusted_over_crude-1).max()*100),
        "crude_change_per_100000":float(change),
        "composition_per_100000":float(composition.sum()),
        "within_age_per_100000":float(within.sum()),
        "unallocated_age_residual_change_per_100000":float(residual_change),
        "composition_share_of_net_change_pct":float(composition.sum()/change*100),
        "fixed_1999_rate_change":float(((r.loc[2016]-r.loc[1999])*p.loc[1999]).sum()),
        "fixed_2016_rate_change":float(((r.loc[2016]-r.loc[1999])*p.loc[2016]).sum()),
        "all_broad_age_rates_increased":bool((r.loc[2016]>r.loc[1999]).all()),
        "annual_direction_agrees":bool((np.sign(a.crude_death_rate.diff().dropna())==np.sign(a.age_adjusted_rate.diff().dropna())).all()),
        "validation":"All assertions passed; unallocated-age residual retained separately."
    }
    (out/"results.json").write_text(json.dumps(summary,indent=2),encoding="utf-8")
    with plt.rc_context({"font.size":10,"pdf.fonttype":42,"axes.spines.top":False,"axes.spines.right":False}):
        fig, axes = plt.subplots(2,2,figsize=(12,8),layout="constrained")
        ax=axes[0,0]
        ax.plot(a.index,a.crude_death_rate,color="#0072B2",marker="o",markersize=3,label="Crude")
        ax.plot(a.index,a.age_adjusted_rate,color="#A34A00",linestyle="--",label="Age-adjusted (2000 US)")
        ax.set(title="A  Published national rates",ylabel="Deaths per 100,000",ylim=(0,22))
        ax.legend(frameon=False)
        ax=axes[0,1]
        ax.plot(a.index,annual.adjusted_minus_crude,color="#333333",marker="s",markersize=3)
        ax.axhline(0,color="gray",linewidth=0.8)
        ax.set(title="B  Adjusted minus crude",ylabel="Rate difference per 100,000")
        ax=axes[1,0]
        ax.plot(a.index,exact,color="#0072B2",label="Observed crude")
        ax.plot(a.index,annual.fixed_1999_composition_known_age_rate,color="#A34A00",linestyle="--",label="Fixed 1999 age shares")
        ax.plot(a.index,annual.fixed_2016_composition_known_age_rate,color="#555555",linestyle=":",label="Fixed 2016 age shares")
        ax.set(title="C  Broad-group composition sensitivity",ylabel="Deaths per 100,000",ylim=(0,22))
        ax.legend(frameon=False,fontsize=9)
        ax=axes[1,1]
        y=np.arange(len(bands))
        ax.barh(y-0.19,within,height=0.36,color="#0072B2",label="Within-age rate component")
        ax.barh(y+0.19,composition,height=0.36,color="#A34A00",hatch="//",label="Age-share component")
        ax.set_yticks(y,bands)
        ax.invert_yaxis()
        ax.axvline(0,color="gray",linewidth=0.8)
        ax.set(title="D  Contributions to 1999–2016 change",xlabel="Crude-rate change per 100,000")
        ax.legend(frameon=False,fontsize=8,loc="lower right")
        for ax in axes.flat: ax.grid(axis="y",alpha=0.15)
        for ax in [axes[0,0],axes[0,1],axes[1,0]]:
            ax.set_xlabel("Year")
            ax.set_xticks([1999,2004,2008,2012,2016])
        fig.suptitle("US drug-overdose mortality, 1999–2016",fontsize=16)
        fig.get_layout_engine().set(rect=(0,0.065,1,0.95))
        fig.text(0.02,0.015,"Source: CDC/NCHS xbxb-epbu; both sexes, all races/origins. No smoothing.\nC excludes unallocated-age deaths in fixed-share curves; D residual = −0.00264 per 100,000. Descriptive estimates; CIs in CSV.",fontsize=9)
        fig.savefig(out/"national_trends.png",dpi=200,facecolor="white")
        fig.savefig(out/"national_trends.pdf",facecolor="white")
        plt.close(fig)
    return summary

if __name__=="__main__":
    parser=argparse.ArgumentParser()
    parser.add_argument("--input",default=".")
    parser.add_argument("--output",default="rerun")
    args=parser.parse_args()
    print(json.dumps(analyze(args.input,args.output),indent=2))
