"""End-to-end workflow: train -> calibrate -> conformal / abstention -> evaluate.

The ``Bundle`` is everything a deployed tool needs, saved as one joblib file:
the fitted model, the calibrator, and the held-out probability arrays from
which conformal thresholds and abstention thresholds are (re)computed for any
alpha / risk target at prediction time.

Data flow and why each split exists::

    train       -> model selection (internal 80/20) and final model fit
    calib_fit   -> fit the probability calibrator (temperature by default)
    calib_conf  -> conformal quantiles + abstention threshold (guarantees)
    test        -> honest evaluation of everything above + drift reference

Security note: joblib/pickle files execute code on load. Only load bundles
you created yourself.
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass, field
from pathlib import Path

import joblib
import numpy as np
import pandas as pd

from . import __version__
from .calibration import CALIBRATORS, calibration_metrics, compare_calibrators
from .conformal import SplitConformal, coverage_upper_bound, evaluate_sets
from .data import ColumnTypes, build_preprocessor, make_split
from .drift import (
    bbse_confusion, bbse_estimate, classifier_two_sample_test, feature_drift, score_drift, summarize,
)
from .models import MODEL_NAMES, compare_models, full_proba, make_model
from .robustness import certified_accuracy_curve, certified_radius, empirical_stability, linear_weights
from .selective import aurc, choose_threshold, risk_coverage_curve  # noqa: F401 (re-exported)

BUNDLE_FILE = "bundle.joblib"
MAX_REF_ROWS = 3000  # cap on the stored drift reference sample


@dataclass
class Bundle:
    model: object
    model_name: str
    classes: np.ndarray
    types: ColumnTypes
    target: str
    calibrator: object
    linear_model: object | None
    comparison: list
    split_sizes: dict
    seed: int
    # Held-out arrays (labels are integer class indices).
    y_calfit: np.ndarray
    P_calfit_raw: np.ndarray
    y_conf: np.ndarray
    P_conf: np.ndarray
    y_test: np.ndarray
    P_test_raw: np.ndarray
    P_test: np.ndarray
    X_ref: pd.DataFrame
    y_ref: np.ndarray
    p_train: np.ndarray
    vocab: set = field(default_factory=set)
    col_std: dict = field(default_factory=dict)
    train_seconds: float = 0.0
    version: str = __version__

    @property
    def n_classes(self) -> int:
        return len(self.classes)


def _limit_threads():
    """Cap native thread pools (OpenMP in HGB, BLAS) at one thread.

    Applied even if numpy was imported before clfkit (so the environment
    variables set in ``clfkit/__init__`` came too late).
    """
    try:
        from threadpoolctl import threadpool_limits

        return threadpool_limits(1)
    except Exception:  # threadpoolctl ships with sklearn; this is just defensive
        import contextlib

        return contextlib.nullcontext()


def train(
    X: pd.DataFrame,
    y_str: np.ndarray,
    types: ColumnTypes,
    target: str = "label",
    model: str = "auto",
    calibration: str = "temperature",
    test_size: float = 0.2,
    calib_size: float = 0.2,
    seed: int = 0,
    compare: bool = True,
    verbose: bool = True,
) -> Bundle:
    """Fit everything and return a ``Bundle``."""
    t0 = time.perf_counter()
    log = print if verbose else (lambda *a, **k: None)
    classes, y = np.unique(np.asarray(y_str).astype(str), return_inverse=True)
    K = len(classes)
    split = make_split(y, test_size, calib_size, seed)
    log(f"[clfkit] {len(y)} rows, {K} classes, features: "
        f"{len(types.numeric)} numeric / {len(types.categorical)} categorical / {len(types.text)} text")
    log(f"[clfkit] split sizes {split.sizes()}")
    pre = build_preprocessor(types)
    sparse = bool(types.text)
    Xtr, ytr = X.iloc[split.train], y[split.train]

    with _limit_threads():
        comparison = []
        if compare or model == "auto":
            comparison = compare_models(pre, Xtr, ytr, sparse, seed=seed)
            for r in comparison:
                log(f"[clfkit]   {r['model']:7s} val_acc={r['val_accuracy']:.3f} "
                    f"val_logloss={r['val_log_loss']:.3f} val_ece={r['val_ece']:.3f} fit={r['fit_seconds']:.1f}s")
        name = comparison[0]["model"] if model == "auto" else model
        if name not in MODEL_NAMES:
            raise ValueError(f"unknown model {name!r}")
        log(f"[clfkit] final model: {name} (refit on all training rows)")
        final = make_model(name, pre, sparse, seed).fit(Xtr, ytr)
        linear = final if name == "logreg" else make_model("logreg", pre, sparse, seed).fit(Xtr, ytr)

        def proba(idx):
            return full_proba(final, X.iloc[idx], K)

        P_calfit_raw = proba(split.calib_fit)
        cal = CALIBRATORS[calibration]().fit(P_calfit_raw, y[split.calib_fit])
        P_conf = cal.transform(proba(split.calib_conf))
        P_test_raw = proba(split.test)
        P_test = cal.transform(P_test_raw)

    ref_idx = split.test[:MAX_REF_ROWS]
    vocab = set()
    for c in types.text:
        for doc in Xtr[c].astype(str).str.lower().str.split():
            vocab.update(doc)
    col_std = {c: float(Xtr[c].std()) or 1.0 for c in types.numeric}
    b = Bundle(
        model=final, model_name=name, classes=classes, types=types, target=target,
        calibrator=cal, linear_model=linear, comparison=comparison,
        split_sizes=split.sizes(), seed=seed,
        y_calfit=y[split.calib_fit], P_calfit_raw=P_calfit_raw,
        y_conf=y[split.calib_conf], P_conf=P_conf,
        y_test=y[split.test], P_test_raw=P_test_raw, P_test=P_test,
        X_ref=X.iloc[ref_idx].reset_index(drop=True), y_ref=y[ref_idx],
        p_train=np.bincount(ytr, minlength=K) / len(ytr),
        vocab=vocab, col_std=col_std, train_seconds=time.perf_counter() - t0,
    )
    return b


def save(bundle: Bundle, out_dir: str | Path, summary: dict | None = None) -> Path:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    joblib.dump(bundle, out / BUNDLE_FILE)
    if summary is not None:
        (out / "summary.json").write_text(json.dumps(summary, indent=2, default=_json_default))
    return out


def load(model_dir: str | Path) -> Bundle:
    return joblib.load(Path(model_dir) / BUNDLE_FILE)


def _json_default(o):
    if isinstance(o, (np.integer,)):
        return int(o)
    if isinstance(o, (np.floating,)):
        return float(o)
    if isinstance(o, np.ndarray):
        return o.tolist()
    if isinstance(o, set):
        return sorted(o)
    return str(o)


# ------------------------------------------------------------ evaluation ---
def evaluate(
    b: Bundle, alpha: float = 0.1, target_risk: float = 0.05, delta: float = 0.05,
    robustness: bool = True,
) -> dict:
    """All held-out diagnostics. Numbers are on the TEST split unless noted."""
    K = b.n_classes
    res: dict = {"model": b.model_name, "classes": b.classes.tolist(), "split_sizes": b.split_sizes,
                 "comparison": b.comparison, "alpha": alpha}

    # Calibration: every method fitted on calib_fit, evaluated on test.
    res["calibration"] = compare_calibrators(b.P_calfit_raw, b.y_calfit, b.P_test_raw, b.y_test)
    res["calibration_used"] = b.calibrator.name
    res["test_metrics"] = calibration_metrics(b.P_test, b.y_test)

    # Conformal: thresholds from calib_conf, coverage measured on test.
    conf_rows = {}
    for score in ("lac", "aps"):
        for mondrian in (False, True):
            cp = SplitConformal(alpha, score, mondrian, seed=b.seed).fit(b.P_conf, b.y_conf)
            ev = evaluate_sets(cp.predict_sets(b.P_test), b.y_test)
            ev["qhat"] = cp.qhat_.tolist()
            conf_rows[f"{score}{'_mondrian' if mondrian else ''}"] = ev
    res["conformal"] = conf_rows
    res["conformal_n_cal"] = int(len(b.y_conf))
    res["conformal_upper"] = coverage_upper_bound(len(b.y_conf), alpha)

    # Selective classification: threshold certified on calib_conf, checked on test.
    conf_c = b.P_conf.max(1); correct_c = b.P_conf.argmax(1) == b.y_conf
    sel = choose_threshold(conf_c, correct_c, target_risk, delta)
    conf_t = b.P_test.max(1); correct_t = b.P_test.argmax(1) == b.y_test
    if sel["threshold"] is not None:
        acc = conf_t >= sel["threshold"]
        sel["test_coverage"] = float(acc.mean())
        sel["test_selective_risk"] = float(1 - correct_t[acc].mean()) if acc.any() else float("nan")
    cov, risk, _ = risk_coverage_curve(conf_t, correct_t)
    sel["test_curve"] = {"coverage": cov.tolist(), "risk": risk.tolist()}
    sel["test_aurc"] = aurc(conf_t, correct_t)
    sel.pop("table")
    res["selective"] = sel

    # Robustness (linear certificate + empirical stability of the final model).
    if robustness and b.linear_model is not None:
        res.update(_robustness(b))
    return res


def _robustness(b: Bundle) -> dict:
    res = {}
    with _limit_threads():
        lin = b.linear_model
        Z = lin[:-1].transform(b.X_ref)
        W, bias = linear_weights(lin[-1])
        pred, rad = certified_radius(Z, W, bias)
        rmax = float(np.quantile(rad[np.isfinite(rad)], 0.95)) if np.isfinite(rad).any() else 1.0
        radii = np.linspace(0, max(rmax, 1e-3), 40)
        res["certified"] = {
            "model": "logreg",
            "space": "preprocessed features (standardised numerics / one-hot / TF-IDF)",
            "clean_accuracy": float(np.mean(pred == b.y_ref)),
            "median_radius": float(np.median(rad)),
            "radii": radii.tolist(),
            "certified_accuracy": certified_accuracy_curve(pred, rad, b.y_ref, radii).tolist(),
        }
        Xs = b.X_ref.iloc[:500]
        res["stability"] = {
            "model": b.model_name,
            "rows": empirical_stability(b.model, Xs, b.types.numeric, b.types.text,
                                        n_repeats=3, seed=b.seed, col_std=b.col_std),
        }
    return res


# ------------------------------------------------------------ prediction ---
def predict(
    b: Bundle, X: pd.DataFrame, alpha: float = 0.1, score: str = "aps", mondrian: bool = False,
    target_risk: float = 0.05, delta: float = 0.05,
) -> tuple[pd.DataFrame, dict]:
    """Label + conformal set + abstain flag for each row of ``X``."""
    K = b.n_classes
    with _limit_threads():
        P = b.calibrator.transform(full_proba(b.model, X[b.types.all()], K))
    cp = SplitConformal(alpha, score, mondrian, seed=b.seed).fit(b.P_conf, b.y_conf)
    sets = cp.predict_sets(P)
    sel = choose_threshold(b.P_conf.max(1), b.P_conf.argmax(1) == b.y_conf, target_risk, delta)
    lam = sel["threshold"]
    conf = P.max(1)
    abstain = np.ones(len(X), bool) if lam is None else conf < lam
    out = pd.DataFrame({
        "pred_label": b.classes[P.argmax(1)],
        "confidence": np.round(conf, 4),
        "prediction_set": ["|".join(b.classes[m]) for m in sets],
        "set_size": sets.sum(1),
        "abstain": abstain,
    })
    for k, c in enumerate(b.classes):
        out[f"p_{c}"] = np.round(P[:, k], 4)
    info = {"alpha": alpha, "score": score, "mondrian": mondrian, "abstain_threshold": lam,
            "target_risk": target_risk, "delta": delta, "n_cal": int(len(b.y_conf))}
    return out, info


# ----------------------------------------------------------------- drift ---
def drift(b: Bundle, X_new: pd.DataFrame, level: float = 0.01, seed: int = 0) -> dict:
    """Run every drift check of a new unlabelled batch against the stored reference."""
    K = b.n_classes
    Xn = X_new[b.types.all()]
    with _limit_threads():
        P_ref = b.calibrator.transform(full_proba(b.model, b.X_ref, K))
        P_new = b.calibrator.transform(full_proba(b.model, Xn, K))
        f_rows = feature_drift(b.X_ref, Xn, b.types, b.vocab)
        s_res = score_drift(P_ref.max(1), P_new.max(1))
        pre = b.linear_model[:-1]  # fitted preprocessor (same columns as the model)
        c2 = classifier_two_sample_test(pre.transform(b.X_ref), pre.transform(Xn), seed=seed)
    C = bbse_confusion(b.y_ref, P_ref.argmax(1), K)
    bb = bbse_estimate(C, np.bincount(b.y_ref, minlength=K) / len(b.y_ref), P_new.argmax(1))
    return {
        "n_ref": int(len(b.X_ref)), "n_new": int(len(Xn)),
        "summary": summarize(f_rows, s_res, c2, level),
        "features": f_rows, "scores": s_res, "c2st": c2,
        "label_shift": {**bb, "classes": b.classes.tolist()},
    }