"""Inference helpers for Dialectica."""

ORDERED_LABELS = ["Surface", "Mechanistic", "Critical"]


class CognitiveClassifier:
    """DistilBERT question classifier."""

    def __init__(self, product_config):
        """Load model and tokenizer."""
        import torch
        from transformers import (
            AutoModelForSequenceClassification,
            AutoTokenizer,
        )

        self.cfg = product_config
        self.torch = torch
        self.tokenizer = AutoTokenizer.from_pretrained(self.cfg.classifier_dir)
        self.model = AutoModelForSequenceClassification.from_pretrained(
            self.cfg.classifier_dir
        )
        self.model.eval()
        self.id2label = self.model.config.id2label or {
            i: label for i, label in enumerate(ORDERED_LABELS)
        }

    def classify(self, question):
        """Predict level and confidence."""
        inputs = self.tokenizer(
            question,
            truncation=True,
            max_length=self.cfg.max_question_length,
            return_tensors="pt",
        )
        with self.torch.no_grad():
            logits = self.model(**inputs).logits
            probs = self.torch.softmax(logits, dim=1)[0]
        predicted_id = int(self.torch.argmax(probs))
        return {
            "level": self.id2label[predicted_id],
            "confidence": float(probs[predicted_id]),
        }


class ConceptMatcher:
    """Embedding-based concept matcher."""

    def __init__(self, product_config):
        """Load embedding model."""
        from sentence_transformers import SentenceTransformer

        self.cfg = product_config
        self.model = SentenceTransformer(self.cfg.embed_model)
        self.concepts = []
        self.concept_embeddings = None

    def set_concepts(self, concepts):
        """Cache concept embeddings."""
        self.concepts = concepts
        if concepts:
            self.concept_embeddings = self.model.encode(
                concepts, convert_to_tensor=True
            )
        else:
            self.concept_embeddings = None

    def match(self, question):
        """Return up to two matching concepts."""
        from sentence_transformers import util

        if not self.concepts or self.concept_embeddings is None:
            return []
        query = self.model.encode(question, convert_to_tensor=True)
        scores = util.cos_sim(query, self.concept_embeddings)[0]

        ranked = sorted(
            range(len(self.concepts)),
            key=lambda i: float(scores[i]),
            reverse=True,
        )
        threshold = self.cfg.concept_match_threshold
        # Keep a second concept if it's close to top or strong enough.
        runner_up_margin = 0.05
        runner_up_absolute = 0.45

        top_index = ranked[0]
        top_score = float(scores[top_index])
        if top_score < threshold:
            return []

        matched = [self.concepts[top_index]]
        if len(ranked) > 1:
            second_index = ranked[1]
            second_score = float(scores[second_index])
            close_to_top = top_score - second_score <= runner_up_margin
            strong_alone = second_score >= runner_up_absolute
            if second_score >= threshold and (close_to_top or strong_alone):
                matched.append(self.concepts[second_index])
        return matched