#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""Step 02 — Estimate the M1–M5 hedonic specification ladder.

Writes to ``results/reproduced/``:
  fit.json          R^2 and N for each specification
  coef_M{1,2,3,5}.csv  coefficient tables (coef, se, p)
  grand_model.parquet  per-listing grand-model (M5) fitted values/residuals
  fsa_premia.csv    FSA location premia (grand-model fixed effects, >=50 listings)

Usage:  python scripts/02_estimate_core.py
"""
import json
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 wp9 import models, sample  # noqa: E402
from wp9.config import REPRODUCED, ensure_dirs  # noqa: E402


def main() -> None:
    ensure_dirs()
    s = sample.load_sample()
    ladder = models.specification_ladder(s)

    fit = {}
    for name, (res, cols, data) in ladder.items():
        r2 = float(res.rsquared)
        fit[name] = {"r2": r2, "n": int(res.nobs)}
        print(f"{name}: N={int(res.nobs):,}  R2={r2:.4f}")
        if name != "M4":
            models.coef_table(res, cols).to_csv(REPRODUCED / f"coef_{name}.csv")
    json.dump(fit, open(REPRODUCED / "fit.json", "w"), indent=2)

    # Grand model (M5): per-listing predictions, residuals and FSA premia
    res5, cols5, _ = ladder["M5"]
    fe = models.fsa_fixed_effects(s, res5, cols5)
    pred = models.predict_with_fe(s, res5, cols5, fe)
    frame = s[["fsa_c", "prov", "lat", "lon", "ln_price"]].copy()
    frame["pred_grand"] = pred
    frame["resid_grand"] = frame["ln_price"] - frame["pred_grand"]
    frame.to_parquet(REPRODUCED / "grand_model.parquet", index=False)

    premia = (frame.assign(fe=s["fsa_c"].map(fe))
              .groupby("fsa_c")
              .agg(fe=("fe", "first"), n=("fe", "size"), prov=("prov", "first"),
                   lat=("lat", "mean"), lon=("lon", "mean")))
    premia = premia[premia["n"] >= 50].copy()
    reference_level = premia["fe"].median()
    premia["premium_pct"] = (np.exp(premia["fe"] - reference_level) - 1) * 100
    premia.index.name = "fsa_c"
    premia.sort_values("premium_pct", ascending=False).to_csv(REPRODUCED / "fsa_premia.csv")
    print(f"FSA premia written for {len(premia):,} neighbourhoods (>=50 listings)")


if __name__ == "__main__":
    main()