# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""
06_hedonic_models.py
--------------------
Hedonic regression models for the Airbnb-rent analysis.

Model 1 (1a-1e): Baseline hedonic rent model (OLS, HC1)
Model 2 (2a-2c): Hedonic Airbnb pricing model
Model 3:         City-level interaction (forward + reverse)

Outputs:
    results/tables/hedonic_rent_baseline.tex
    results/tables/hedonic_airbnb_pricing.tex
    results/tables/city_level_interaction.tex
"""

import sys
from pathlib import Path

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

import numpy as np
import pandas as pd
import statsmodels.api as sm

from src.config import AIRBNB_CLEAN, MERGED_ANALYSIS, TABLE_DIR, require
from src.latex_tables import significance_star, results_to_latex

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


def run_ols(df: pd.DataFrame, y_col: str, x_cols: list[str], label: str):
    """Run OLS with HC1 robust SE, dropping NaN rows for relevant cols."""
    cols = [y_col] + x_cols
    sub = df[cols].dropna()
    Y = sub[y_col]
    X = sm.add_constant(sub[x_cols])
    model = sm.OLS(Y, X).fit(cov_type="HC1")
    print(f"\n  [{label}]  N={int(model.nobs):,}  R2={model.rsquared:.4f}  "
          f"Adj-R2={model.rsquared_adj:.4f}")
    return model


# ═════════════════════════════════════════════════════════════════════════════
# MODEL 1: Baseline Hedonic Rent Model (1a – 1e)
# ═════════════════════════════════════════════════════════════════════════════

def model1_hedonic_rent(rent: pd.DataFrame) -> None:
    print("\n" + "-" * 72)
    print("MODEL 1: Baseline Hedonic Rent Model")
    print("-" * 72)

    # --- prepare variables -------------------------------------------------
    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 (fixed effects)
    if "city" in rent_m.columns:
        city_dummies = pd.get_dummies(rent_m["city"], prefix="city",
                                      drop_first=True, dtype=float)
        rent_m = pd.concat([rent_m, city_dummies], axis=1)
        city_fe_cols = list(city_dummies.columns)
    else:
        city_fe_cols = []

    controls = ["bedrooms", "bathrooms"] + bt_cols

    # 1a: No controls, no FE
    res_1a = run_ols(rent_m, "log_rent", ["airbnb_count_500m"], "1a")

    # 1b: With controls, no FE
    res_1b = run_ols(rent_m, "log_rent", ["airbnb_count_500m"] + controls, "1b")

    # 1c: With controls + city FE
    res_1c = run_ols(
        rent_m, "log_rent", ["airbnb_count_500m"] + controls + city_fe_cols, "1c"
    )

    # 1d: airbnb_density_500m instead
    res_1d = run_ols(
        rent_m, "log_rent", ["airbnb_density_500m"] + controls + city_fe_cols, "1d"
    )

    # 1e: share_entire_home_500m instead
    res_1e = run_ols(
        rent_m, "log_rent", ["share_entire_home_500m"] + controls + city_fe_cols, "1e"
    )

    # -- Key display variables (not all city/building dummies) --------------
    display_rent = (
        ["const", "airbnb_count_500m", "airbnb_density_500m",
         "share_entire_home_500m", "bedrooms", "bathrooms"]
        + bt_cols
    )
    all_res = [res_1a, res_1b, res_1c, res_1d, res_1e]

    results_to_latex(
        all_res,
        ["(1a)", "(1b)", "(1c)", "(1d)", "(1e)"],
        dep_var="log\\_rent",
        display_vars=[v for v in display_rent
                      if any(v in r.params.index for r in all_res)],
        out_path=TABLE_DIR / "hedonic_rent_baseline.tex",
        note="Models (1c)-(1e) include city fixed effects (not shown).",
    )


# ═════════════════════════════════════════════════════════════════════════════
# MODEL 2: Hedonic Airbnb Pricing Model (2a – 2c)
# ═════════════════════════════════════════════════════════════════════════════

def model2_airbnb_pricing(rent: pd.DataFrame, airbnb: pd.DataFrame) -> None:
    print("\n" + "-" * 72)
    print("MODEL 2: Hedonic Airbnb Pricing Model")
    print("-" * 72)

    # Compute mean rent per city from rent data
    city_col_rent = "city"
    # Determine city column name in airbnb data
    city_col_ab = "city_clean" if "city_clean" in airbnb.columns else "city"

    mean_rent_city = (
        rent.groupby(city_col_rent)["log_rent"]
        .mean()
        .rename("mean_rent_city")
        .reset_index()
        .rename(columns={city_col_rent: city_col_ab})
    )
    print(f"\n  Mean rent computed for {len(mean_rent_city)} cities")

    ab = airbnb.merge(mean_rent_city, on=city_col_ab, how="inner")
    print(f"  Airbnb rows after merge: {len(ab):,}")

    # Ensure log_price exists
    if "log_price" not in ab.columns and "price_numeric" in ab.columns:
        ab["log_price"] = np.log(ab["price_numeric"].clip(lower=1))

    # City FE for Airbnb
    ab_city_dum = pd.get_dummies(ab[city_col_ab], prefix="acity",
                                 drop_first=True, dtype=float)
    ab = pd.concat([ab, ab_city_dum], axis=1)
    ab_city_fe = list(ab_city_dum.columns)

    ab_controls = [
        c for c in ["bedrooms", "bathrooms", "guests_count",
                    "amenities_count", "is_superhost", "is_entire_home"]
        if c in ab.columns
    ]

    # Convert boolean controls to float
    for c in ab_controls:
        if ab[c].dtype == bool:
            ab[c] = ab[c].astype(float)

    # 2a: No controls
    res_2a = run_ols(ab, "log_price", ["mean_rent_city"], "2a")

    # 2b: With controls
    res_2b = run_ols(ab, "log_price", ["mean_rent_city"] + ab_controls, "2b")

    # 2c: With controls + city FE
    res_2c = run_ols(ab, "log_price",
                     ["mean_rent_city"] + ab_controls + ab_city_fe, "2c")

    display_ab = ["const", "mean_rent_city"] + ab_controls
    all_res = [res_2a, res_2b, res_2c]

    results_to_latex(
        all_res,
        ["(2a)", "(2b)", "(2c)"],
        dep_var="log\\_price",
        display_vars=[v for v in display_ab
                      if any(v in r.params.index for r in all_res)],
        out_path=TABLE_DIR / "hedonic_airbnb_pricing.tex",
        note="Model (2c) includes city fixed effects (not shown).",
    )


# ═════════════════════════════════════════════════════════════════════════════
# MODEL 3: City-level Interaction Model
# ═════════════════════════════════════════════════════════════════════════════

def model3_city_interaction(rent: pd.DataFrame) -> None:
    print("\n" + "-" * 72)
    print("MODEL 3: City-level Interaction")
    print("-" * 72)

    # Aggregate rent data to city level
    city_agg = (
        rent.groupby("city")
        .agg(
            mean_log_rent=("log_rent", "mean"),
            mean_bedrooms=("bedrooms", "mean"),
            mean_bathrooms=("bathrooms", "mean"),
            n_rent_listings=("log_rent", "count"),
        )
        .reset_index()
    )

    # Get airbnb_count_city from rent data (should be constant within a city)
    airbnb_city_vars = [c for c in rent.columns
                        if c.startswith("airbnb_") and c.endswith("_city")]
    if airbnb_city_vars:
        city_airbnb = rent.groupby("city")[airbnb_city_vars].first().reset_index()
        city_agg = city_agg.merge(city_airbnb, on="city", how="left")

    print(f"  City-level rows: {len(city_agg)}")
    print(f"  Columns: {list(city_agg.columns)}")

    # Determine which airbnb count variable is available
    ab_count_col = "airbnb_count_city" if "airbnb_count_city" in city_agg.columns else None
    city_controls = ["mean_bedrooms", "mean_bathrooms"]

    if ab_count_col is None:
        print("  WARNING: airbnb_count_city not found — skipping Model 3.")
        return

    # 3-forward: mean_log_rent ~ airbnb_count_city + controls
    res_3fwd = run_ols(
        city_agg, "mean_log_rent", [ab_count_col] + city_controls, "3-fwd"
    )
    # 3-reverse: airbnb_count_city ~ mean_log_rent + controls
    res_3rev = run_ols(
        city_agg, ab_count_col, ["mean_log_rent"] + city_controls, "3-rev"
    )

    # LaTeX table (fragment: the paper supplies the table environment)
    lines: list[str] = []
    lines.append(r"\begin{tabular}{lcc}")
    lines.append(r"\toprule")
    lines.append(r" & \textbf{(3-fwd)} & \textbf{(3-rev)} \\")
    lines.append(
        r"Dep.\ var: & \textit{mean\_log\_rent} & \textit{airbnb\_count\_city} \\"
    )
    lines.append(r"\midrule")

    # Show all variables for both models
    all_vars_3 = list(dict.fromkeys(
        list(res_3fwd.params.index) + list(res_3rev.params.index)
    ))
    for var in all_vars_3:
        cells_coef = []
        cells_se = []
        for res in [res_3fwd, res_3rev]:
            if var in res.params.index:
                b = res.params[var]
                se = res.bse[var]
                p = res.pvalues[var]
                cells_coef.append(f"{b:.4f}{significance_star(p)}")
                cells_se.append(f"({se:.4f})")
            else:
                cells_coef.append("")
                cells_se.append("")
        vn = var.replace("_", r"\_")
        lines.append(f"{vn} & {cells_coef[0]} & {cells_coef[1]} " + r"\\")
        lines.append(f" & {cells_se[0]} & {cells_se[1]} " + r"\\[4pt]")

    lines.append(r"\midrule")
    lines.append(
        f"Observations & {int(res_3fwd.nobs)} & {int(res_3rev.nobs)} " + r"\\"
    )
    lines.append(
        f"R$^2$ & {res_3fwd.rsquared:.4f} & {res_3rev.rsquared:.4f} " + r"\\"
    )
    lines.append(r"\bottomrule")
    lines.append(r"\end{tabular}")
    lines.append(
        r"\parbox{\textwidth}{\footnotesize Robust (HC1) standard errors in "
        r"parentheses. $^{***}p<0.01$; $^{**}p<0.05$; $^{*}p<0.10$.}"
    )

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


def main() -> None:
    print("=" * 72)
    print("06  HEDONIC REGRESSION MODELS")
    print("=" * 72)

    rent = pd.read_parquet(require(MERGED_ANALYSIS, PROCESSED_HINT))
    airbnb = pd.read_parquet(require(AIRBNB_CLEAN, PROCESSED_HINT))

    print(f"\nRent data: {rent.shape[0]:,} rows, {rent.shape[1]} cols")
    print(f"Airbnb data: {airbnb.shape[0]:,} rows, {airbnb.shape[1]} cols")

    model1_hedonic_rent(rent)
    model2_airbnb_pricing(rent, airbnb)
    model3_city_interaction(rent)

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


if __name__ == "__main__":
    main()