"""
Text classifier for category, root_cause, and subcategory.

Architecture (v2.1 — "tri-vote"):
  P = (2·P_tfidf + P_emb_lr + P_emb_knn) / 4
    P_tfidf   — TF-IDF (word+char) → soft vote of calibrated LR + calibrated
                LinearSVC + ComplementNB, trained on serve-shape text +
                investigation-augmented text
    P_emb_lr  — calibrated LogisticRegression on multilingual sentence
                embeddings of the raw report text (models/embedder.py)
    P_emb_knn — cosine kNN (k=5, distance-weighted) on the same embeddings —
                template-heavy operational text suits retrieval
  The reported confidence is Platt-recalibrated on out-of-fold predictions so
  "confidence 0.8" empirically means ≈80% correct. If the embedding model is
  unavailable the classifier degrades to TF-IDF-only automatically.

Honesty rules (chosen after measuring the deployed train/serve skew):
  1. Train/serve parity — evaluated and served on exactly the input shape the
     API receives. Post-investigation columns (Root Caused, Action Taken) are
     training augmentation only, never serve features.
  2. Dedup — duplicate (report text, label) pairs are collapsed before
     training/CV so template texts can't leak across folds.
  3. No fabricated answers — below the confidence floor the API returns
     UNCERTAIN with ranked candidates instead of guessing.
"""
from __future__ import annotations

import re
from pathlib import Path
from typing import Optional

import joblib
import numpy as np
import pandas as pd
from sklearn.calibration import CalibratedClassifierCV
from sklearn.ensemble import VotingClassifier
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, f1_score
from sklearn.model_selection import StratifiedKFold
from sklearn.naive_bayes import ComplementNB
from sklearn.neighbors import KNeighborsClassifier
from sklearn.pipeline import Pipeline, FeatureUnion
from sklearn.svm import LinearSVC

from config import DATA_DIR, get_logger
from models import embedder

log = get_logger("classifier")

CONFIDENCE_THRESHOLD = 0.35
HIGH_CONFIDENCE = 0.60
MIN_TRAIN_ROWS = 10
MIN_CLASS_SAMPLES = 3
# Candidate blend weights (tfidf, emb_lr, emb_knn); the best combo per role is
# selected on out-of-fold predictions at train time.
WEIGHT_GRID = [(2, 1, 1), (1, 1, 1), (1, 2, 1), (1, 1, 2), (2, 2, 1)]
DEFAULT_WEIGHTS = (2.0, 1.0, 1.0)

INDONESIAN_STOPWORDS = {
    "yang", "di", "ke", "dan", "atau", "dengan", "untuk", "dari", "ini", "itu",
    "pada", "oleh", "karena", "dalam", "bahwa", "tidak", "ada", "akan", "sudah",
    "bisa", "juga", "tersebut", "dapat", "saat", "telah", "setelah", "ketika",
    "seperti", "lebih", "sangat", "masih", "sesuai", "agar", "sehingga",
    "namun", "tetapi", "maka", "jika", "apabila", "walaupun", "meskipun",
    "sebelum", "antara", "terhadap", "melalui", "berdasarkan", "atas",
    "secara", "hal", "serta", "pihak", "selama", "kepada", "para", "suatu",
    "nya", "kita", "kami", "mereka", "dia", "belum", "baru", "sebuah",
    "demikian", "kemudian", "selanjutnya",
    "a", "an", "the", "is", "are", "was", "were", "be", "been", "being",
    "have", "has", "had", "do", "does", "did", "will", "would", "could",
    "should", "may", "might", "shall", "can", "need", "to", "of", "in",
    "on", "at", "by", "for", "with", "about",
}


def _clean_text(text: str) -> str:
    if not text or str(text).lower() in ("none", "nan", ""):
        return ""
    text = str(text).lower()
    text = re.sub(r"[^\w\s]", " ", text)
    text = re.sub(r"\b\d{3,}\b", " NUM ", text)
    tokens = [t for t in text.split() if t not in INDONESIAN_STOPWORDS and len(t) > 1]
    return " ".join(tokens)


def _ctx_token(prefix: str, value) -> str:
    """Context value → stable feature token, e.g. __air_garuda_indonesia__."""
    if value is None or str(value).strip().lower() in ("", "none", "nan", "no delay"):
        return ""
    return f"__{prefix}_{str(value).strip().lower().replace(' ', '_')}__"


def build_serve_text(
    text: str,
    category: Optional[str] = None,
    airline: Optional[str] = None,
    branch: Optional[str] = None,
    area: Optional[str] = None,
    sheet: Optional[str] = None,
) -> str:
    """
    The one serve-shape input builder — used identically at train and predict
    time. Only fields known when the report is filed may appear here.
    """
    parts = [_clean_text(text)]
    parts.append(_ctx_token("cat", category))
    parts.append(_ctx_token("air", airline))
    parts.append(_ctx_token("br", branch))
    parts.append(_ctx_token("area", area))
    parts.append(_ctx_token("sheet", sheet))
    return " ".join(p for p in parts if p)


# Backwards-compatible alias (older imports)
def _build_combined_text(desc: str, category=None, airline=None, branch=None,
                         area=None, delay_code=None, severity=None,
                         case_classification=None) -> str:
    return build_serve_text(desc, category=category, airline=airline,
                            branch=branch, area=area)


def _tfidf_union() -> FeatureUnion:
    return FeatureUnion([
        ("word", TfidfVectorizer(
            analyzer="word", ngram_range=(1, 2), min_df=1, max_df=0.95,
            max_features=20000, sublinear_tf=True,
            stop_words=list(INDONESIAN_STOPWORDS))),
        ("char", TfidfVectorizer(
            analyzer="char_wb", ngram_range=(2, 5), min_df=1, max_df=0.95,
            max_features=15000, sublinear_tf=True)),
    ])


def _make_pipeline() -> Pipeline:
    """TF-IDF → soft vote of calibrated LR + calibrated SVC + ComplementNB."""
    return Pipeline([
        ("tfidf", _tfidf_union()),
        ("clf", VotingClassifier(
            estimators=[
                ("lr", CalibratedClassifierCV(
                    LogisticRegression(max_iter=3000, class_weight="balanced", C=5.0),
                    method="sigmoid", cv=3)),
                ("svc", CalibratedClassifierCV(
                    LinearSVC(class_weight="balanced", C=1.0),
                    method="sigmoid", cv=3)),
                ("nb", ComplementNB(alpha=0.3)),
            ],
            voting="soft",
        )),
    ])


def _make_emb_lr() -> CalibratedClassifierCV:
    return CalibratedClassifierCV(
        LogisticRegression(max_iter=3000, class_weight="balanced", C=10.0),
        method="sigmoid", cv=3)


def _fit_emb_lr(E, y):
    """Calibrated LR on embeddings; degrades to plain LR when the smallest
    class is too small for the calibration folds."""
    min_class = int(pd.Series(y).value_counts().min())
    if min_class >= 3:
        return _make_emb_lr().fit(E, y)
    if min_class >= 2:
        return CalibratedClassifierCV(
            LogisticRegression(max_iter=3000, class_weight="balanced", C=10.0),
            method="sigmoid", cv=2).fit(E, y)
    return LogisticRegression(max_iter=3000, class_weight="balanced", C=10.0).fit(E, y)


def _make_emb_knn() -> KNeighborsClassifier:
    return KNeighborsClassifier(n_neighbors=5, metric="cosine", weights="distance")


def _model_path(role: str) -> Path:
    return DATA_DIR / f"classifier_{role}.joblib"


def _align(P: np.ndarray, member_classes, target_classes) -> np.ndarray:
    out = np.zeros((P.shape[0], len(target_classes)))
    order = {c: i for i, c in enumerate(target_classes)}
    for j, c in enumerate(member_classes):
        out[:, order[c]] = P[:, j]
    return out


class _Classifier:
    """Tri-vote (TF-IDF ensemble + embedding LR + embedding kNN) with
    Platt-recalibrated confidence and an abstention floor."""

    def __init__(self, role: str):
        self.role = role
        self.pipe: Optional[Pipeline] = None
        self.emb_lr = None
        self.emb_knn = None
        self.platt: Optional[LogisticRegression] = None
        self.uses_embeddings = False
        self.blend_weights: tuple = DEFAULT_WEIGHTS
        # Inputs whose nearest training neighbour is less similar than this
        # are out-of-distribution: abstain instead of hallucinating a label.
        self.ood_sim_floor: Optional[float] = None

    # ── internals ─────────────────────────────────────────────────────────

    def _blend(self, X_serve_texts, E: Optional[np.ndarray]) -> tuple[np.ndarray, list]:
        wt, wl, wk = self.blend_weights
        P = wt * self.pipe.predict_proba(X_serve_texts)
        classes = list(self.pipe.classes_)
        total = wt
        if E is not None and self.emb_lr is not None and self.emb_knn is not None:
            P = P + wl * _align(self.emb_lr.predict_proba(E), self.emb_lr.classes_, classes)
            P = P + wk * _align(self.emb_knn.predict_proba(E), self.emb_knn.classes_, classes)
            total += wl + wk
        return P / total, classes

    def _recalibrate(self, conf: np.ndarray) -> np.ndarray:
        if self.platt is None:
            return conf
        return self.platt.predict_proba(conf.reshape(-1, 1))[:, 1]

    # ── training ──────────────────────────────────────────────────────────

    def fit(self, X_serve: pd.Series, y: pd.Series,
            X_aug: Optional[pd.Series] = None,
            X_raw: Optional[pd.Series] = None) -> dict:
        X_serve = X_serve.fillna("").astype(str).reset_index(drop=True)
        y = y.fillna("").astype(str).reset_index(drop=True)
        X_aug = X_aug.fillna("").astype(str).reset_index(drop=True) if X_aug is not None else None
        X_raw = X_raw.fillna("").astype(str).reset_index(drop=True) if X_raw is not None else None

        counts = y.value_counts()
        valid = counts[counts >= MIN_CLASS_SAMPLES].index
        mask = y.isin(valid).values
        X_serve, y = X_serve[mask].reset_index(drop=True), y[mask].reset_index(drop=True)
        if X_aug is not None:
            X_aug = X_aug[mask].reset_index(drop=True)
        if X_raw is not None:
            X_raw = X_raw[mask].reset_index(drop=True)

        if len(X_serve) < MIN_TRAIN_ROWS or y.nunique() < 2:
            raise ValueError(f"Insufficient data: {len(X_serve)} rows, {y.nunique()} classes")

        # Embeddings for the whole (deduped) training set — None if unavailable
        E_all = embedder.encode(X_raw.tolist()) if X_raw is not None else None
        self.uses_embeddings = E_all is not None

        # ── OOF cross-validation replicating the production recipe ────────
        # Member probabilities are collected per fold so the blend weights can
        # be selected on out-of-fold predictions. Note: metrics are reported
        # for the selected weights on the same OOF — a small (~5-candidate)
        # selection optimism, documented and accepted.
        cv_m: dict = {}
        n_folds = int(min(5, y.value_counts().min()))
        if n_folds >= 2:
            skf = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=42)
            all_classes = np.unique(y)
            n, C = len(y), len(all_classes)
            OP_t, OP_l, OP_k = np.zeros((n, C)), np.zeros((n, C)), np.zeros((n, C))
            for tr, te in skf.split(X_serve, y):
                fold = _Classifier(self.role)
                fold._fit_members(X_serve.iloc[tr], y.iloc[tr],
                                  X_aug.iloc[tr] if X_aug is not None else None,
                                  E_all[tr] if E_all is not None else None)
                Xe = X_serve.iloc[te]
                OP_t[te] = _align(fold.pipe.predict_proba(Xe), fold.pipe.classes_, all_classes)
                if E_all is not None and fold.emb_lr is not None:
                    OP_l[te] = _align(fold.emb_lr.predict_proba(E_all[te]),
                                      fold.emb_lr.classes_, all_classes)
                    OP_k[te] = _align(fold.emb_knn.predict_proba(E_all[te]),
                                      fold.emb_knn.classes_, all_classes)

            grid = WEIGHT_GRID if E_all is not None else [(1.0, 0.0, 0.0)]
            best_w, best_acc, best_pred, best_conf = None, -1.0, None, None
            for w in grid:
                wt, wl, wk = w
                P = (wt * OP_t + wl * OP_l + wk * OP_k) / (wt + wl + wk)
                idx = np.argmax(P, axis=1)
                pred_w = all_classes[idx]
                acc_w = float(accuracy_score(y, pred_w))
                if acc_w > best_acc:
                    best_w, best_acc = w, acc_w
                    best_pred = pred_w
                    best_conf = P[np.arange(n), idx]
            self.blend_weights = tuple(float(v) for v in best_w)
            pred, conf = best_pred, best_conf

            # Platt recalibration: raw blended confidence → P(top label correct)
            correct = (pred == y.values).astype(int)
            if 0 < correct.sum() < len(correct):
                self.platt = LogisticRegression(max_iter=1000)
                self.platt.fit(conf.reshape(-1, 1), correct)
            conf_cal = self._recalibrate(conf)

            hi = conf_cal >= HIGH_CONFIDENCE
            sel_acc = float(accuracy_score(y.values[hi], pred[hi])) if hi.sum() else 0.0
            cv_m = {
                "cv_accuracy":     round(float(accuracy_score(y, pred)), 4),
                "cv_f1_weighted":  round(float(f1_score(y, pred, average="weighted", zero_division=0)), 4),
                "cv_folds":        n_folds,
                "blend_weights":   list(self.blend_weights),
                "hiconf_coverage": round(float(hi.mean()), 4),
                "hiconf_accuracy": round(sel_acc, 4),
                "effectiveness":   round(float(hi.mean()) * sel_acc, 4),
            }

        # ── final members on all data ──────────────────────────────────────
        self._fit_members(X_serve, y, X_aug, E_all)

        log.info(
            "[%s] cv_acc=%s hiconf_cov=%s hiconf_acc=%s | %d unique rows, %d classes, embeddings=%s",
            self.role, cv_m.get("cv_accuracy", "?"), cv_m.get("hiconf_coverage", "?"),
            cv_m.get("hiconf_accuracy", "?"), len(X_serve), y.nunique(), self.uses_embeddings,
        )
        return {
            "accuracy":    cv_m.get("cv_accuracy"),
            "f1_weighted": cv_m.get("cv_f1_weighted"),
            "n_classes":   int(y.nunique()),
            "n_samples":   int(len(X_serve)),
            "embeddings":  self.uses_embeddings,
            **cv_m,
        }

    def _fit_members(self, X_serve, y, X_aug, E):
        if X_aug is not None:
            X_fit = pd.concat([X_serve, X_aug], ignore_index=True)
            y_fit = pd.concat([y, y], ignore_index=True)
        else:
            X_fit, y_fit = X_serve, y
        self.pipe = _make_pipeline()
        self.pipe.fit(X_fit, y_fit)
        if E is not None:
            self.emb_lr = _fit_emb_lr(E, y)
            self.emb_knn = _make_emb_knn()
            self.emb_knn.fit(E, y)
            # Data-driven OOD floor: 1st percentile of each training row's
            # nearest-neighbour similarity (2nd neighbour = excludes self),
            # clipped to a sane band.
            dist, _ = self.emb_knn.kneighbors(E, n_neighbors=min(2, len(E)))
            nn_sim = 1.0 - dist[:, -1]
            self.ood_sim_floor = float(np.clip(np.percentile(nn_sim, 1), 0.30, 0.55))
        else:
            self.emb_lr = self.emb_knn = None
            self.ood_sim_floor = None

    # ── inference ─────────────────────────────────────────────────────────

    def predict(self, serve_text: str, raw_text: Optional[str] = None) -> dict:
        if not serve_text or not serve_text.strip():
            return {"label": None, "confidence": None, "status": "empty_input"}
        if self.pipe is None:
            return {"label": None, "confidence": None, "status": "model_not_trained"}

        E = None
        degraded = False
        if self.uses_embeddings:
            E = embedder.encode([raw_text or serve_text])
            if E is None:
                degraded = True  # trained with embeddings, serving without

        P, classes = self._blend([serve_text], E)
        proba = P[0]
        order = np.argsort(proba)[::-1]
        conf_raw = float(proba[order[0]])
        conf = float(self._recalibrate(np.array([conf_raw]))[0])

        top = [
            {"label": classes[i], "confidence": round(float(proba[i]), 4)}
            for i in order[:3]
        ]
        out: dict = {"confidence": round(conf, 4), "top_candidates": top}
        if degraded:
            out["degraded"] = "embeddings_unavailable"

        # OOD guard: an input unlike anything in training must not get a
        # confident label, no matter what the calibrated blend says.
        if E is not None and self.emb_knn is not None and self.ood_sim_floor is not None:
            dist, _ = self.emb_knn.kneighbors(E, n_neighbors=1)
            nn_sim = float(1.0 - dist[0][0])
            out["nn_similarity"] = round(nn_sim, 3)
            if nn_sim < self.ood_sim_floor:
                out["confidence"] = round(min(conf, CONFIDENCE_THRESHOLD - 0.01), 4)
                return {**out, "label": "UNCERTAIN", "status": "low_confidence",
                        "reason": "input_dissimilar_to_training_data"}

        if conf < CONFIDENCE_THRESHOLD:
            return {**out, "label": "UNCERTAIN", "status": "low_confidence"}
        return {
            **out,
            "label": classes[order[0]],
            "status": "ok" if conf >= HIGH_CONFIDENCE else "medium_confidence",
        }

    def __getstate__(self):
        # embedder singleton is module-level; nothing torch-y lives on self
        return self.__dict__


# ── Public API ────────────────────────────────────────────────────────────────

# In-memory cache so /classify doesn't joblib.load on every request.
_CACHE: dict[str, tuple[float, "_Classifier"]] = {}


def train(X: pd.Series, y: pd.Series, role: str,
          X_aug: Optional[pd.Series] = None,
          X_raw: Optional[pd.Series] = None) -> dict:
    clf = _Classifier(role)
    metrics = clf.fit(X, y, X_aug=X_aug, X_raw=X_raw)
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    joblib.dump(clf, _model_path(role))
    _CACHE.pop(role, None)
    return metrics


def _load(role: str) -> Optional["_Classifier"]:
    path = _model_path(role)
    if not path.exists():
        return None
    mtime = path.stat().st_mtime
    cached = _CACHE.get(role)
    if cached and cached[0] == mtime:
        return cached[1]
    clf: _Classifier = joblib.load(path)
    _CACHE[role] = (mtime, clf)
    return clf


def predict(text: str, role: str,
            category: Optional[str] = None,
            airline: Optional[str] = None,
            branch: Optional[str] = None,
            area: Optional[str] = None,
            sheet: Optional[str] = None,
            category_hint: Optional[str] = None) -> dict:
    clf = _load(role)
    if clf is None:
        return {"label": None, "confidence": None, "status": "model_not_trained"}
    serve_text = build_serve_text(
        text,
        category=category or category_hint,
        airline=airline, branch=branch, area=area, sheet=sheet,
    )
    return clf.predict(serve_text, raw_text=str(text).strip())


def model_exists(role: str) -> bool:
    return _model_path(role).exists()