# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""
08_quantile_models.py
---------------------
Quantile regression analysis for the Airbnb-rent study.

Model 5: Quantile regression across the conditional rent distribution.

  Q_tau(log_rent | X) = alpha_tau + beta_tau * airbnb_count_500m + X'gamma_tau

for tau in {0.10, 0.25, 0.50, 0.75, 0.90}

Outputs:
    results/tables/quantile_regression.tex
    figures/quantile_coefficients.pdf
"""

import sys
from pathlib import Path

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

import matplotlib
matplotlib.use("Agg")
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 statsmodels.regression.quantile_regression import QuantReg  # noqa: E402

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

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

TAUS = [0.10, 0.25, 0.50, 0.75, 0.90]


def prepare_data(rent: pd.DataFrame):
    """Build dummies and return the complete-case design (Y, X, x_cols, bt_cols)."""
    rent_m = rent.copy()

    # Building-type dummies
    if "building_type" in rent_m.columns:
        bt_dummies = pd.get_dummies(
            rent_m["building_type"], prefix="bt", drop_first=True, dtype=float
        )
        rent_m = pd.concat([rent_m, bt_dummies], axis=1)
        bt_cols = list(bt_dummies.columns)
    else:
        bt_cols = []

    # City dummies — use only the top 5 cities to keep manageable
    if "city" in rent_m.columns:
        top5 = rent_m["city"].value_counts().nlargest(5).index.tolist()
        print(f"  Top-5 cities for quantile FE: {top5}")
        rent_m["city_top5"] = rent_m["city"].where(rent_m["city"].isin(top5), other="Other")
        city_dum = pd.get_dummies(
            rent_m["city_top5"], prefix="city", drop_first=True, dtype=float
        )
        rent_m = pd.concat([rent_m, city_dum], axis=1)
        city_fe_cols = list(city_dum.columns)
    else:
        city_fe_cols = []

    controls = ["bedrooms", "bathrooms"] + bt_cols
    x_cols = ["airbnb_count_500m"] + controls + city_fe_cols

    # Drop missing
    use_cols = ["log_rent"] + x_cols
    rent_q = rent_m[use_cols].dropna().reset_index(drop=True)
    print(f"  Complete cases: {len(rent_q):,}")

    Y = rent_q["log_rent"]
    X = sm.add_constant(rent_q[x_cols])
    return Y, X, bt_cols


def build_table(qr_results: dict, res_ols, bt_cols) -> None:
    """Write the quantile-regression LaTeX table."""
    print("\n  Building LaTeX table ...")

    # Variables to display (omit city dummies for readability)
    display_vars = ["const", "airbnb_count_500m", "bedrooms", "bathrooms"] + bt_cols

    n_models = len(TAUS) + 1  # quantiles + OLS
    col_spec = "l" + "c" * n_models
    all_res = [qr_results[t] for t in TAUS] + [res_ols]
    col_headers = [f"$\\tau={t:.2f}$" for t in TAUS] + ["OLS"]

    lines: list[str] = []
    lines.append(r"\begin{tabular}{" + col_spec + "}")
    lines.append(r"\toprule")
    lines.append(
        " & ".join([""] + [f"\\textbf{{{h}}}" for h in col_headers]) + r" \\"
    )
    lines.append(
        " & ".join(["Dep.\\ var:"] + [r"\textit{log\_rent}"] * n_models) + r" \\"
    )
    lines.append(r"\midrule")

    for var in display_vars:
        cells_c = []
        cells_s = []
        for res in all_res:
            if var in res.params.index:
                b = res.params[var]
                se = res.bse[var]
                p = res.pvalues[var]
                cells_c.append(f"{b:.4f}{_star(p)}")
                cells_s.append(f"({se:.4f})")
            else:
                cells_c.append("")
                cells_s.append("")
        vn = var.replace("_", r"\_")
        lines.append(f"{vn} & " + " & ".join(cells_c) + r" \\")
        lines.append(f" & " + " & ".join(cells_s) + r" \\[4pt]")

    lines.append(r"\midrule")

    # City FE indicator
    lines.append(
        "City FE (top 5) & " + " & ".join(["Yes"] * n_models) + r" \\"
    )

    # N
    lines.append(
        "Observations & "
        + " & ".join([f"{int(res.nobs)}" for res in all_res])
        + r" \\"
    )

    # Pseudo-R2 / R2
    r2_cells = []
    for i, res in enumerate(all_res):
        if i < len(TAUS):
            r2_cells.append(f"{res.prsquared:.4f}")
        else:
            r2_cells.append(f"{res.rsquared:.4f}")
    lines.append("(Pseudo-)R$^2$ & " + " & ".join(r2_cells) + r" \\")

    lines.append(r"\bottomrule")
    lines.append(r"\end{tabular}")
    lines.append(
        r"\parbox{\textwidth}{\footnotesize Standard errors in parentheses. "
        r"OLS uses HC1 robust SE. "
        r"$^{***}p<0.01$; $^{**}p<0.05$; $^{*}p<0.10$. "
        r"City FE limited to the 5 largest cities (others grouped).}"
    )

    tex = "\n".join(lines) + "\n"
    out_table = TABLE_DIR / "quantile_regression.tex"
    out_table.write_text(tex, encoding="utf-8")
    print(f"  -> saved {out_table}")


def build_plot(qr_model: QuantReg, res_ols) -> None:
    """Coefficient plot: beta_tau(airbnb_count_500m) across a fine tau grid."""
    print("  Building coefficient plot ...")

    fine_taus = np.arange(0.05, 0.96, 0.05)
    fine_betas = []
    fine_ci_lo = []
    fine_ci_hi = []

    for tau in fine_taus:
        res = qr_model.fit(q=tau)
        beta = res.params["airbnb_count_500m"]
        se = res.bse["airbnb_count_500m"]
        fine_betas.append(beta)
        fine_ci_lo.append(beta - 1.96 * se)
        fine_ci_hi.append(beta + 1.96 * se)

    fine_betas = np.array(fine_betas)
    fine_ci_lo = np.array(fine_ci_lo)
    fine_ci_hi = np.array(fine_ci_hi)

    fig, ax = plt.subplots(figsize=(7, 4.5))

    # 95% CI band
    ax.fill_between(fine_taus, fine_ci_lo, fine_ci_hi, alpha=0.2, color="steelblue",
                    label="95% CI")

    # Quantile regression line
    ax.plot(fine_taus, fine_betas, "o-", color="steelblue", markersize=4,
            label=r"$\beta_\tau$ (Quantile Reg.)")

    # OLS reference
    ols_beta = res_ols.params["airbnb_count_500m"]
    ols_se = res_ols.bse["airbnb_count_500m"]
    ax.axhline(ols_beta, color="firebrick", linestyle="--", linewidth=1.2,
               label=f"OLS estimate ({ols_beta:.5f})")
    ax.axhspan(ols_beta - 1.96 * ols_se, ols_beta + 1.96 * ols_se,
               color="firebrick", alpha=0.08)

    ax.axhline(0, color="grey", linestyle=":", linewidth=0.7)
    ax.set_xlabel(r"Quantile ($\tau$)")
    ax.set_ylabel(r"$\beta_\tau$ (airbnb\_count\_500m)")
    ax.set_title("Effect of Airbnb Count (500m) on Log Rent Across Quantiles")
    ax.legend(fontsize=8, loc="best")
    fig.tight_layout()

    fig_path = FIG_DIR / "quantile_coefficients.pdf"
    fig.savefig(fig_path, dpi=300)
    plt.close(fig)
    print(f"  -> saved {fig_path}")


def main() -> None:
    print("=" * 72)
    print("08  QUANTILE REGRESSION MODELS")
    print("=" * 72)

    rent = pd.read_parquet(require(MERGED_ANALYSIS, PROCESSED_HINT))
    print(f"\nRent data: {rent.shape[0]:,} rows")

    Y, X, bt_cols = prepare_data(rent)

    qr_model = QuantReg(Y, X)
    qr_results: dict[float, object] = {}

    print("\n" + "-" * 72)
    print("MODEL 5: Quantile Regressions")
    print("-" * 72)

    for tau in TAUS:
        res = qr_model.fit(q=tau)
        qr_results[tau] = res
        beta = res.params["airbnb_count_500m"]
        se = res.bse["airbnb_count_500m"]
        p = res.pvalues["airbnb_count_500m"]
        print(
            f"  tau={tau:.2f}  β(airbnb)={beta:.6f}  SE={se:.6f}  "
            f"p={p:.4f}{_star(p)}  pseudo-R2={res.prsquared:.4f}"
        )

    # Also run OLS for comparison
    res_ols = sm.OLS(Y, X).fit(cov_type="HC1")
    print(
        f"  OLS      β(airbnb)={res_ols.params['airbnb_count_500m']:.6f}  "
        f"SE={res_ols.bse['airbnb_count_500m']:.6f}  "
        f"R2={res_ols.rsquared:.4f}"
    )

    build_table(qr_results, res_ols, bt_cols)
    build_plot(qr_model, res_ols)

    print("\n" + "=" * 72)
    print("08  DONE")
    print("=" * 72)


if __name__ == "__main__":
    main()