"""
Debatra — Worker 3: Stance Tracker Inference
=============================================
Base model:  cross-encoder/nli-deberta-v3-small
Adapter:     LoRA r=16 trained on ibm/claim_stance + ibm/argq_30k + climate_fever
Task:        3-class sequence classification
Labels:      FAVOR=0  AGAINST=1  NONE=2
Test F1:     0.848

Input:  (text: str, topic: str)
        Internally formatted as: "[TOPIC] {topic} [ARGUMENT] {text}"
Output: {
  "label":      0,
  "stance":     "FAVOR",
  "confidence": 0.91,
  "score":      9.0,         # 0-10 (high = clear committed stance)
  "uncertain":  False,
}
"""

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel


LABEL2ID  = {"FAVOR":0, "AGAINST":1, "NONE":2}
ID2LABEL  = {v: k for k, v in LABEL2ID.items()}
BASE_MODEL = "cross-encoder/nli-deberta-v3-small"
MAX_LENGTH = 256


class StanceTrackerWorker:
    def __init__(self, model_path: str, confidence_threshold: float = 0.60):
        self.model_path           = model_path
        self.confidence_threshold = confidence_threshold
        self._loaded              = False
        self._load()

    def _load(self):
        try:
            try:
                self.tokenizer = AutoTokenizer.from_pretrained(
                    self.model_path, use_fast=True
                )
            except Exception:
                # Some LoRA export folders omit config.json; fall back to base tokenizer.
                self.tokenizer = AutoTokenizer.from_pretrained(
                    BASE_MODEL, use_fast=True
                )
            base = AutoModelForSequenceClassification.from_pretrained(
                BASE_MODEL,
                num_labels=len(LABEL2ID),
                ignore_mismatched_sizes=True,
            )
            self.model = PeftModel.from_pretrained(base, self.model_path)
            self.model.eval()
            if torch.cuda.is_available():
                self.model = self.model.cuda()
            self._loaded = True
        except Exception as e:
            raise RuntimeError(f"Stance Tracker failed to load from {self.model_path}: {e}")

    def status(self) -> str:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        return f"loaded ({device})" if self._loaded else "not loaded"

    def predict(self, text: str, topic: str = "") -> dict:
        """
        Format input as "[TOPIC] {topic} [ARGUMENT] {text}"
        matching the training format from data_prep_v3.py.
        """
        if topic:
            formatted = f"[TOPIC] {topic} [ARGUMENT] {text}"
        else:
            formatted = text

        enc = self.tokenizer(
            formatted,
            truncation=True,
            max_length=MAX_LENGTH,
            return_tensors="pt",
        )
        device = next(self.model.parameters()).device
        enc    = {k: v.to(device) for k, v in enc.items()}

        with torch.no_grad():
            outputs    = self.model(**enc)
            probs      = torch.softmax(outputs.logits[0], dim=-1)
            label_id   = torch.argmax(probs).item()
            confidence = probs[label_id].item()

        stance    = ID2LABEL[label_id]
        uncertain = confidence < self.confidence_threshold

        # Score: clear FAVOR or AGAINST = high, NONE = mid, uncertain = low
        if stance == "NONE":
            score = 4.0
        else:
            score = round(confidence * 10.0, 2)  # max 10 for 100% conf

        if uncertain:
            score = round(score * 0.6, 2)

        return {
            "label":      label_id,
            "stance":     stance if not uncertain else "NONE",
            "confidence": round(confidence, 3),
            "score":      min(score, 10.0),
            "uncertain":  uncertain,
        }