# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Task A models (property-level price at a future sale date). All expose fit(X, y, meta) / predict(X, meta) → point (log price)
and, where available, predict_quantiles(X, meta, taus) → (n, len(taus)). `meta` carries the categorical keys for the
hierarchical model (municipality / MRC / region codes, horizon)."""
from __future__ import annotations

import os

import numpy as np

# NOTE (macOS): LightGBM / XGBoost (libomp) and PyTorch cannot train in the same process without a segmentation fault or a
# deadlock in the OpenMP runtime. Tree models live here (lazy imports); neural models live in models_a_torch.py; the runner
# executes the two families in separate processes (`--family trees|torch`).
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import ElasticNet, Ridge
from sklearn.svm import SVR

TAUS = np.array([0.05, 0.25, 0.5, 0.75, 0.95])


class OLS:
    name = "hedonic_ols"
    def fit(self, X, y, meta=None):
        Xa = np.column_stack([np.ones(len(X)), X]); self.b = np.linalg.lstsq(Xa, y, rcond=None)[0]
        r = y - Xa @ self.b; self.sigma = r.std(); return self
    def predict(self, X, meta=None):
        return np.column_stack([np.ones(len(X)), X]) @ self.b
    def predict_quantiles(self, X, meta=None, taus=TAUS):
        from scipy.stats import norm
        m = self.predict(X); return m[:, None] + self.sigma * norm.ppf(taus)[None, :]


class RidgeM(OLS):
    name = "ridge"
    def fit(self, X, y, meta=None):
        self.mu = X.mean(0); self.sd = X.std(0) + 1e-9; self.m = Ridge(alpha=1.0).fit((X - self.mu) / self.sd, y)
        self.sigma = (y - self.m.predict((X - self.mu) / self.sd)).std(); return self
    def predict(self, X, meta=None):
        return self.m.predict((X - self.mu) / self.sd)


class ElasticNetM(RidgeM):
    name = "elastic_net"
    def fit(self, X, y, meta=None):
        self.mu = X.mean(0); self.sd = X.std(0) + 1e-9; self.m = ElasticNet(alpha=0.001, l1_ratio=0.5, max_iter=3000).fit((X - self.mu) / self.sd, y)
        self.sigma = (y - self.m.predict((X - self.mu) / self.sd)).std(); return self


class RF:
    name = "random_forest"
    def fit(self, X, y, meta=None):
        self.m = RandomForestRegressor(n_estimators=150, max_depth=18, min_samples_leaf=5, max_features=0.5, n_jobs=int(os.environ.get("WP23_JOBS", "8")), random_state=0).fit(X, y); return self
    def predict(self, X, meta=None):
        return self.m.predict(X)


class XGB:
    name = "xgboost"
    def fit(self, X, y, meta=None):
        import xgboost as xgb
        self.m = xgb.XGBRegressor(n_estimators=600, learning_rate=0.05, max_depth=8, subsample=0.8, colsample_bytree=0.8, tree_method="hist", n_jobs=int(os.environ.get("WP23_JOBS", "8")), random_state=0).fit(X, y); return self
    def predict(self, X, meta=None):
        return self.m.predict(X)


class LGBM:
    name = "lightgbm"
    def __init__(self, quantiles: bool = True):
        self.quantiles = quantiles
    def fit(self, X, y, meta=None):
        import lightgbm as lgb
        kw = dict(n_estimators=800, learning_rate=0.04, num_leaves=63, min_child_samples=30, subsample=0.8, subsample_freq=1, colsample_bytree=0.8, n_jobs=int(os.environ.get("WP23_JOBS", "8")), random_state=0, verbose=-1)
        self.m = lgb.LGBMRegressor(objective="l2", **kw).fit(X, y)
        if self.quantiles:
            self.q = {t: lgb.LGBMRegressor(objective="quantile", alpha=float(t), **{**kw, "n_estimators": 400}).fit(X, y) for t in TAUS}
        return self
    def predict(self, X, meta=None):
        return self.m.predict(X)
    def predict_quantiles(self, X, meta=None, taus=TAUS):
        Q = np.column_stack([self.q[t].predict(X) for t in taus]); return np.sort(Q, axis=1)


class SVRM:
    name = "svr"
    def fit(self, X, y, meta=None, n_max=15000):
        rng = np.random.default_rng(0); idx = rng.choice(len(X), min(n_max, len(X)), replace=False)
        self.mu = X.mean(0); self.sd = X.std(0) + 1e-9; self.m = SVR(C=3.0, epsilon=0.05, gamma="scale").fit((X[idx] - self.mu) / self.sd, y[idx]); return self
    def predict(self, X, meta=None):
        return self.m.predict((X - self.mu) / self.sd)


def registry(family: str = "trees", fast: bool = False) -> list:
    if family == "torch":
        from wp23.models_a_torch import HSHN, MLP
        return [MLP(), HSHN("partial"), HSHN("full"), HSHN("none")]
    base = [OLS(), RidgeM(), ElasticNetM(), LGBM(), XGB()]
    if not fast:
        base += [RF(), SVRM()]
    return base