#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""Step 03 — Vertical-inequity regressions, heterogeneity, horizontal
inequity and the implied tax shift.

Writes to ``results/reproduced/``:
  vertical.csv        Cheng OLS / Cheng FE / Clapp IV / Paglin–Fogarty
  quantile.csv        quantile-regression β(τ) on within-cell data
  heterogeneity.csv   Cheng-FE γ by property class, age, land share, year,
                      roll lag and municipality size
  binscatter.csv      within-cell mean ln ratio by price vigintile
  horizontal.csv      |deviation| regression (who gets noisy assessments)
  taxshift.csv        over/under-taxation by within-cell price decile
  robustness.csv      γ across sample and measurement variants

Usage:  python scripts/03_estimate_regressions.py
"""
import sys
from pathlib import Path

import numpy as np
import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from wp10 import config, models, sample  # noqa: E402


def main() -> None:
    config.ensure_dirs()
    df = sample.load()
    out = config.REPRODUCED

    # ------------------------------------------------------------ main table
    print("Vertical-inequity estimators")
    res = [models.cheng_pooled(df), models.cheng_fe(df), models.clapp_iv(df)]
    for r in res:
        print(f"  {r['estimator']:<28} beta={r['beta']:.4f} ({r['se']:.4f})  "
              f"gamma={r['gamma']:+.4f}  n={r['n']:,}")
    pf = models.paglin_fogarty(df)
    print(f"  {pf['estimator']:<28} a={pf['intercept']:,.0f} "
          f"({pf['intercept_se']:,.0f})  b={pf['slope']:.4f}")
    pd.DataFrame(res).to_csv(out / "vertical.csv", index=False)
    pd.Series(pf).to_csv(out / "paglin_fogarty.csv")

    # ------------------------------------------------------------ quantiles
    qt = models.quantile_betas(df)
    qt.to_csv(out / "quantile.csv", index=False)
    print("Quantile betas:", {f"{r.tau:.2f}": round(r.beta, 3)
                              for r in qt.itertuples()})

    # ------------------------------------------------------------ binscatter
    d = df.copy()
    d["lnr_w"] = d["ln_ratio"] - d.groupby("cell")["ln_ratio"].transform("mean")
    d["lnp_w"] = d["ln_price"] - d.groupby("cell")["ln_price"].transform("mean")
    d["bin"] = pd.qcut(d["lnp_w"], 20, labels=False)
    (d.groupby("bin")
       .agg(x=("lnp_w", "mean"), y=("lnr_w", "mean"),
            se=("lnr_w", lambda s: s.std() / np.sqrt(len(s))), n=("lnr_w", "size"))
       .reset_index()
       .to_csv(out / "binscatter.csv", index=False))

    # ------------------------------------------------------------ heterogeneity
    muni_sales = df.groupby("muni")["muni"].transform("size")
    groups = {
        "Single-family": df["prop_class"] == "single_family",
        "Condominium": df["prop_class"] == "condo",
        "Plex (2–5 units)": df["prop_class"] == "plex",
        "Cottage": df["prop_class"] == "cottage",
        "Age < 20 y": df["age"] < 20,
        "Age 20–60 y": df["age"].between(20, 60),
        "Age > 60 y": df["age"] > 60,
        "Land share < 0.2": df["land_share"] < 0.2,
        "Land share 0.2–0.4": df["land_share"].between(0.2, 0.4),
        "Land share > 0.4": df["land_share"] > 0.4,
        "Roll lag < 24 m": df["lag_months"] < 24,
        "Roll lag 24–48 m": df["lag_months"].between(24, 48),
        "Roll lag > 48 m": df["lag_months"] > 48,
        "Muni < 1k sales": muni_sales < 1_000,
        "Muni 1k–10k sales": muni_sales.between(1_000, 10_000),
        "Muni > 10k sales": muni_sales > 10_000,
    }
    groups.update({f"Sales {y}": df["sale_year"] == y
                   for y in sorted(df["sale_year"].unique())})
    het = models.gamma_by_group(df, groups)
    het.to_csv(out / "heterogeneity.csv", index=False)
    print(f"Heterogeneity: {len(het)} subgroups estimated")

    # ------------------------------------------------------------ horizontal
    tab, meta = models.horizontal_dispersion(df)
    tab.to_csv(out / "horizontal.csv")
    pd.Series(meta).to_csv(out / "horizontal_meta.csv")
    print("Horizontal-dispersion regression:", meta)

    # ------------------------------------------------------------ tax shift
    ts = models.tax_shift(df)
    ts.to_csv(out / "taxshift.csv", index=False)
    print("Tax shift by decile (mean %):",
          {int(r.decile): f"{r.mean_rel:+.1%}" for r in ts.itertuples()})

    # ------------------------------------------------------------ robustness
    variants = {
        "Baseline": df,
        "Condominiums only": df[df["prop_class"] == "condo"],
        "Excluding sales < $100k": df[df["amount"] >= 100_000],
        "Match score = 220 (max)": df[df["match_score"] >= 219.9],
        "Match distance <= 10 m": df[df["match_dist_m"] <= 10],
        "Single-family only": df[df["prop_class"] == "single_family"],
        "Ratio trim 5/95": df[df["ratio"].between(
            df["ratio"].quantile(.05), df["ratio"].quantile(.95))],
        "Sales 2021-2023": df[df["sale_year"] <= 2023],
        "Sales 2024-2026": df[df["sale_year"] >= 2024],
        "Munis >= 300 sales": df[df.groupby("muni")["muni"]
                                 .transform("size") >= 300],
        "Cells >= 50 sales": df[df.groupby("cell")["cell"]
                                .transform("size") >= 50],
    }
    rows = []
    for label, sub in variants.items():
        counts = sub.groupby("cell")["cell"].transform("size")
        sub = sub[counts >= config.CELL_MIN_OBS]
        est = models.cheng_fe(sub)
        iv = models.clapp_iv(sub)
        rows.append({"variant": label, "gamma_fe": est["gamma"],
                     "se_fe": est["se"], "gamma_iv": iv["gamma"],
                     "se_iv": iv["se"], "n": est["n"]})
        print(f"  {label:<26} gamma_FE={est['gamma']:+.4f} "
              f"gamma_IV={iv['gamma']:+.4f}  n={est['n']:,}")
    pd.DataFrame(rows).to_csv(out / "robustness.csv", index=False)


if __name__ == "__main__":
    main()