# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""
10_robustness_tables_figures.py
-------------------------------
Generate additional publication-quality robustness tables and figures.

Robustness tables
-----------------
1. Winsorization / outlier robustness  -> results/tables/robustness_outliers.tex
2. Subsample analysis                  -> results/tables/robustness_subsamples.tex

Additional figures
------------------
1. figures/rent_airbnb_heatmap.pdf      — hexbin of Montreal listings
2. figures/coefficient_robustness.pdf   — coefficient comparison plot
3. figures/rent_by_airbnb_bins.pdf      — mean rent by airbnb quintiles
"""

import sys
import warnings
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from src.plotting import use_publication_style

use_publication_style()

import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402
import statsmodels.api as sm  # noqa: E402

from src.config import MERGED_ANALYSIS, TABLE_DIR, FIG_DIR, require  # noqa: E402
from src.latex_tables import significance_star as _stars_raw  # noqa: E402

warnings.filterwarnings("ignore")

PROCESSED_HINT = "Run scripts/04_merge_data.py first."


def _stars(pval: float) -> str:
    """Return significance stars (empty for NaN p-values)."""
    if pd.isna(pval):
        return ""
    return _stars_raw(pval)


# ─────────────────────────────────────────────────────────────────────────────
# Helpers: OLS with city fixed effects
# ─────────────────────────────────────────────────────────────────────────────

def run_ols_city_fe(data: pd.DataFrame, label: str,
                    airbnb_var: str = "airbnb_count_500m",
                    with_fit_stats: bool = True) -> dict:
    """
    OLS: log_rent ~ airbnb_var + bedrooms + bathrooms + city_FE.
    Returns dict with coefficient, std error, p-value (and N, R² when
    with_fit_stats is True).
    """
    needed = ["log_rent", airbnb_var, "bedrooms", "bathrooms", "city"]
    sub = data[needed].dropna()
    if len(sub) < 30:
        out = {"label": label, "coef": np.nan, "se": np.nan, "pval": np.nan}
        if with_fit_stats:
            out.update({"N": len(sub), "R2": np.nan})
        return out
    city_dummies = pd.get_dummies(sub["city"], prefix="city", drop_first=True, dtype=float)
    X = pd.concat([
        sub[[airbnb_var, "bedrooms", "bathrooms"]].astype(float).reset_index(drop=True),
        city_dummies.reset_index(drop=True),
    ], axis=1)
    X = sm.add_constant(X.astype(float))
    y = sub["log_rent"].astype(float).reset_index(drop=True)
    model = sm.OLS(y, X).fit(cov_type="HC1")
    out = {
        "label": label,
        "coef": model.params[airbnb_var],
        "se": model.bse[airbnb_var],
        "pval": model.pvalues[airbnb_var],
    }
    if with_fit_stats:
        out.update({"N": int(model.nobs), "R2": model.rsquared})
    return out


# ═════════════════════════════════════════════════════════════════════════════
# TABLE 1: Winsorization / outlier robustness
# ═════════════════════════════════════════════════════════════════════════════

def table_outlier_robustness(df: pd.DataFrame) -> None:
    print("\n--- Winsorization robustness ---")

    # (a) Full sample
    res_full = run_ols_city_fe(df, "Full sample")

    # (b) Winsorized at 5th / 95th percentile on monthly_rent
    p5 = df["monthly_rent"].quantile(0.05)
    p95 = df["monthly_rent"].quantile(0.95)
    df_win = df[(df["monthly_rent"] >= p5) & (df["monthly_rent"] <= p95)].copy()
    res_win = run_ols_city_fe(df_win, "Rent winsorized (5/95)")

    # (c) Excluding top / bottom 1% of airbnb_count_500m
    p1 = df["airbnb_count_500m"].quantile(0.01)
    p99 = df["airbnb_count_500m"].quantile(0.99)
    df_trim = df[(df["airbnb_count_500m"] >= p1) & (df["airbnb_count_500m"] <= p99)].copy()
    res_trim = run_ols_city_fe(df_trim, "Airbnb trimmed (1/99)")

    outlier_results = [res_full, res_win, res_trim]
    for r in outlier_results:
        print(f"  {r['label']:30s}  coef={r['coef']:.6f}  N={r['N']}")

    # Build LaTeX table
    lines = []
    lines.append(r"\begin{tabular}{lccc}")
    lines.append(r"\toprule")
    lines.append(r" & (1) Full Sample & (2) Rent 5/95 & (3) Airbnb 1/99 \\")
    lines.append(r"\midrule")
    # Coefficient row
    coef_cells = " & ".join(
        f"{r['coef']:.6f}{_stars(r['pval'])}" for r in outlier_results
    )
    lines.append(f"Airbnb count (500m) & {coef_cells} \\\\")
    # SE row
    se_cells = " & ".join(f"({r['se']:.6f})" for r in outlier_results)
    lines.append(f" & {se_cells} \\\\")
    lines.append(r"\midrule")
    # N
    n_cells = " & ".join(f"{r['N']:,}" for r in outlier_results)
    lines.append(f"N & {n_cells} \\\\")
    # R²
    r2_cells = " & ".join(f"{r['R2']:.4f}" for r in outlier_results)
    lines.append(f"$R^2$ & {r2_cells} \\\\")
    lines.append(r"City FE & Yes & Yes & Yes \\")
    lines.append(r"\bottomrule")
    lines.append(r"\end{tabular}")
    lines.append(r"\begin{tablenotes}\small")
    lines.append(r"\item \textit{Notes:} Robust standard errors (HC1) in parentheses. "
                 r"* $p<0.10$, ** $p<0.05$, *** $p<0.01$. "
                 r"Column (2) winsorizes monthly rent at the 5th and 95th percentiles. "
                 r"Column (3) trims the top and bottom 1\% of Airbnb listing counts.")
    lines.append(r"\end{tablenotes}")

    (TABLE_DIR / "robustness_outliers.tex").write_text("\n".join(lines) + "\n",
                                                       encoding="utf-8")
    print(f"Saved {TABLE_DIR / 'robustness_outliers.tex'}")


# ═════════════════════════════════════════════════════════════════════════════
# TABLE 2: Subsample analysis
# ═════════════════════════════════════════════════════════════════════════════

def table_subsamples(df: pd.DataFrame) -> list[dict]:
    print("\n--- Subsample analysis ---")

    # Identify Montreal
    montreal_mask = df["city"].str.contains("Montr", case=False, na=False)

    sub_results = []

    # (a) Montreal only
    res_mtl = run_ols_city_fe(df[montreal_mask], "Montreal only")
    sub_results.append(res_mtl)

    # (b) Outside Montreal
    res_non = run_ols_city_fe(df[~montreal_mask], "Outside Montreal")
    sub_results.append(res_non)

    # (c) Apartments only
    apt_mask = df["building_type"].str.contains("Apartment", case=False, na=False)
    res_apt = run_ols_city_fe(df[apt_mask], "Apartments only")
    sub_results.append(res_apt)

    # (d) Houses only
    house_mask = df["building_type"].str.contains("House", case=False, na=False)
    res_house = run_ols_city_fe(df[house_mask], "Houses only")
    sub_results.append(res_house)

    for r in sub_results:
        print(f"  {r['label']:30s}  coef={r['coef']:.6f}  N={r['N']}")

    # Build LaTeX table
    lines2 = []
    lines2.append(r"\begin{tabular}{lcccc}")
    lines2.append(r"\toprule")
    lines2.append(r" & (1) Montreal & (2) Outside Mtl & (3) Apartments & (4) Houses \\")
    lines2.append(r"\midrule")
    coef_cells2 = " & ".join(
        f"{r['coef']:.6f}{_stars(r['pval'])}" if not pd.isna(r['coef']) else "---"
        for r in sub_results
    )
    lines2.append(f"Airbnb count (500m) & {coef_cells2} \\\\")
    se_cells2 = " & ".join(
        f"({r['se']:.6f})" if not pd.isna(r['se']) else ""
        for r in sub_results
    )
    lines2.append(f" & {se_cells2} \\\\")
    lines2.append(r"\midrule")
    n_cells2 = " & ".join(
        f"{r['N']:,}" for r in sub_results
    )
    lines2.append(f"N & {n_cells2} \\\\")
    r2_cells2 = " & ".join(
        f"{r['R2']:.4f}" if not pd.isna(r['R2']) else "---"
        for r in sub_results
    )
    lines2.append(f"$R^2$ & {r2_cells2} \\\\")
    lines2.append(r"City FE & Yes & Yes & Yes & Yes \\")
    lines2.append(r"\bottomrule")
    lines2.append(r"\end{tabular}")
    lines2.append(r"\begin{tablenotes}\small")
    lines2.append(r"\item \textit{Notes:} Robust standard errors (HC1) in parentheses. "
                  r"* $p<0.10$, ** $p<0.05$, *** $p<0.01$. "
                  r"All specifications include city fixed effects, bedrooms, and bathrooms "
                  r"as controls.")
    lines2.append(r"\end{tablenotes}")

    (TABLE_DIR / "robustness_subsamples.tex").write_text("\n".join(lines2) + "\n",
                                                         encoding="utf-8")
    print(f"Saved {TABLE_DIR / 'robustness_subsamples.tex'}")

    return sub_results


# ═════════════════════════════════════════════════════════════════════════════
# FIGURE 1: Hexbin heatmap — Montreal area
# ═════════════════════════════════════════════════════════════════════════════

def figure_heatmap(df: pd.DataFrame) -> None:
    print("\n--- Hexbin heatmap (Montreal) ---")

    mtl_area = df[
        (df["lat"] >= 45.4) & (df["lat"] <= 45.6)
        & (df["lon"] >= -73.8) & (df["lon"] <= -73.5)
    ].dropna(subset=["lat", "lon", "airbnb_count_500m"])

    fig, ax = plt.subplots(figsize=(8, 6))
    hb = ax.hexbin(
        mtl_area["lon"], mtl_area["lat"],
        C=mtl_area["airbnb_count_500m"],
        reduce_C_function=np.mean,
        gridsize=40, cmap="YlOrRd", mincnt=1,
    )
    fig.colorbar(hb, ax=ax, label="Mean Airbnb count (500m)")
    ax.set_xlabel("Longitude")
    ax.set_ylabel("Latitude")
    ax.set_title("Airbnb Density — Montreal Area")
    fig.savefig(FIG_DIR / "rent_airbnb_heatmap.pdf")
    plt.close(fig)
    print(f"Saved {FIG_DIR / 'rent_airbnb_heatmap.pdf'}")


# ═════════════════════════════════════════════════════════════════════════════
# FIGURE 2: Coefficient robustness plot
# ═════════════════════════════════════════════════════════════════════════════

def figure_coefficient_robustness(df: pd.DataFrame, sub_results: list[dict]) -> None:
    print("\n--- Coefficient robustness plot ---")

    res_mtl, res_non, res_apt = sub_results[0], sub_results[1], sub_results[2]

    specs = []
    # Baseline (500m)
    specs.append(run_ols_city_fe(df, "Baseline (500m)"))

    # Buffer variants
    for buf, label in [("250m", "250m buffer"), ("1km", "1km buffer"), ("2km", "2km buffer")]:
        var = f"airbnb_count_{buf}"
        if var in df.columns:
            specs.append(run_ols_city_fe(df, label, airbnb_var=var,
                                         with_fit_stats=False))

    # Subsamples
    specs.append(res_mtl)
    specs.append(res_non)
    specs.append(res_apt)

    # Filter out NaN results
    specs = [s for s in specs if not pd.isna(s.get("coef", np.nan))]

    labels = [s["label"] for s in specs]
    coefs = np.array([s["coef"] for s in specs])
    ses = np.array([s["se"] for s in specs])

    fig, ax = plt.subplots(figsize=(7, 0.5 * len(specs) + 2))
    y_pos = np.arange(len(specs))
    ax.errorbar(
        coefs, y_pos, xerr=1.96 * ses,
        fmt="o", color="#2166ac", ecolor="#92c5de", capsize=4,
        markersize=6, elinewidth=1.5,
    )
    ax.axvline(0, color="grey", linestyle="--", linewidth=0.8)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(labels)
    ax.set_xlabel(r"Coefficient on Airbnb count ($\beta$)")
    ax.set_title("Coefficient Robustness Across Specifications")
    ax.invert_yaxis()
    fig.savefig(FIG_DIR / "coefficient_robustness.pdf")
    plt.close(fig)
    print(f"Saved {FIG_DIR / 'coefficient_robustness.pdf'}")


# ═════════════════════════════════════════════════════════════════════════════
# FIGURE 3: Mean rent by Airbnb count quintiles
# ═════════════════════════════════════════════════════════════════════════════

def figure_rent_by_bins(df: pd.DataFrame) -> None:
    print("\n--- Rent by Airbnb quintiles ---")

    sub_q = df[["monthly_rent", "airbnb_count_500m"]].dropna()
    sub_q["quintile"] = pd.qcut(
        sub_q["airbnb_count_500m"], q=5, labels=False, duplicates="drop"
    )
    quintile_means = sub_q.groupby("quintile")["monthly_rent"].mean()
    quintile_labels = []
    for q in sorted(sub_q["quintile"].unique()):
        lo = sub_q.loc[sub_q["quintile"] == q, "airbnb_count_500m"].min()
        hi = sub_q.loc[sub_q["quintile"] == q, "airbnb_count_500m"].max()
        quintile_labels.append(f"Q{int(q)+1}\n[{lo:.0f}-{hi:.0f}]")

    fig, ax = plt.subplots(figsize=(7, 5))
    bars = ax.bar(
        range(len(quintile_means)),
        quintile_means.values,
        color="#2166ac",
        edgecolor="white",
        width=0.65,
    )
    # Add value labels on bars
    for bar, val in zip(bars, quintile_means.values):
        ax.text(
            bar.get_x() + bar.get_width() / 2, bar.get_height() + 10,
            f"${val:,.0f}", ha="center", va="bottom", fontsize=9,
        )
    ax.set_xticks(range(len(quintile_means)))
    ax.set_xticklabels(quintile_labels)
    ax.set_xlabel("Airbnb Count (500m) Quintile")
    ax.set_ylabel("Mean Monthly Rent ($)")
    ax.set_title("Mean Rent by Airbnb Exposure Quintile")
    fig.savefig(FIG_DIR / "rent_by_airbnb_bins.pdf")
    plt.close(fig)
    print(f"Saved {FIG_DIR / 'rent_by_airbnb_bins.pdf'}")


def main() -> None:
    print("Loading data ...")
    df = pd.read_parquet(require(MERGED_ANALYSIS, PROCESSED_HINT))
    print(f"Full sample: {len(df):,} observations")

    table_outlier_robustness(df)
    sub_results = table_subsamples(df)
    figure_heatmap(df)
    figure_coefficient_robustness(df, sub_results)
    figure_rent_by_bins(df)

    print("\n=== 10_robustness_tables_figures.py complete ===")


if __name__ == "__main__":
    main()