# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""Design matrices and the M1–M5 hedonic specification ladder.

The semi-logarithmic hedonic equation regresses ln(price) on structural
attributes, dwelling-type/ownership dummies and location fixed effects.
FSA fixed effects are absorbed with ``linearmodels.AbsorbingLS`` (numerically
identical to full-dummy OLS); standard errors are clustered by FSA.
"""
import numpy as np
import pandas as pd
import statsmodels.api as sm
from linearmodels.iv.absorbing import AbsorbingLS

from .config import STRUCT


def design(data: pd.DataFrame, extra: pd.DataFrame | None = None) -> pd.DataFrame:
    """Structural regressors + dwelling-type, ownership and category dummies."""
    blocks = [data[STRUCT].astype(float)]
    if extra is not None:
        blocks.append(extra)
    blocks += [pd.get_dummies(data[col], prefix=pfx, drop_first=True).astype(float)
               for col, pfx in (("btype_c", "bt"), ("own_c", "ow"), ("cat", "cat"))]
    return sm.add_constant(pd.concat(blocks, axis=1))


def fit_absorbing(data: pd.DataFrame, extra: pd.DataFrame | None = None):
    """FSA-fixed-effects hedonic model with FSA-clustered standard errors.

    Returns the fitted AbsorbingLS results and the design column order.
    """
    X = design(data, extra=extra)
    res = AbsorbingLS(data["ln_price"], X,
                      absorb=data[["fsa_c"]].astype("category"),
                      drop_absorbed=True).fit(cov_type="clustered",
                                              clusters=data[["fsa_c"]])
    return res, X.columns.tolist()


def coef_table(res, columns=None) -> pd.DataFrame:
    """Coefficient / SE / p-value table indexed by regressor name.

    Works for both statsmodels (``bse``) and linearmodels (``std_errors``);
    uses the estimator's own parameter index, which may exclude regressors
    dropped as collinear with the absorbed effects.
    """
    se = res.std_errors if hasattr(res, "std_errors") else res.bse
    index = res.params.index if hasattr(res.params, "index") else columns
    return pd.DataFrame({
        "coef": pd.Series(np.asarray(res.params).ravel(), index=index),
        "se": pd.Series(np.asarray(se).ravel(), index=index),
        "p": pd.Series(np.asarray(res.pvalues).ravel(), index=index),
    })


def fsa_fixed_effects(data: pd.DataFrame, res, columns=None) -> pd.Series:
    """Recover the absorbed FSA effects as within-FSA mean residuals of X'b."""
    beta = pd.Series(np.asarray(res.params).ravel(), index=res.params.index)
    xb = design(data).reindex(columns=beta.index, fill_value=0.0).values @ beta.values
    return (data["ln_price"] - xb).groupby(data["fsa_c"]).mean()


def predict_with_fe(data: pd.DataFrame, res, columns, fe: pd.Series) -> np.ndarray:
    """Linear prediction X'b + absorbed FSA effect."""
    beta = pd.Series(np.asarray(res.params).ravel(), index=res.params.index)
    xb = design(data).reindex(columns=beta.index, fill_value=0.0).values @ beta.values
    return xb + data["fsa_c"].map(fe).values


def specification_ladder(sample: pd.DataFrame) -> dict:
    """Estimate the M1–M5 ladder; returns {name: (results, columns, data)}.

    M1–M3 use the house subsample (structural; +type/ownership; +province);
    M4 is houses with FSA fixed effects; M5 is the grand model over all
    residential dwellings with FSA fixed effects.
    """
    houses = sample[sample["cat"] == "house"]
    out = {}

    X1 = sm.add_constant(houses[STRUCT].astype(float))
    out["M1"] = (sm.OLS(houses["ln_price"], X1)
                 .fit(cov_type="cluster", cov_kwds={"groups": houses["fsa_c"]}),
                 X1.columns.tolist(), houses)

    X2 = design(houses)
    out["M2"] = (sm.OLS(houses["ln_price"], X2)
                 .fit(cov_type="cluster", cov_kwds={"groups": houses["fsa_c"]}),
                 X2.columns.tolist(), houses)

    X3 = X2.join(pd.get_dummies(houses["prov"], prefix="pv", drop_first=True).astype(float))
    out["M3"] = (sm.OLS(houses["ln_price"], X3)
                 .fit(cov_type="cluster", cov_kwds={"groups": houses["fsa_c"]}),
                 X3.columns.tolist(), houses)

    res4, cols4 = fit_absorbing(houses)
    out["M4"] = (res4, cols4, houses)

    res5, cols5 = fit_absorbing(sample)
    out["M5"] = (res5, cols5, sample)
    return out


def duan_smearing(residuals: np.ndarray) -> float:
    """Duan (1983) smearing factor for retransformation from logs."""
    return float(np.mean(np.exp(residuals)))