# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Arm C: LightGBM on log price (CatBoost / XGBoost as robustness), tuned on validation only (random search +
early stopping). Arm D: AVM = LightGBM point model on the full feature set (including as-of comparable features)
+ LightGBM quantile models for q05..q95 + conformal rescaling of the 90 % interval on validation."""
from __future__ import annotations

import json
import random

import numpy as np
import polars as pl

from .common import Prediction, QUANTILES, apply_scale, categorical_features, conformal_scale, numeric_features

SEARCH_SPACE = {
    "num_leaves": [31, 63, 127, 255],
    "learning_rate": [0.02, 0.03, 0.05, 0.08],
    "min_child_samples": [20, 50, 100, 200],
    "feature_fraction": [0.6, 0.8, 1.0],
    "bagging_fraction": [0.7, 0.9, 1.0],
    "lambda_l2": [0.0, 1.0, 10.0],
}


def _to_pandas(df: pl.DataFrame, num: list[str], cat: list[str]):
    pdf = df.select(num + cat).to_pandas()
    for c in cat:
        pdf[c] = pdf[c].astype("category")
    for c in num:
        pdf[c] = pdf[c].astype(float)
    return pdf


class LightGBMModel:
    def __init__(self, n_trials: int = 12, seed: int = 20261027, objective: str = "regression", alpha: float | None = None, max_rounds: int = 4000, exclude_prefixes: tuple[str, ...] = ()):
        self.exclude_prefixes = exclude_prefixes
        self.n_trials = n_trials
        self.seed = seed
        self.objective = objective
        self.alpha = alpha
        self.max_rounds = max_rounds
        self.params: dict = {}
        self.model = None
        self.num: list[str] = []
        self.cat: list[str] = []
        self.search_log: list[dict] = []

    def _base_params(self, p: dict) -> dict:
        out = {"objective": self.objective, "verbose": -1, "seed": self.seed, "num_threads": 0, "bagging_freq": 1, **p}
        if self.objective == "quantile":
            out["alpha"] = self.alpha
        return out

    def fit(self, train: pl.DataFrame, valid: pl.DataFrame, params: dict | None = None) -> "LightGBMModel":
        import lightgbm as lgb

        self.num = [c for c in numeric_features(train) if not c.startswith(self.exclude_prefixes)]
        self.cat = categorical_features(train)
        Xtr, Xva = _to_pandas(train, self.num, self.cat), _to_pandas(valid, self.num, self.cat)
        # align categories
        for c in self.cat:
            cats = Xtr[c].cat.categories
            Xva[c] = Xva[c].cat.set_categories(cats)
        ytr, yva = train["log_price"].to_numpy(), valid["log_price"].to_numpy()
        dtr = lgb.Dataset(Xtr, ytr, categorical_feature=self.cat, free_raw_data=False, params={"feature_pre_filter": False})
        dva = lgb.Dataset(Xva, yva, reference=dtr, categorical_feature=self.cat, free_raw_data=False)
        rng = random.Random(self.seed)
        candidates = [params] if params else [{k: rng.choice(v) for k, v in SEARCH_SPACE.items()} for _ in range(self.n_trials)]
        best = None
        for p in candidates:
            m = lgb.train(self._base_params(p), dtr, num_boost_round=self.max_rounds, valid_sets=[dva], callbacks=[lgb.early_stopping(100, verbose=False)])
            pred = m.predict(Xva, num_iteration=m.best_iteration)
            score = float(np.mean(np.abs(pred - yva))) if self.objective != "quantile" else float(_pinball(yva, pred, self.alpha))
            self.search_log.append({"params": p, "best_iteration": m.best_iteration, "valid_score": score})
            if best is None or score < best[0]:
                best = (score, p, m)
        self.params = {**best[1], "num_boost_round": best[2].best_iteration}
        self.model = best[2]
        return self

    def predict_log(self, df: pl.DataFrame) -> np.ndarray:
        X = _to_pandas(df, self.num, self.cat)
        for c in self.cat:
            X[c] = X[c].cat.set_categories(self.model.pandas_categorical[self.cat.index(c)])
        return self.model.predict(X, num_iteration=self.model.best_iteration)

    def predict(self, df: pl.DataFrame) -> Prediction:
        return Prediction(tx_id=df["tx_id"].to_numpy(), point=np.exp(self.predict_log(df)), quantiles=None)


def _pinball(y, q, a):
    d = y - q
    return np.mean(np.maximum(a * d, (a - 1) * d))


class AVM:
    """Point model + five quantile models + conformal scaling (validation)."""

    def __init__(self, n_trials: int = 12, seed: int = 20261027):
        self.point = LightGBMModel(n_trials=n_trials, seed=seed)
        self.qmodels: dict[float, LightGBMModel] = {}
        self.scale = 1.0
        self.n_trials = n_trials
        self.seed = seed

    def fit(self, train: pl.DataFrame, valid: pl.DataFrame) -> "AVM":
        self.point.fit(train, valid)
        p = {k: v for k, v in self.point.params.items() if k != "num_boost_round"}
        for a in QUANTILES:
            m = LightGBMModel(n_trials=1, seed=self.seed, objective="quantile", alpha=a)
            m.fit(train, valid, params=p)
            self.qmodels[a] = m
        qv = self._q_log(valid)
        self.scale = conformal_scale(self.point.predict_log(valid), qv, valid["log_price"].to_numpy(), 0.90)
        return self

    def _q_log(self, df: pl.DataFrame) -> np.ndarray:
        q = np.column_stack([self.qmodels[a].predict_log(df) for a in QUANTILES])
        return np.sort(q, axis=1)  # enforce monotone quantiles

    def predict(self, df: pl.DataFrame) -> Prediction:
        p = self.point.predict_log(df)
        q = apply_scale(self._q_log(df), self.scale)
        # keep the point inside [q05, q95]
        q[:, 0] = np.minimum(q[:, 0], p - 1e-4)
        q[:, 4] = np.maximum(q[:, 4], p + 1e-4)
        q = np.sort(q, axis=1)
        return Prediction(tx_id=df["tx_id"].to_numpy(), point=np.exp(p), quantiles=np.exp(q))


class CatBoostModel:
    def __init__(self, seed: int = 20261027, iterations: int = 4000, exclude_prefixes: tuple[str, ...] = ()):
        self.exclude_prefixes = exclude_prefixes
        self.seed, self.iterations = seed, iterations
        self.model = None
        self.num: list[str] = []
        self.cat: list[str] = []

    def fit(self, train: pl.DataFrame, valid: pl.DataFrame) -> "CatBoostModel":
        from catboost import CatBoostRegressor, Pool

        self.num, self.cat = [c for c in numeric_features(train) if not c.startswith(self.exclude_prefixes)], categorical_features(train)
        Xtr = train.select(self.num + self.cat).to_pandas()
        Xva = valid.select(self.num + self.cat).to_pandas()
        for c in self.cat:
            Xtr[c] = Xtr[c].astype(str)
            Xva[c] = Xva[c].astype(str)
        self.model = CatBoostRegressor(iterations=self.iterations, learning_rate=0.05, depth=8, loss_function="MAE", random_seed=self.seed, verbose=0, od_type="Iter", od_wait=100, thread_count=-1)
        self.model.fit(Pool(Xtr, train["log_price"].to_numpy(), cat_features=self.cat), eval_set=Pool(Xva, valid["log_price"].to_numpy(), cat_features=self.cat), use_best_model=True)
        return self

    def predict(self, df: pl.DataFrame) -> Prediction:
        X = df.select(self.num + self.cat).to_pandas()
        for c in self.cat:
            X[c] = X[c].astype(str)
        return Prediction(tx_id=df["tx_id"].to_numpy(), point=np.exp(self.model.predict(X)), quantiles=None)


class XGBoostModel:
    def __init__(self, seed: int = 20261027, rounds: int = 4000, exclude_prefixes: tuple[str, ...] = ()):
        self.exclude_prefixes = exclude_prefixes
        self.seed, self.rounds = seed, rounds
        self.model = None
        self.num: list[str] = []
        self.cat: list[str] = []

    def fit(self, train: pl.DataFrame, valid: pl.DataFrame) -> "XGBoostModel":
        import xgboost as xgb

        self.num, self.cat = [c for c in numeric_features(train) if not c.startswith(self.exclude_prefixes)], categorical_features(train)
        Xtr, Xva = _to_pandas(train, self.num, self.cat), _to_pandas(valid, self.num, self.cat)
        for c in self.cat:
            Xva[c] = Xva[c].cat.set_categories(Xtr[c].cat.categories)
        self.model = xgb.XGBRegressor(n_estimators=self.rounds, learning_rate=0.05, max_depth=8, subsample=0.8, colsample_bytree=0.8, enable_categorical=True, tree_method="hist", random_state=self.seed, early_stopping_rounds=100, objective="reg:absoluteerror")
        self.model.fit(Xtr, train["log_price"].to_numpy(), eval_set=[(Xva, valid["log_price"].to_numpy())], verbose=False)
        self._cats = {c: Xtr[c].cat.categories for c in self.cat}
        return self

    def predict(self, df: pl.DataFrame) -> Prediction:
        X = _to_pandas(df, self.num, self.cat)
        for c in self.cat:
            X[c] = X[c].cat.set_categories(self._cats[c])
        return Prediction(tx_id=df["tx_id"].to_numpy(), point=np.exp(self.model.predict(X)), quantiles=None)


def dump_search_log(model: LightGBMModel, path) -> None:
    with open(path, "w") as fh:
        json.dump({"params": model.params, "search_log": model.search_log}, fh, indent=1, default=float)