# Author: Simon-Pierre Boucher — contact@spboucher.ai
#
"""Hedonic OLS specifications A-E (standardized covariates, HC3 errors)."""

import pandas as pd
import statsmodels.api as sm
from sklearn.preprocessing import StandardScaler

from .config import STRUCTURAL_VARS
from .references import SIM_COLS


def standardized_design(df, columns):
    """Z-score the selected columns and prepend a constant."""
    X = pd.DataFrame(
        StandardScaler().fit_transform(df[columns]),
        columns=columns,
        index=df.index,
    )
    return sm.add_constant(X)


def fit_ols(df, columns, y=None):
    """OLS of log-price on standardized covariates with HC3 robust errors."""
    if y is None:
        y = df["log_price"]
    return sm.OLS(y, standardized_design(df, columns)).fit(cov_type="HC3")


def fit_all_models(df):
    """Fit specifications A-E and return them with the significant-similarity list.

    A: structural only            B: A + description length
    C: A + 20 similarities        D: B + 20 similarities (full)
    E: B + similarities significant at 5% in D (parsimonious)
    """
    specs = {
        "A": STRUCTURAL_VARS,
        "B": STRUCTURAL_VARS + ["remarks_length"],
        "C": STRUCTURAL_VARS + SIM_COLS,
        "D": STRUCTURAL_VARS + ["remarks_length"] + SIM_COLS,
    }
    models = {name: fit_ols(df, cols) for name, cols in specs.items()}
    sig_sims = [v for v in SIM_COLS if models["D"].pvalues.get(v, 1) < 0.05]
    specs["E"] = STRUCTURAL_VARS + ["remarks_length"] + sig_sims
    models["E"] = fit_ols(df, specs["E"])
    return models, specs, sig_sims