"""
Debatra — Worker 2: Fallacy Detector Inference
================================================
Base model:  microsoft/deberta-v3-base
Adapter:     LoRA r=32 trained on tasksource/logical-fallacy + MAFALDA
Task:        6-class sequence classification
Labels:      ad_hominem=0  appeal_to_authority=1  false_dichotomy=2
             strawman=3  hasty_generalization=4  no_fallacy=5
Test F1:     0.609

Input:  text string (sentence or short passage)
Output: {
  "label":      0,
  "label_name": "ad_hominem",
  "confidence": 0.87,
  "is_fallacy": True,
  "score":      2.5,        # 1-10 (low = fallacious)
  "detail":     "Detected: ad_hominem (87% confidence)",
  "uncertain":  False,
}
"""

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


LABEL2ID = {
    "ad_hominem":0, "appeal_to_authority":1,
    "false_dichotomy":2, "strawman":3,
    "hasty_generalization":4, "no_fallacy":5,
}
ID2LABEL  = {v: k for k, v in LABEL2ID.items()}

# Fallacy score: fallacious labels get low score, no_fallacy gets high score
LABEL_TO_SCORE = {
    "ad_hominem":0,          "appeal_to_authority":1,
    "false_dichotomy":2,     "strawman":2,
    "hasty_generalization":3, "no_fallacy":5,
}

BASE_MODEL = "microsoft/deberta-v3-base"
MAX_LENGTH = 256


class FallacyDetectorWorker:
    def __init__(self, model_path: str, confidence_threshold: float = 0.55):
        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"Fallacy Detector 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) -> dict:
        enc = self.tokenizer(
            text,
            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()

        label_name = ID2LABEL[label_id]
        is_fallacy = label_name != "no_fallacy"
        uncertain  = confidence < self.confidence_threshold

        # Score: base from label severity, scaled by confidence
        base_score = LABEL_TO_SCORE[label_name]
        # Scale to 1-10: no_fallacy → 10, worst fallacy → 1
        # base_score 0=worst → 1, base_score 5=clean → 10
        score_1_10 = round(1 + (base_score / 5.0) * 9.0, 2)
        # Confidence-adjusted: uncertain predictions get middled
        if uncertain:
            score_1_10 = round(score_1_10 * 0.7 + 5.0 * 0.3, 2)

        detail = (
            f"Detected: {label_name.replace('_', ' ')} ({confidence:.0%} confidence)"
            if is_fallacy and not uncertain
            else "No fallacy detected" if not is_fallacy
            else f"Low confidence fallacy signal ({label_name.replace('_', ' ')}, {confidence:.0%})"
        )

        return {
            "label":      label_id,
            "label_name": label_name,
            "confidence": round(confidence, 3),
            "is_fallacy": is_fallacy and not uncertain,
            "score":      score_1_10,
            "detail":     detail,
            "uncertain":  uncertain,
        }