"""Ridge / elastic-net and LightGBM baselines on hand-crafted origin features (identical information set).
Targets: (i) cumulative H-minute return (location), (ii) log realised variance over H (scale). Predictive
distribution: Student-t with location from the mean model (ridge/LGBM), scale from exp(log-variance model) and nu
from the kurtosis of standardised training residuals (clipped to [2.5, 50])."""
from __future__ import annotations

import numpy as np

from ..data.windows import SymbolGrid, N_DAY
from .naive import _hist_returns

FEATURE_NAMES = ["log_rv30", "log_rv390", "log_rv5d", "ret30", "ret390", "ret_day", "overnight", "relvol30",
                 "range30", "minute_frac", "time_to_close", "weekday", "early_close", "log_sigma", "skew390", "kurt390",
                 "sign_last", "ret5", "rv_ratio_30_390"]


def origin_features(g: SymbolGrid, s: int, m: int, sigma: float) -> np.ndarray | None:
    r, sl = _hist_returns(g, s, m, 6)
    if len(r) < 60:
        return None
    r30, r390, r5 = r[-30:], r[-390:], r[-5:]
    lc_today = g.flat_lc[s * N_DAY: s * N_DAY + m + 1]
    ok = ~np.isnan(lc_today)
    ret_day = float(lc_today[ok][-1] - lc_today[ok][0]) if ok.sum() >= 2 else 0.0
    gap = g.gap[s]; gap = 0.0 if np.isnan(gap) else float(gap)
    base = g.flat_base[s * N_DAY + max(0, m - 29): s * N_DAY + m + 1]; mk = g.flat_mask[s * N_DAY + max(0, m - 29): s * N_DAY + m + 1]
    seas = g.flat_seas[s * N_DAY + max(0, m - 29): s * N_DAY + m + 1]
    rel = base[mk, 4] - seas[mk]; rel = rel[~np.isnan(rel)]
    relvol30 = float(rel.mean()) if len(rel) else 0.0
    rng30 = float((base[mk, 1] - base[mk, 2]).mean()) if mk.any() else 0.0
    sd390 = r390.std() + 1e-12
    f = [np.log(np.mean(r30 ** 2) + 1e-12), np.log(np.mean(r390 ** 2) + 1e-12), np.log(np.mean(r ** 2) + 1e-12),
         r30.sum() / sigma, r390.sum() / sigma, ret_day / sigma, gap / sigma, relvol30, rng30 / sigma, m / 389.0,
         (g.n_slots[s] - 1 - m) / 389.0, g.weekday[s] / 4.0, float(g.early[s]), np.log(sigma),
         float(np.mean(((r390 - r390.mean()) / sd390) ** 3)), float(np.mean(((r390 - r390.mean()) / sd390) ** 4)),
         float(np.sign(r[-1])), r5.sum() / sigma, np.log((np.mean(r30 ** 2) + 1e-12) / (np.mean(r390 ** 2) + 1e-12))]
    return np.array(f, dtype=np.float64)


class _TwoPart:
    """Common fit/predict for location + log-variance models with Student-t assembly."""

    def __init__(self, name: str):
        self.name = name
        self.loc: dict[int, object] = {}; self.lv: dict[int, object] = {}; self.nu: dict[int, float] = {}
        self.scaler = None

    def _new(self):
        raise NotImplementedError

    def fit(self, X: np.ndarray, cum: np.ndarray, rv: np.ndarray, sigma: np.ndarray, H: int, Xv=None, cumv=None, rvv=None, sigv=None) -> None:
        y_loc = cum / sigma
        y_lv = np.log(rv / sigma ** 2 + 1e-9)
        self.loc[H] = self._new().fit(X, y_loc)
        self.lv[H] = self._new().fit(X, y_lv)
        var = np.exp(self.lv[H].predict(X)) * sigma ** 2
        z = (cum - self.loc[H].predict(X) * sigma) / np.sqrt(var)
        k = float(np.mean(z ** 4) / np.mean(z ** 2) ** 2)
        self.nu[H] = float(np.clip(4 + 6 / max(k - 3, 0.12), 2.5, 50))

    def predict_batch(self, X: np.ndarray, sigma: np.ndarray, H: int) -> dict:
        mu = self.loc[H].predict(X) * sigma
        var = np.exp(self.lv[H].predict(X)) * sigma ** 2
        nu = self.nu[H]
        return {"mu": mu, "scale": np.sqrt(var * (nu - 2) / nu), "nu": np.full(len(X), nu), "var": var}


class RidgeBaseline(_TwoPart):
    def __init__(self, alpha: float = 1.0):
        super().__init__("ridge"); self.alpha = alpha

    def _new(self):
        from sklearn.linear_model import Ridge
        from sklearn.pipeline import make_pipeline
        from sklearn.preprocessing import StandardScaler
        return make_pipeline(StandardScaler(), Ridge(alpha=self.alpha))


class ElasticNetBaseline(_TwoPart):
    def __init__(self, alpha: float = 1e-3, l1_ratio: float = 0.5):
        super().__init__("elastic_net"); self.alpha, self.l1 = alpha, l1_ratio

    def _new(self):
        from sklearn.linear_model import ElasticNet
        from sklearn.pipeline import make_pipeline
        from sklearn.preprocessing import StandardScaler
        return make_pipeline(StandardScaler(), ElasticNet(alpha=self.alpha, l1_ratio=self.l1, max_iter=5000))


class LightGBMBaseline(_TwoPart):
    def __init__(self, n_estimators: int = 400, learning_rate: float = 0.03, num_leaves: int = 31, seed: int = 0):
        super().__init__("lightgbm"); self.kw = dict(n_estimators=n_estimators, learning_rate=learning_rate,
                                                     num_leaves=num_leaves, subsample=0.8, subsample_freq=1,
                                                     colsample_bytree=0.8, random_state=seed, verbose=-1, n_jobs=8)

    def _new(self):
        import lightgbm as lgb
        return lgb.LGBMRegressor(**self.kw)