\"\"\"
inference.py  — Deep Dive v2: Self-Contained Inference Module
Model       : itsLu/mentalbert-v5-deep-dive-v2
Task        : Binary Depression vs Suicidal re-ranker (second-tier specialist)

⚠️  WARNING: This model MUST NOT be used standalone.
    It is invoked by Quick Vibe (itsLu/mentalbert-v5-source-aware) only when:
      (a) Quick Vibe's top1 ∈ {Depression, Suicidal} AND margin < 0.20, OR
      (b) Quick Vibe abstains.
    It has NO coverage for: Normal, Anxiety, Stress, Bipolar,
    Personality Disorder, Directed Aggression.

Usage:
    from inference import DeepDiveV2
    clf = DeepDiveV2.from_hub()
    result = clf.predict("I don't want to be here anymore.")
    # {'label': 'Suicidal', 'p_suicidal': 0.91, 'crisis_evidence_found': True,
    #  'crisis_tokens_matched': ['want to die']}
\"\"\"

import json, torch, torch.nn as nn
from typing import Any, Dict, List
from transformers    import BertModel, BertTokenizerFast
from huggingface_hub import hf_hub_download

HF_REPO  = "itsLu/mentalbert-v5-deep-dive-v2"
CLASSES  = ["Depression", "Suicidal"]


class CrisisEvidenceMentalBERT(nn.Module):
    \"\"\"BertModel + crisis-evidence token max-pooling + binary head.\"\"\"
    def __init__(self, model_name: str, hidden_size: int = 768):
        super().__init__()
        self.encoder    = BertModel.from_pretrained(model_name)
        self.dropout    = nn.Dropout(0.2)
        self.classifier = nn.Sequential(
            nn.Linear(hidden_size * 2, 256),
            nn.GELU(),
            nn.Dropout(0.2),
            nn.Linear(256, 2),
        )

    def forward(self, input_ids, attention_mask, crisis_mask):
        h_seq        = self.encoder(input_ids=input_ids,
                                    attention_mask=attention_mask).last_hidden_state
        h_cls        = h_seq[:, 0, :]
        masked       = h_seq.masked_fill(~crisis_mask.unsqueeze(-1), float('-inf'))
        any_crisis   = crisis_mask.any(dim=1)
        h_crisis_raw = masked.max(dim=1).values
        h_crisis     = torch.where(any_crisis.unsqueeze(-1), h_crisis_raw, h_cls)
        return self.classifier(self.dropout(torch.cat([h_cls, h_crisis], dim=-1)))


class DeepDiveV2:
    \"\"\"High-level inference wrapper — load from Hub, call .predict(text).\"\"\"

    def __init__(self, model: CrisisEvidenceMentalBERT, tokenizer: BertTokenizerFast,
                 crisis_keywords: List[str], threshold: float, max_len: int = 256):
        self.model           = model.eval()
        self.tokenizer       = tokenizer
        self.crisis_keywords = [k.lower() for k in crisis_keywords]
        self.threshold       = threshold
        self.max_len         = max_len
        self.device          = next(model.parameters()).device

    @classmethod
    def from_hub(cls, repo_id: str = HF_REPO, device: str = "cpu") -> "DeepDiveV2":
        tokenizer = BertTokenizerFast.from_pretrained(repo_id)
        cfg_path  = hf_hub_download(repo_id=repo_id, filename="inference_config.json")
        kw_path   = hf_hub_download(repo_id=repo_id, filename="crisis_keywords.json")
        clf_path  = hf_hub_download(repo_id=repo_id, filename="classifier.pt")
        with open(cfg_path) as f: cfg = json.load(f)
        with open(kw_path)  as f: kws = json.load(f)["crisis_keywords"]
        model = CrisisEvidenceMentalBERT(model_name=repo_id)
        model.classifier.load_state_dict(torch.load(clf_path, map_location=device))
        model.to(device)
        return cls(model=model, tokenizer=tokenizer, crisis_keywords=kws,
                   threshold=cfg["threshold"], max_len=cfg["max_len"])

    def _crisis_mask(self, text: str, offsets: list) -> "torch.BoolTensor":
        tl = text.lower()
        spans = []
        for kw in self.crisis_keywords:
            s = 0
            while (i := tl.find(kw, s)) >= 0:
                spans.append((i, i + len(kw))); s = i + 1
        mask = torch.zeros(self.max_len, dtype=torch.bool)
        for ti, (s, e) in enumerate(offsets):
            if s == e == 0: continue
            for cs, ce in spans:
                if s < ce and e > cs: mask[ti] = True; break
        return mask

    def predict(self, text: str) -> Dict[str, Any]:
        \"\"\"
        Returns:
          label                 : 'Depression' | 'Suicidal'
          p_suicidal            : float  — P(Suicidal)
          crisis_evidence_found : bool   — any crisis keyword matched
          crisis_tokens_matched : list[str]
        \"\"\"
        enc  = self.tokenizer(text, max_length=self.max_len, truncation=True,
                              padding="max_length", return_offsets_mapping=True,
                              return_tensors="pt")
        ids  = enc["input_ids"].to(self.device)
        amsk = enc["attention_mask"].to(self.device)
        offs = enc["offset_mapping"].squeeze(0).tolist()
        cmsk = self._crisis_mask(text, offs).unsqueeze(0).to(self.device)
        matched = [kw for kw in self.crisis_keywords if kw in text.lower()]
        with torch.no_grad():
            probs = torch.softmax(self.model(ids, amsk, cmsk), dim=-1).squeeze(0).cpu()
        p = float(probs[1])
        return {"label": "Suicidal" if p >= self.threshold else "Depression",
                "p_suicidal": round(p, 6),
                "crisis_evidence_found": bool(cmsk.any().item()),
                "crisis_tokens_matched": matched}

    def predict_batch(self, texts: List[str]) -> List[Dict[str, Any]]:
        return [self.predict(t) for t in texts]