"""
HuggingFace Inference Endpoints handler for the V6 Hierarchical cascade.

Loads all 5 stages plus Platt calibrators (Stage 1A, Stage 3) and the 3-seed
Stage 2 ensemble. Routes a single string through the cascade.

Supports two operating points selected via `data["mode"]`:
  - "balanced" (default): F1-optimal balance with Sui-miss / Sui-FP penalty.
  - "safety": stricter Suicidal recall (val Sui->Dep <= 70, calibrated for drift).
"""
import os, json, glob, joblib
import numpy as np
import torch
import torch.nn.functional as F
from transformers import (
    AutoTokenizer,
    RobertaForSequenceClassification,
    BertTokenizerFast, BertForSequenceClassification,
    LongformerTokenizerFast, LongformerForSequenceClassification,
)


def _apply_platt(calibrator, p0, p1, eps=1e-7):
    """Apply Platt scaling to a single binary (p0, p1) probability."""
    p0 = float(np.clip(p0, eps, 1 - eps))
    p1 = float(np.clip(p1, eps, 1 - eps))
    logit = np.log(p1) - np.log(p0)
    cal_p1 = float(calibrator.predict_proba(np.array([[logit]]))[:, 1][0])
    return [1.0 - cal_p1, cal_p1]


class EndpointHandler:
    def __init__(self, path=""):
        with open(os.path.join(path, "config.json"), "r") as f:
            self.cfg = json.load(f)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.s2_classes = self.cfg["stage2_class_order"]

        # ── Stage 0: Cardiff RoBERTa ──
        s0_path = os.path.join(path, "stage0")
        self.tok0 = AutoTokenizer.from_pretrained(s0_path)
        self.m0   = RobertaForSequenceClassification.from_pretrained(s0_path).to(self.device).eval()

        # ── Stage 1A: MentalBERT + Platt calibrator ──
        s1a_path = os.path.join(path, "stage1a")
        self.tok1a = BertTokenizerFast.from_pretrained(s1a_path)
        self.m1a   = BertForSequenceClassification.from_pretrained(s1a_path).to(self.device).eval()
        cal_s1a = os.path.join(s1a_path, "platt_calibrator.joblib")
        self.platt_s1a = joblib.load(cal_s1a) if os.path.exists(cal_s1a) else None

        # ── Stage 1B: MentalBERT ──
        s1b_path = os.path.join(path, "stage1b")
        self.tok1b = BertTokenizerFast.from_pretrained(s1b_path)
        self.m1b   = BertForSequenceClassification.from_pretrained(s1b_path).to(self.device).eval()

        # ── Stage 2: 3-seed MentalBERT ensemble ──
        seed_dirs = sorted(glob.glob(os.path.join(path, "stage2", "seed_*")))
        if not seed_dirs:
            # Backwards compat: single stage2 folder
            seed_dirs = [os.path.join(path, "stage2")]
        self.tok2_list = [BertTokenizerFast.from_pretrained(d) for d in seed_dirs]
        self.m2_list   = [BertForSequenceClassification.from_pretrained(d).to(self.device).eval()
                          for d in seed_dirs]
        self.n_s2_models = len(self.m2_list)

        # ── Stage 3: Longformer + Platt calibrator ──
        s3_path = os.path.join(path, "stage3")
        self.tok3 = LongformerTokenizerFast.from_pretrained(s3_path)
        self.m3   = LongformerForSequenceClassification.from_pretrained(s3_path).to(self.device).eval()
        cal_s3 = os.path.join(s3_path, "platt_calibrator.joblib")
        self.platt_s3 = joblib.load(cal_s3) if os.path.exists(cal_s3) else None

        st = self.cfg["stages"]
        self.ml0  = st["stage0"]["max_len"]
        self.ml1a = st["stage1a"]["max_len"]
        self.ml1b = st["stage1b"]["max_len"]
        self.ml2  = st["stage2"]["max_len"]
        self.ml3  = st["stage3"]["max_len"]

        thr = self.cfg["thresholds"]
        self.t0        = float(thr["stage0"])
        self.balanced  = {"stage1a": float(thr["balanced"]["stage1a"]),
                          "stage3":  float(thr["balanced"]["stage3"])}
        self.safety    = {"stage1a": float(thr["safety"]["stage1a"]),
                          "stage3":  float(thr["safety"]["stage3"])}
        self.default_mode = thr.get("default_mode", "balanced")

    @torch.no_grad()
    def _probs_bert(self, m, tok, text, max_len):
        enc = tok(text, max_length=max_len, padding="max_length",
                  truncation=True, return_tensors="pt").to(self.device)
        out = m(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"])
        return F.softmax(out.logits, dim=-1)[0].cpu().tolist()

    @torch.no_grad()
    def _probs_s2_ensemble(self, text):
        """Average softmax probs across all stage-2 seed models."""
        acc = None
        for m, tok in zip(self.m2_list, self.tok2_list):
            enc = tok(text, max_length=self.ml2, padding="max_length",
                      truncation=True, return_tensors="pt").to(self.device)
            out = m(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"])
            p   = F.softmax(out.logits, dim=-1)[0].cpu().numpy()
            acc = p if acc is None else acc + p
        return (acc / self.n_s2_models).tolist()

    @torch.no_grad()
    def _probs_lf(self, m, tok, text, max_len):
        enc = tok(text, max_length=max_len, padding="max_length",
                  truncation=True, return_tensors="pt").to(self.device)
        gmask = torch.zeros_like(enc["attention_mask"])
        gmask[0, 0] = 1
        out = m(input_ids=enc["input_ids"],
                attention_mask=enc["attention_mask"],
                global_attention_mask=gmask)
        return F.softmax(out.logits, dim=-1)[0].cpu().tolist()

    def __call__(self, data):
        if isinstance(data, str):
            text = data; mode = self.default_mode
        else:
            text = data.get("inputs", "")
            if isinstance(text, list):
                text = text[0] if len(text) > 0 else ""
            mode = data.get("mode", self.default_mode)

        if mode not in ("balanced", "safety"):
            mode = self.default_mode
        thr = self.safety if mode == "safety" else self.balanced
        t1a = thr["stage1a"]; t3 = thr["stage3"]

        stage_probs = {}

        # Stage 0: DA gate
        p0 = self._probs_bert(self.m0, self.tok0, text, self.ml0)
        stage_probs["stage0"] = p0
        if p0[1] >= self.t0:
            return {"label": "Directed Aggression", "exit_stage": "stage0",
                    "mode": mode, "stage_probs": stage_probs}

        # Stage 1A: Sui gate (Platt-calibrated)
        p1a_raw = self._probs_bert(self.m1a, self.tok1a, text, self.ml1a)
        p1a = _apply_platt(self.platt_s1a, p1a_raw[0], p1a_raw[1]) if self.platt_s1a else p1a_raw
        stage_probs["stage1a"] = p1a
        stage_probs["stage1a_raw"] = p1a_raw
        if p1a[1] >= t1a:
            return {"label": "Suicidal", "exit_stage": "stage1a",
                    "mode": mode, "stage_probs": stage_probs}

        # Stage 1B: Normal vs Distress (argmax)
        p1b = self._probs_bert(self.m1b, self.tok1b, text, self.ml1b)
        stage_probs["stage1b"] = p1b
        if p1b[0] > p1b[1]:
            return {"label": "Normal", "exit_stage": "stage1b",
                    "mode": mode, "stage_probs": stage_probs}

        # Stage 2: 5-class ensemble argmax
        p2 = self._probs_s2_ensemble(text)
        stage_probs["stage2"] = p2
        s2_idx = int(max(range(len(p2)), key=lambda i: p2[i]))
        s2_label = self.s2_classes[s2_idx]
        if s2_label != "Depression":
            return {"label": s2_label, "exit_stage": "stage2",
                    "mode": mode, "stage_probs": stage_probs}

        # Stage 3: Dep vs Sui (Platt-calibrated)
        p3_raw = self._probs_lf(self.m3, self.tok3, text, self.ml3)
        p3 = _apply_platt(self.platt_s3, p3_raw[0], p3_raw[1]) if self.platt_s3 else p3_raw
        stage_probs["stage3"] = p3
        stage_probs["stage3_raw"] = p3_raw
        final = "Suicidal" if p3[1] >= t3 else "Depression"
        return {"label": final, "exit_stage": "stage3",
                "mode": mode, "stage_probs": stage_probs}