"""Drift detection for a deployed classifier.

Three complementary checks on a new, *unlabelled* batch versus a reference
sample drawn from the training distribution:

1. Per-feature tests. Numeric: two-sample Kolmogorov-Smirnov. Categorical:
   chi-square test of homogeneity on the category counts. Text: KS on the
   document length (in tokens) and on the out-of-vocabulary rate. We run many
   tests, so p-values are adjusted with Holm's step-down method (controls the
   family-wise error rate: P(any false alarm) <= level when nothing drifted).
2. Score drift. KS test on the model's confidence (max probability). Cheap,
   model-aware, and often the first thing to move.
3. Classifier two-sample test (C2ST; Lopez-Paz & Oquab, ICLR 2017). Train a
   classifier to tell reference rows from new rows on half the data; under
   H0 (same distribution) its held-out accuracy is Binomial(n, 1/2), so a
   one-sided binomial test gives an exact p-value. Catches multivariate drift
   that per-feature tests miss.

Label shift. If p(y) changes but p(x | y) does not, the model's predictions
on the new batch let us estimate the new class proportions with Black Box
Shift Estimation (BBSE; Lipton, Wang & Smola, ICML 2018): with the confusion
matrix C[i, j] = P(yhat = i, y = j) estimated on labelled reference data,
and mu[i] = P_new(yhat = i) observed on the new batch, solve C w = mu for the
importance weights w_j = q(y=j)/p(y=j); then q = w * p. It is only valid under
the label-shift assumption and needs an invertible (well-conditioned) C; we
clip negative weights and renormalise, and report the condition number.
"""

from __future__ import annotations

import numpy as np
import pandas as pd
from scipy.stats import binomtest, chi2_contingency, ks_2samp
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split

from .data import ColumnTypes


def holm_adjust(pvals) -> np.ndarray:
    """Holm-Bonferroni adjusted p-values (FWER control)."""
    p = np.asarray(pvals, dtype=float)
    m = len(p)
    if m == 0:
        return p
    order = np.argsort(p)
    adj = np.empty(m)
    running = 0.0
    for rank, i in enumerate(order):
        running = max(running, (m - rank) * p[i])
        adj[i] = min(1.0, running)
    return adj


def _tokens(s: pd.Series):
    return s.fillna("").astype(str).str.lower().str.split()


def feature_drift(
    ref: pd.DataFrame, new: pd.DataFrame, types: ColumnTypes, vocab: set | None = None
) -> list[dict]:
    """One row per test with raw and Holm-adjusted p-values, sorted by p-value."""
    rows = []
    for col in types.numeric:
        a, b = ref[col].dropna().to_numpy(), new[col].dropna().to_numpy()
        if len(a) and len(b):
            res = ks_2samp(a, b)
            rows.append({"feature": col, "test": "KS", "statistic": float(res.statistic), "p": float(res.pvalue)})
    for col in types.categorical:
        a, b = ref[col].astype(str), new[col].astype(str)
        cats = sorted(set(a) | set(b))
        table = np.array(
            [[(a == c).sum() for c in cats], [(b == c).sum() for c in cats]], dtype=float
        )
        table = table[:, table.sum(0) > 0]
        if table.shape[1] < 2:
            continue
        chi2, p, _, _ = chi2_contingency(table)
        rows.append({"feature": col, "test": "chi2", "statistic": float(chi2), "p": float(p)})
    for col in types.text:
        ta, tb = _tokens(ref[col]), _tokens(new[col])
        res = ks_2samp(ta.str.len(), tb.str.len())
        rows.append({"feature": f"{col}[length]", "test": "KS", "statistic": float(res.statistic), "p": float(res.pvalue)})
        if vocab:
            def oov(toks):
                return np.array([np.mean([t not in vocab for t in d]) if d else 0.0 for d in toks])
            res = ks_2samp(oov(ta), oov(tb))
            rows.append({"feature": f"{col}[oov_rate]", "test": "KS", "statistic": float(res.statistic), "p": float(res.pvalue)})
    adj = holm_adjust([r["p"] for r in rows])
    for r, a in zip(rows, adj):
        r["p_holm"] = float(a)
    rows.sort(key=lambda r: r["p"])
    return rows


def score_drift(conf_ref: np.ndarray, conf_new: np.ndarray) -> dict:
    """KS test on the model's top-class confidence."""
    res = ks_2samp(conf_ref, conf_new)
    return {
        "statistic": float(res.statistic),
        "p": float(res.pvalue),
        "mean_conf_ref": float(np.mean(conf_ref)),
        "mean_conf_new": float(np.mean(conf_new)),
    }


def classifier_two_sample_test(Z_ref, Z_new, seed: int = 0, max_per_side: int = 2000) -> dict:
    """C2ST with logistic regression on the model's feature matrix.

    Both samples are subsampled to the same size so chance accuracy is 1/2.
    """
    rng = np.random.default_rng(seed)
    n = min(Z_ref.shape[0], Z_new.shape[0], max_per_side)
    ia = rng.choice(Z_ref.shape[0], n, replace=False)
    ib = rng.choice(Z_new.shape[0], n, replace=False)
    import scipy.sparse as sp

    stack = sp.vstack if sp.issparse(Z_ref) else np.vstack
    Z = stack([Z_ref[ia], Z_new[ib]])
    lab = np.r_[np.zeros(n), np.ones(n)]
    tr, te = train_test_split(np.arange(2 * n), test_size=0.5, random_state=seed, stratify=lab)
    clf = LogisticRegression(C=1.0, max_iter=2000).fit(Z[tr], lab[tr])
    correct = int((clf.predict(Z[te]) == lab[te]).sum())
    res = binomtest(correct, len(te), 0.5, alternative="greater")
    return {"accuracy": correct / len(te), "n_test": int(len(te)), "p": float(res.pvalue)}


def bbse_confusion(y_true: np.ndarray, y_pred: np.ndarray, K: int) -> np.ndarray:
    """Joint confusion matrix C[i, j] = P(yhat = i, y = j) on labelled reference data."""
    C = np.zeros((K, K))
    np.add.at(C, (y_pred, y_true), 1.0)
    return C / len(y_true)


def bbse_estimate(C: np.ndarray, p_ref: np.ndarray, y_pred_new: np.ndarray) -> dict:
    """Estimate new class proportions under label shift (hard-prediction BBSE)."""
    K = C.shape[0]
    mu = np.bincount(y_pred_new, minlength=K) / len(y_pred_new)
    w, *_ = np.linalg.lstsq(C, mu, rcond=None)
    w = np.clip(w, 0.0, None)
    q = w * p_ref
    q = q / q.sum() if q.sum() > 0 else np.full(K, 1.0 / K)
    return {
        "weights": w.tolist(),
        "p_ref": p_ref.tolist(),
        "q_est": q.tolist(),
        "pred_dist_new": mu.tolist(),
        "condition_number": float(np.linalg.cond(C)),
        "tv_distance": float(0.5 * np.abs(q - p_ref).sum()),
    }


def summarize(feature_rows, score_res, c2st_res, level: float = 0.01) -> dict:
    """Combine the checks into one alarm flag (Bonferroni over the 3 families)."""
    lvl = level / 3
    n_feat = sum(r["p_holm"] < lvl for r in feature_rows)
    alarms = {
        "features": n_feat > 0,
        "scores": score_res["p"] < lvl,
        "c2st": c2st_res["p"] < lvl,
    }
    return {"drift": any(alarms.values()), "alarms": alarms, "n_drifted_features": int(n_feat), "level": level}