"""
External validation examples for the StatsPAI software-journal materials.

Proves StatsPAI's quality by:
  1. Replicating classic published results (Card 1995, Lee 2008, Prop 99)
  2. Cross-validating against EconML on the same data
  3. Comparing OLS vs IV on real Card (1995) data

These examples are evidence inputs, not blanket validation claims: they
compare selected StatsPAI outputs with established packages or published
benchmarks and leave method-specific validation status to the registry and
JSS audit ledgers.

Usage:
    python run_replication.py
"""

import numpy as np
import pandas as pd
import time
import warnings

warnings.filterwarnings("ignore")

import statspai as sp


# ====================================================================
# EXPERIMENT 1: Card (1995) — IV returns to schooling
# Published result: OLS ~0.073, IV ~0.132 (Angrist & Pischke MHE Table 4.1.1)
# ====================================================================
def replication_card_1995():
    print("\n" + "=" * 75)
    print("REPLICATION 1: Card (1995) — Returns to Schooling (IV)")
    print("Published benchmark: OLS ≈ 0.073, IV ≈ 0.132")
    print("Source: Angrist & Pischke, MHE Table 4.1.1")
    print("=" * 75)

    data, _ = sp.replicate("card_1995")

    # OLS
    ols = sp.regress(
        "lwage ~ educ + exper + expersq + black + south + smsa",
        data=data, robust="hc1",
    )
    ols_educ = ols.params["educ"]
    ols_se = ols.std_errors["educ"]

    # IV: educ instrumented by nearc4
    iv = sp.ivreg(
        "lwage ~ (educ ~ nearc4) + exper + expersq + black + south + smsa",
        data=data,
    )
    iv_educ = iv.params["educ"]
    iv_se = iv.std_errors["educ"]

    print(f"\n  {'':20} {'Coefficient':>12} {'SE':>10} {'Published':>12}")
    print(f"  {'-'*56}")
    print(f"  {'OLS (educ)':<20} {ols_educ:>12.4f} {ols_se:>10.4f} {'~0.073':>12}")
    print(f"  {'IV  (educ)':<20} {iv_educ:>12.4f} {iv_se:>10.4f} {'~0.132':>12}")
    print(f"\n  IV/OLS ratio: {iv_educ/ols_educ:.2f}x (published: ~1.8x)")
    print(f"  N = {len(data)}")

    # First-stage F-statistic
    first_stage = sp.regress(
        "educ ~ nearc4 + exper + expersq + black + south + smsa",
        data=data, robust="hc1",
    )
    fs_coef = first_stage.params["nearc4"]
    fs_se = first_stage.std_errors["nearc4"]
    fs_t = fs_coef / fs_se
    print(f"  First-stage: nearc4 coef = {fs_coef:.4f} (t = {fs_t:.2f})")

    return ols_educ, iv_educ


# ====================================================================
# EXPERIMENT 2: Lee (2008) — RD incumbency advantage
# Published result: RD estimate ≈ 0.08 (8 percentage points)
# ====================================================================
def replication_lee_2008():
    print("\n" + "=" * 75)
    print("REPLICATION 2: Lee (2008) — RD Incumbency Advantage")
    print("Published benchmark: LATE ≈ 0.08 (8 pp)")
    print("Source: Lee (2008), Journal of Econometrics")
    print("=" * 75)

    data, _ = sp.replicate("lee_2008")

    rd = sp.rdrobust(data, y="voteshare_next", x="margin", c=0)

    print(f"\n  RD Estimate:     {rd.estimate:.4f}")
    print(f"  SE:              {rd.se:.4f}")
    ci = rd.ci if hasattr(rd, "ci") and rd.ci else (
        rd.estimate - 1.96 * rd.se, rd.estimate + 1.96 * rd.se
    )
    print(f"  95% CI:          [{ci[0]:.4f}, {ci[1]:.4f}]")
    print(f"  p-value:         {rd.pvalue:.4f}")
    print(f"  N = {len(data)}")
    print(f"  Published LATE ≈ 0.08")
    print(f"  Difference from published: {abs(rd.estimate - 0.08):.4f}")

    # McCrary density test
    try:
        mccrary = sp.mccrary_test(data, x="margin", c=0)
        print(f"  McCrary test p-value: {mccrary.pvalue:.4f} (no manipulation)")
    except Exception as e:
        print(f"  McCrary test: {str(e)[:60]}")

    return rd.estimate


# ====================================================================
# EXPERIMENT 3: California Prop 99 — Synthetic Control
# Published result: ~-26 packs per capita by 2000 (Abadie et al. 2010)
# ====================================================================
def replication_prop99():
    print("\n" + "=" * 75)
    print("REPLICATION 3: Abadie et al. (2010) — Prop 99 Synthetic Control")
    print("Published benchmark: ≈ −26 packs per capita by year 2000")
    print("Source: Abadie, Diamond & Hainmueller (2010), JASA")
    print("=" * 75)

    df = sp.california_prop99()
    print(f"  Data: {df['state'].nunique()} states, {df['year'].nunique()} years")

    # Synthetic control
    try:
        r = sp.synth(
            df,
            outcome="packspercapita",
            unit="state",
            time="year",
            treated_unit="California",
            treatment_time=1989,
        )
        print(f"\n  Estimated gap (year 2000):  {r.estimate:.2f} packs")
        print(f"  Published gap (year 2000): ≈ −26 packs")

        # Print gaps for key years
        if hasattr(r, "gaps") or hasattr(r, "detail"):
            detail = r.detail if hasattr(r, "detail") else {}
            if "gaps" in detail:
                gaps = detail["gaps"]
                for yr in [1989, 1995, 2000]:
                    if yr in gaps.index:
                        print(f"  Gap in {yr}: {gaps.loc[yr]:.2f}")
    except Exception as e:
        print(f"  Synthetic control error: {str(e)[:80]}")
        # Try SDID as fallback
        try:
            r = sp.sdid(
                df, outcome="packspercapita", unit="state",
                time="year", treated_unit="California", treatment_time=1989,
            )
            print(f"  SDID estimate: {r.estimate:.2f}")
        except Exception as e2:
            print(f"  SDID fallback error: {str(e2)[:80]}")


# ====================================================================
# EXPERIMENT 4: LaLonde (1986) — NSW experimental benchmark
# Published: Experimental ITT ≈ $1,794 (Dehejia & Wahba 1999)
# ====================================================================
def replication_lalonde_1986():
    print("\n" + "=" * 75)
    print("REPLICATION 4: LaLonde (1986) / Dehejia & Wahba (1999)")
    print("Published benchmark: Experimental ITT ≈ $1,794")
    print("=" * 75)

    data, _ = sp.replicate("lalonde_1986")
    covs = ["age", "education", "black", "hispanic", "married", "nodegree", "re74", "re75"]

    # Experimental estimate (simple difference in means)
    treated = data[data["treat"] == 1]["re78"]
    control = data[data["treat"] == 0]["re78"]
    raw_diff = treated.mean() - control.mean()
    raw_se = np.sqrt(treated.var() / len(treated) + control.var() / len(control))

    # OLS with controls
    ols = sp.regress(
        f"re78 ~ treat + {' + '.join(covs)}",
        data=data, robust="hc1",
    )
    ols_est = ols.params["treat"]
    ols_se = ols.std_errors["treat"]

    # PSM
    psm = sp.match(data, y="re78", treat="treat", covariates=covs)

    # DML
    dml = sp.dml(data, y="re78", treat="treat", covariates=covs)

    # AIPW
    aipw = sp.aipw(data, y="re78", treat="treat", covariates=covs)

    print(f"\n  {'Estimator':<25} {'Estimate':>10} {'SE':>10}")
    print(f"  {'-'*47}")
    print(f"  {'Raw difference':<25} {raw_diff:>10.1f} {raw_se:>10.1f}")
    print(f"  {'OLS + controls':<25} {ols_est:>10.1f} {ols_se:>10.1f}")
    print(f"  {'PSM':<25} {psm.estimate:>10.1f} {psm.se:>10.1f}")
    print(f"  {'DML':<25} {dml.estimate:>10.1f} {dml.se:>10.1f}")
    print(f"  {'AIPW':<25} {aipw.estimate:>10.1f} {aipw.se:>10.1f}")
    print(f"\n  Published (D&W 1999):    ≈ $1,794")
    print(f"  N = {len(data)} (treated: {len(treated)}, control: {len(control)})")
    print(f"  Note: Using simulated data matching original structure.")
    print(f"  Exact replication requires the original NSW/PSID data.")

    return raw_diff


# ====================================================================
# EXPERIMENT 5: Cross-validation with EconML
# Same data, same task — do both packages agree?
# ====================================================================
def cross_validate_econml():
    print("\n" + "=" * 75)
    print("CROSS-VALIDATION: StatsPAI vs EconML (DML on same data)")
    print("=" * 75)

    try:
        from econml.dml import LinearDML
    except ImportError:
        print("  EconML not installed — skipping cross-validation.")
        return

    from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier

    # Generate data with known effect
    np.random.seed(42)
    df = sp.dgp_observational(n=2000, effect=0.5, confounding=0.3, seed=42)
    Y = df["y"].values
    T = df["treatment"].values
    X = df[["x1", "x2"]].values

    # --- StatsPAI DML ---
    t0 = time.time()
    sp_dml = sp.dml(df, y="y", treat="treatment", covariates=["x1", "x2"])
    sp_time = time.time() - t0

    # --- EconML DML ---
    t0 = time.time()
    econml_dml = LinearDML(
        model_y=RandomForestRegressor(n_estimators=100, random_state=42),
        model_t=RandomForestRegressor(n_estimators=100, random_state=42),
        discrete_treatment=True,
        cv=5,
        random_state=42,
    )
    econml_dml.fit(Y, T, X=X)
    econml_est = econml_dml.ate(X)
    econml_ci = econml_dml.ate_interval(X)
    econml_time = time.time() - t0

    print(f"\n  True ATE = 0.500")
    print(f"\n  {'Package':<15} {'ATE':>10} {'95% CI':>24} {'Time(s)':>10}")
    print(f"  {'-'*62}")

    sp_ci_l = sp_dml.estimate - 1.96 * sp_dml.se
    sp_ci_u = sp_dml.estimate + 1.96 * sp_dml.se
    print(f"  {'StatsPAI':<15} {sp_dml.estimate:>10.4f} {'[' + f'{sp_ci_l:.4f}, {sp_ci_u:.4f}' + ']':>24} {sp_time:>10.3f}")
    print(f"  {'EconML':<15} {econml_est:>10.4f} {'[' + f'{econml_ci[0]:.4f}, {econml_ci[1]:.4f}' + ']':>24} {econml_time:>10.3f}")

    diff = abs(sp_dml.estimate - econml_est)
    print(f"\n  Difference: {diff:.4f}")
    print(f"  Agreement: {'Yes (< 0.05)' if diff < 0.05 else 'Close' if diff < 0.1 else 'Divergent'}")

    # --- Also compare on Card (1995) IV ---
    print(f"\n  --- IV comparison on Card (1995) ---")
    card_data, _ = sp.replicate("card_1995")

    # StatsPAI IV
    sp_iv = sp.ivreg(
        "lwage ~ (educ ~ nearc4) + exper + expersq + black + south + smsa",
        data=card_data,
    )

    # EconML IV (DMLIV)
    try:
        from econml.iv.dml import DMLIV

        Y_card = card_data["lwage"].values
        T_card = card_data["educ"].values
        Z_card = card_data[["nearc4"]].values
        X_card = card_data[["exper", "expersq", "black", "south", "smsa"]].values

        econml_iv = DMLIV(
            model_y_xw=RandomForestRegressor(n_estimators=100, random_state=42),
            model_t_xw=RandomForestRegressor(n_estimators=100, random_state=42),
            model_t_xwz=RandomForestRegressor(n_estimators=100, random_state=42),
            model_final=LinearDML(
                model_y=RandomForestRegressor(n_estimators=50, random_state=42),
                model_t=RandomForestClassifier(n_estimators=50, random_state=42),
            ),
            cv=3,
            random_state=42,
        )
        econml_iv.fit(Y_card, T_card, Z=Z_card, X=X_card)
        econml_iv_est = econml_iv.ate(X_card)

        sp_iv_est = sp_iv.params["educ"]
        print(f"  StatsPAI IV (educ): {sp_iv_est:.4f}")
        print(f"  EconML DMLIV:       {econml_iv_est:.4f}")
        print(f"  Note: EconML uses nonparametric DML-IV, StatsPAI uses classical 2SLS")
        print(f"  Both methods are valid; differences reflect estimator choice, not error")
    except Exception as e:
        sp_iv_est = sp_iv.params["educ"]
        print(f"  StatsPAI IV (educ): {sp_iv_est:.4f}")
        print(f"  EconML DMLIV: skipped ({str(e)[:60]})")


# ====================================================================
# MAIN
# ====================================================================
if __name__ == "__main__":
    np.random.seed(42)

    print("StatsPAI Paper — External Validation Experiments")
    print("=" * 75)
    print(f"StatsPAI version: {sp.__version__}")
    print()

    replication_card_1995()
    replication_lee_2008()
    replication_prop99()
    replication_lalonde_1986()
    cross_validate_econml()

    print("\n" + "=" * 75)
    print("All replication experiments completed.")
    print("=" * 75)
