"""Baseline classifiers and a quick, honest model comparison.

All models are sklearn ``Pipeline`` objects: ``preprocessor -> [svd] -> estimator``
so they accept the raw DataFrame directly and can be pickled as one object.
"""

from __future__ import annotations

import time

import numpy as np
from sklearn.base import clone
from sklearn.decomposition import TruncatedSVD
from sklearn.ensemble import HistGradientBoostingClassifier, RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, log_loss
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline

from . import MAX_JOBS
from .calibration import ece

MODEL_NAMES = ("logreg", "hgb", "rf")


def make_model(name: str, preprocessor, sparse_input: bool, seed: int = 0) -> Pipeline:
    """Build one of the baseline pipelines.

    * ``logreg``: multinomial logistic regression (L2, C=1). Linear, so it also
      supports exact robustness certificates (see ``robustness``).
    * ``hgb``: histogram gradient boosting. It needs dense input, so for sparse
      text features we first project to 50 dims with TruncatedSVD (LSA).
    * ``rf``: random forest, 200 trees, at most ``MAX_JOBS`` workers.
    """
    pre = clone(preprocessor)
    if name == "logreg":
        est = LogisticRegression(C=1.0, max_iter=2000)
        steps = [("pre", pre), ("clf", est)]
    elif name == "hgb":
        est = HistGradientBoostingClassifier(
            max_iter=150, learning_rate=0.1, early_stopping=True, validation_fraction=0.1,
            n_iter_no_change=10, random_state=seed,
        )
        steps = [("pre", pre)]
        if sparse_input:
            steps.append(("svd", TruncatedSVD(n_components=50, random_state=seed)))
        steps.append(("clf", est))
    elif name == "rf":
        est = RandomForestClassifier(
            n_estimators=200, min_samples_leaf=1, n_jobs=MAX_JOBS, random_state=seed
        )
        steps = [("pre", pre), ("clf", est)]
    else:
        raise ValueError(f"unknown model {name!r}; choose from {MODEL_NAMES}")
    return Pipeline(steps)


def compare_models(
    preprocessor,
    X_train,
    y_train: np.ndarray,
    sparse_input: bool,
    names=MODEL_NAMES,
    val_frac: float = 0.2,
    seed: int = 0,
) -> list[dict]:
    """Fit each model on 80% of the training rows and score on the other 20%.

    Only training data is used, so the calibration and test splits stay
    untouched for the guarantees computed later. Rows are sorted best-first by
    validation log-loss (a proper scoring rule: rewards good probabilities, not
    just good argmax).
    """
    from .data import _safe_stratify

    idx = np.arange(len(y_train))
    tr, va = train_test_split(
        idx, test_size=val_frac, random_state=seed, stratify=_safe_stratify(y_train)
    )
    n_classes = int(y_train.max()) + 1
    rows = []
    for name in names:
        model = make_model(name, preprocessor, sparse_input, seed)
        t0 = time.perf_counter()
        model.fit(X_train.iloc[tr], y_train[tr])
        fit_s = time.perf_counter() - t0
        P = full_proba(model, X_train.iloc[va], n_classes)
        rows.append(
            {
                "model": name,
                "val_accuracy": float(accuracy_score(y_train[va], P.argmax(1))),
                "val_log_loss": float(log_loss(y_train[va], np.clip(P, 1e-12, 1), labels=range(n_classes))),
                "val_ece": float(ece(P, y_train[va])),
                "fit_seconds": float(fit_s),
            }
        )
    rows.sort(key=lambda r: r["val_log_loss"])
    return rows


def full_proba(model, X, n_classes: int) -> np.ndarray:
    """``predict_proba`` padded to all ``n_classes`` columns.

    A model trained on a subset that happened to miss a class returns fewer
    columns; we put zeros there so column ``k`` always means class ``k``.
    """
    P = model.predict_proba(X)
    classes = model.classes_ if hasattr(model, "classes_") else model[-1].classes_
    if P.shape[1] == n_classes:
        return P
    out = np.zeros((P.shape[0], n_classes))
    out[:, np.asarray(classes, dtype=int)] = P
    return out