# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""The WP11 horse-race harness.

Twenty hedonic models organised along four design axes:

  A. functional form   — linear, semi-log, log-log, Box–Cox, quadratic,
                         splines (linear estimator, muni + quarter FE)
  B. time effects      — none / year / quarter / month FE (quadratic backbone)
  C. spatial controls  — none / municipality / ~1.1 km grid / neighbourhood
                         unit FE (quadratic backbone, quarter FE)
  D. estimation method — OLS, ridge, random forest, gradient boosting
                         (with and without coordinates), spatial k-NN
                         comparables

Every linear model is estimated as a sparse ridge with a vanishing penalty
(numerically OLS, but it (i) survives collinear dummies and (ii) predicts
unseen categories at the reference level via ``handle_unknown='ignore'``).
Every model is evaluated on identical train/test splits, and all metrics
are computed on *price levels* — log models are retransformed with Duan's
smearing factor, Box–Cox by analytic inversion — so that functional forms
compete on the same scoreboard.
"""
from __future__ import annotations

import time

import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import (HistGradientBoostingRegressor,
                              RandomForestRegressor)
from sklearn.linear_model import Ridge, RidgeCV
from sklearn.neighbors import KNeighborsRegressor
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, SplineTransformer

from . import config

NUM_LEVELS = ["area", "lot", "has_lot", "age", "floors", "units"]
NUM_LOGS = ["ln_area", "ln_lot", "has_lot", "age", "floors", "units"]
NUM_QUAD = NUM_LEVELS + ["area2", "age2"]
CATS = ["prop_class", "link"]
ML_NUM = ["area", "lot", "has_lot", "age", "floors", "units",
          "lat", "lng", "t"]


def add_derived(df: pd.DataFrame) -> pd.DataFrame:
    """Quadratic terms, scaled to keep the design well conditioned."""
    out = df.copy()
    out["area2"] = (out["area"] / 100.0) ** 2
    out["age2"] = (out["age"] / 10.0) ** 2
    return out


# ---------------------------------------------------------------- splits
def split(df: pd.DataFrame, scheme: str):
    """Return (train, test) under the random or forward-in-time scheme."""
    if scheme == "random":
        rng = np.random.default_rng(config.SEED)
        mask = rng.random(len(df)) < config.TEST_SHARE
        return df[~mask], df[mask]
    if scheme == "temporal":
        cut = pd.Timestamp(config.TEMPORAL_CUTOFF)
        return df[df["sale_date"] < cut], df[df["sale_date"] >= cut]
    raise ValueError(scheme)


# ---------------------------------------------------------------- pipelines
def _linear_pipeline(numeric: list[str], cats: list[str],
                     spline_cols: list[str] | None = None,
                     alpha: float = 1e-6) -> Pipeline:
    transformers = [("cat",
                     OneHotEncoder(handle_unknown="ignore", drop="first"),
                     cats)]
    if spline_cols:
        transformers.append(
            ("spl", SplineTransformer(degree=3, n_knots=6), spline_cols))
        passthrough = [c for c in numeric if c not in spline_cols]
    else:
        passthrough = numeric
    if passthrough:
        transformers.append(("num", "passthrough", passthrough))
    ct = ColumnTransformer(transformers, sparse_threshold=1.0)
    return Pipeline([("ct", ct),
                     ("reg", Ridge(alpha=alpha, solver="sparse_cg"))])


def _ml_frame(df: pd.DataFrame, geo: bool = True) -> pd.DataFrame:
    cols = ML_NUM if geo else [c for c in ML_NUM if c not in ("lat", "lng")]
    X = df[cols].copy()
    for c in CATS:
        X[c] = df[c].astype("category").cat.codes
    return X


# ---------------------------------------------------------------- metrics
def _metrics(y_level: np.ndarray, pred_level: np.ndarray) -> dict:
    pred_level = np.clip(pred_level, config.PRED_FLOOR, None)
    ape = np.abs(pred_level - y_level) / y_level
    ln_y, ln_p = np.log(y_level), np.log(pred_level)
    err = ln_y - ln_p
    return {"mdape": float(np.median(ape) * 100),
            "mape": float(np.mean(ape) * 100),
            "rmse_ln": float(np.sqrt(np.mean(err ** 2))),
            "r2_ln": float(1 - np.sum(err ** 2)
                           / np.sum((ln_y - ln_y.mean()) ** 2))}


# ---------------------------------------------------------------- Box–Cox
def _boxcox(y: np.ndarray, lam: float) -> np.ndarray:
    return np.log(y) if lam == 0 else (y ** lam - 1.0) / lam


def _inv_boxcox(z: np.ndarray, lam: float) -> np.ndarray:
    if lam == 0:
        return np.exp(z)
    base = np.clip(lam * z + 1.0, 1e-6, None)
    return base ** (1.0 / lam)


def fit_boxcox(train: pd.DataFrame, numeric: list[str], cats: list[str]):
    """Profile-likelihood choice of λ on the training sample."""
    y = train["price"].to_numpy(float)
    n = len(y)
    sum_ln_y = np.log(y).sum()
    rows = []
    for lam in config.BOXCOX_GRID:
        pipe = _linear_pipeline(numeric, cats)
        z = _boxcox(y, lam)
        pipe.fit(train, z)
        sse = float(np.sum((z - pipe.predict(train)) ** 2))
        ll = -0.5 * n * np.log(sse / n) + (lam - 1.0) * sum_ln_y
        rows.append({"lambda": lam, "loglik": ll})
    prof = pd.DataFrame(rows)
    lam_star = float(prof.loc[prof["loglik"].idxmax(), "lambda"])
    return lam_star, prof


# ---------------------------------------------------------------- registry
def model_registry() -> list[dict]:
    """The twenty models. ``fe`` lists categorical FE columns appended to
    the base dummies; ``numeric``/``spline`` define the design."""
    A = [
        dict(name="A1 Linear (levels)", group="A. Functional form",
             kind="linear_level", numeric=NUM_LEVELS, fe=["muni", "quarter"]),
        dict(name="A2 Semi-log", group="A. Functional form",
             kind="linear_log", numeric=NUM_LEVELS, fe=["muni", "quarter"]),
        dict(name="A3 Log-log", group="A. Functional form",
             kind="linear_log", numeric=NUM_LOGS, fe=["muni", "quarter"]),
        dict(name="A4 Box-Cox", group="A. Functional form",
             kind="boxcox", numeric=NUM_LEVELS, fe=["muni", "quarter"]),
        dict(name="A5 Semi-log + quadratics", group="A. Functional form",
             kind="linear_log", numeric=NUM_QUAD, fe=["muni", "quarter"]),
        dict(name="A6 Semi-log + splines", group="A. Functional form",
             kind="linear_log", numeric=NUM_LEVELS, fe=["muni", "quarter"],
             spline=["area", "age", "lot"]),
    ]
    B = [
        dict(name="B1 No time effects", group="B. Time effects",
             kind="linear_log", numeric=NUM_QUAD, fe=["muni"]),
        dict(name="B2 Year FE", group="B. Time effects",
             kind="linear_log", numeric=NUM_QUAD, fe=["muni", "sale_year_c"]),
        dict(name="B3 Quarter FE", group="B. Time effects",
             kind="linear_log", numeric=NUM_QUAD, fe=["muni", "quarter"]),
        dict(name="B4 Month FE", group="B. Time effects",
             kind="linear_log", numeric=NUM_QUAD, fe=["muni", "month"]),
    ]
    C = [
        dict(name="C1 No spatial controls", group="C. Spatial controls",
             kind="linear_log", numeric=NUM_QUAD, fe=["quarter"]),
        dict(name="C2 Municipality FE", group="C. Spatial controls",
             kind="linear_log", numeric=NUM_QUAD, fe=["muni", "quarter"]),
        dict(name="C3 Grid-cell FE (~5.5 km)", group="C. Spatial controls",
             kind="linear_log", numeric=NUM_QUAD, fe=["grid5", "quarter"]),
        dict(name="C4 Grid-cell FE (~1.1 km)", group="C. Spatial controls",
             kind="linear_log", numeric=NUM_QUAD, fe=["grid", "quarter"]),
    ]
    D = [
        dict(name="D1 OLS (muni + month FE)", group="D. Estimation method",
             kind="linear_log", numeric=NUM_QUAD, fe=["muni", "month"]),
        dict(name="D2 Ridge (cross-validated)", group="D. Estimation method",
             kind="ridge", numeric=NUM_QUAD, fe=["muni", "month"]),
        dict(name="D3 Random forest", group="D. Estimation method",
             kind="rf"),
        dict(name="D4 Gradient boosting", group="D. Estimation method",
             kind="hgb"),
        dict(name="D5 Gradient boosting, no coordinates",
             group="D. Estimation method", kind="hgb_nogeo"),
        dict(name="D6 Spatial k-NN comparables", group="D. Estimation method",
             kind="knn"),
    ]
    return A + B + C + D


# ---------------------------------------------------------------- fit/eval
def fit_predict(spec: dict, train: pd.DataFrame, test: pd.DataFrame):
    """Fit one registry entry, return level predictions on the test set."""
    kind = spec["kind"]
    y_train = train["price"].to_numpy(float)

    if kind in ("linear_level", "linear_log", "boxcox", "ridge"):
        cats = CATS + spec.get("fe", [])
        pipe = _linear_pipeline(spec["numeric"], cats,
                                spline_cols=spec.get("spline"))
        if kind == "linear_level":
            pipe.fit(train, y_train)
            return pipe.predict(test), pipe
        if kind in ("linear_log", "ridge"):
            if kind == "ridge":
                from sklearn.preprocessing import StandardScaler
                pipe.steps[-1:] = [
                    ("sc", StandardScaler(with_mean=False)),
                    ("reg", RidgeCV(alphas=np.logspace(-6, 2, 17),
                                    cv=3))]
            z = np.log(y_train)
            pipe.fit(train, z)
            resid = z - pipe.predict(train)
            smear = float(np.mean(np.exp(resid)))       # Duan (1983)
            return np.exp(pipe.predict(test)) * smear, pipe
        lam, _ = fit_boxcox(train, spec["numeric"], cats)
        pipe = _linear_pipeline(spec["numeric"], cats)
        pipe.fit(train, _boxcox(y_train, lam))
        spec["lambda"] = lam
        return _inv_boxcox(pipe.predict(test), lam), pipe

    if kind == "rf":
        m = RandomForestRegressor(n_estimators=120, min_samples_leaf=5,
                                  max_features=0.5, n_jobs=-1,
                                  random_state=config.SEED)
        m.fit(_ml_frame(train), np.log(y_train))
        return np.exp(m.predict(_ml_frame(test))), m

    if kind in ("hgb", "hgb_nogeo"):
        geo = kind == "hgb"
        m = HistGradientBoostingRegressor(
            max_iter=600, learning_rate=0.08, max_leaf_nodes=63,
            min_samples_leaf=20, l2_regularization=1e-2,
            random_state=config.SEED)
        m.fit(_ml_frame(train, geo), np.log(y_train))
        return np.exp(m.predict(_ml_frame(test, geo))), m

    if kind == "knn":
        # price-per-m2 of the 10 nearest sold neighbours (lat/lng, km-scaled)
        Xtr = train[["lat", "lng"]].to_numpy() * np.array([111.0, 78.0])
        Xte = test[["lat", "lng"]].to_numpy() * np.array([111.0, 78.0])
        m = KNeighborsRegressor(n_neighbors=10, weights="distance", n_jobs=-1)
        m.fit(Xtr, np.log(y_train) - train["ln_area"].to_numpy())
        return np.exp(m.predict(Xte) + test["ln_area"].to_numpy()), m

    raise ValueError(kind)


def run_horserace(df: pd.DataFrame, scheme: str) -> pd.DataFrame:
    """Evaluate the full registry under one split scheme.

    Under the forward-in-time split, test-period time categories (months,
    quarters, years) were never seen in training. A one-hot encoder would
    silently price them at the *baseline* period; the honest forecasting
    rule for a static model is to freeze the price level at the last
    period observed in training, so unseen time labels are remapped to the
    latest training label before prediction.
    """
    d = add_derived(df)
    d["sale_year_c"] = d["sale_year"].astype(str)
    train, test = split(d, scheme)
    if scheme == "temporal":
        test = test.copy()
        for col in ("quarter", "month", "sale_year_c"):
            last = train[col].max()
            test.loc[~test[col].isin(train[col].unique()), col] = last
    y_test = test["price"].to_numpy(float)
    rows = []
    for spec in model_registry():
        t0 = time.time()
        pred, _ = fit_predict(dict(spec), train, test)
        met = _metrics(y_test, pred)
        rows.append({"name": spec["name"], "group": spec["group"],
                     **met, "seconds": round(time.time() - t0, 1),
                     "n_train": len(train), "n_test": len(test)})
        print(f"  [{scheme}] {spec['name']:<38} "
              f"MdAPE={met['mdape']:5.1f}%  R2_ln={met['r2_ln']:.3f}  "
              f"({rows[-1]['seconds']}s)")
    return pd.DataFrame(rows)