"""Baseline ladder B0-B6 with train-only preprocessing and an equal tuning budget per information set."""
from __future__ import annotations

import json
import random
from dataclasses import dataclass, field

import numpy as np
import polars as pl
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.metrics import log_loss

INFO_SETS = {
    "B1": ["fundamentals"],
    "B2": ["fundamentals", "daily"],
    "B3": ["fundamentals", "daily", "intraday_30m"],
    "B4": ["fundamentals", "daily", "intraday_5m"],
    "B5": ["fundamentals", "daily", "intraday_1m"],
    "B5all": ["fundamentals", "daily", "intraday_30m", "intraday_5m", "intraday_1m"],
}


@dataclass
class Preprocessor:
    """Winsorise (1/99 %), impute (median) and standardise using TRAINING statistics only."""
    cols: list[str]
    lo: dict = field(default_factory=dict)
    hi: dict = field(default_factory=dict)
    med: dict = field(default_factory=dict)
    mu: dict = field(default_factory=dict)
    sd: dict = field(default_factory=dict)
    fitted_on_rows: int = 0

    def fit(self, X: pl.DataFrame) -> "Preprocessor":
        for c in self.cols:
            s = X[c].cast(pl.Float64)
            s = s.filter(s.is_finite())
            if s.len() == 0:
                self.lo[c], self.hi[c], self.med[c], self.mu[c], self.sd[c] = 0.0, 0.0, 0.0, 0.0, 1.0
                continue
            lo, hi = float(s.quantile(0.01)), float(s.quantile(0.99))
            w = s.clip(lo, hi)
            self.lo[c], self.hi[c], self.med[c] = lo, hi, float(w.median())
            self.mu[c] = float(w.mean())
            sd = float(w.std()) if w.len() > 1 else 1.0
            self.sd[c] = sd if sd and sd > 1e-12 else 1.0
        self.fitted_on_rows = X.height
        return self

    def transform(self, X: pl.DataFrame) -> np.ndarray:
        out = np.empty((X.height, len(self.cols)), dtype=np.float64)
        for j, c in enumerate(self.cols):
            v = X[c].cast(pl.Float64).to_numpy().astype(np.float64)
            v = np.where(np.isfinite(v), v, np.nan)
            v = np.clip(v, self.lo[c], self.hi[c])
            v = np.where(np.isnan(v), self.med[c], v)
            out[:, j] = (v - self.mu[c]) / self.sd[c]
        return out

    def to_dict(self) -> dict:
        return {"cols": self.cols, "lo": self.lo, "hi": self.hi, "med": self.med, "mu": self.mu, "sd": self.sd,
                "fitted_on_rows": self.fitted_on_rows}


def columns_for(info_set: str, fc: dict[str, list[str]]) -> list[str]:
    cols = []
    for fam in INFO_SETS[info_set]:
        cols += fc[fam]
    return cols


# ----------------------------------------------------------------------------- B0
def base_rate_predict(train_y: np.ndarray, n: int) -> np.ndarray:
    return np.full(n, float(np.clip(train_y.mean(), 1e-4, 1 - 1e-4)))


def firm_history_predict(train: pl.DataFrame, test: pl.DataFrame, y: str) -> np.ndarray:
    """B0b: trailing firm-level adverse rate (from training rows only, shrunk to the pooled rate with prior weight 8)."""
    pooled = float(train[y].mean())
    g = train.group_by("entity_id").agg([pl.col(y).sum().alias("k"), pl.len().alias("n")])
    t = test.select("entity_id").join(g, on="entity_id", how="left").with_columns([pl.col("k").fill_null(0), pl.col("n").fill_null(0)])
    k, n = t["k"].to_numpy().astype(float), t["n"].to_numpy().astype(float)
    return np.clip((k + 8 * pooled) / (n + 8), 1e-4, 1 - 1e-4)


# ------------------------------------------------------------------------ logistic
def fit_logistic(Xtr: np.ndarray, ytr: np.ndarray, Xva: np.ndarray, yva: np.ndarray, C_grid: list[float], seed: int) -> tuple:
    best, best_ll, best_C = None, np.inf, None
    for C in C_grid:
        m = LogisticRegression(C=C, penalty="l2", solver="lbfgs", max_iter=2000, random_state=seed)
        m.fit(Xtr, ytr)
        ll = log_loss(yva, np.clip(m.predict_proba(Xva)[:, 1], 1e-6, 1 - 1e-6))
        if ll < best_ll:
            best, best_ll, best_C = m, ll, C
    return best, {"C": best_C, "val_logloss": best_ll}


def fit_ridge(Xtr, ytr, Xva, yva, alpha_grid=(0.1, 1, 10, 100, 1000)):
    best, best_mse, best_a = None, np.inf, None
    for a in alpha_grid:
        m = Ridge(alpha=a).fit(Xtr, ytr)
        mse = float(np.mean((m.predict(Xva) - yva) ** 2))
        if mse < best_mse:
            best, best_mse, best_a = m, mse, a
    return best, {"alpha": best_a, "val_mse": best_mse}


# ------------------------------------------------------------------------ LightGBM
def _sample_params(space: dict, rng: random.Random) -> dict:
    return {k: rng.choice(v) for k, v in space.items()}


def fit_lightgbm(Xtr, ytr, Xva, yva, cfg: dict, seed: int, objective: str = "binary", n_jobs: int = 8) -> tuple:
    import lightgbm as lgb
    rng = random.Random(seed)
    best, best_score, best_params, trials = None, np.inf, None, []
    for t in range(int(cfg["budget_trials"])):
        p = _sample_params(cfg["search_space"], rng)
        params = dict(objective=objective, learning_rate=p["learning_rate"], num_leaves=p["num_leaves"],
                      min_child_samples=p["min_child_samples"], feature_fraction=p["feature_fraction"],
                      bagging_fraction=p["bagging_fraction"], bagging_freq=1 if p["bagging_fraction"] < 1 else 0,
                      lambda_l2=p["lambda_l2"], n_estimators=int(cfg["n_estimators_max"]), random_state=seed + t,
                      n_jobs=n_jobs, verbose=-1)
        m = lgb.LGBMClassifier(**params) if objective == "binary" else lgb.LGBMRegressor(**params)
        m.fit(Xtr, ytr, eval_set=[(Xva, yva)], callbacks=[lgb.early_stopping(int(cfg["early_stopping_rounds"]), verbose=False)])
        if objective == "binary":
            score = log_loss(yva, np.clip(m.predict_proba(Xva)[:, 1], 1e-6, 1 - 1e-6))
        else:
            score = float(np.mean((m.predict(Xva) - yva) ** 2))
        trials.append({"trial": t, "params": p, "best_iteration": int(m.best_iteration_ or 0), "val_score": score})
        if score < best_score:
            best, best_score, best_params = m, score, {**p, "best_iteration": int(m.best_iteration_ or 0)}
    return best, {"val_score": best_score, "params": best_params, "trials": trials}


def predict_proba(model, X: np.ndarray) -> np.ndarray:
    return np.clip(model.predict_proba(X)[:, 1], 1e-6, 1 - 1e-6)