from __future__ import annotations

from collections.abc import Mapping, Sequence

import torch
from torch import Tensor, nn

from .encoding import TransformerTextEncoder, serialize_state
from .experts import CandidateScorer, NoulHead, RawExpert
from .types import Question, TaskOutput


class ConceptMemoryModel(nn.Module):
    """Raw-only typed judgment model; state representation is shared per request."""

    def __init__(self, encoder: TransformerTextEncoder, decision_size: int = 256):
        super().__init__()
        self.encoder = encoder
        self.raw_expert = RawExpert(encoder.hidden_size, decision_size)
        self.criterion_projection = nn.Linear(encoder.hidden_size, decision_size)
        self.candidate_scorer = CandidateScorer(decision_size)
        self.noul_head = NoulHead(decision_size)
        self.decision_size = decision_size

    @classmethod
    def from_pretrained(
        cls, name: str = "distilbert/distilbert-base-uncased", decision_size: int = 256,
        max_length: int = 512,
    ) -> ConceptMemoryModel:
        return cls(TransformerTextEncoder.from_pretrained(name, max_length), decision_size)

    def encode_state(self, state: Mapping[str, object] | str) -> Tensor:
        return self.encoder([serialize_state(state)])[0]

    def forward_question(self, state_embedding: Tensor, question: Question) -> TaskOutput:
        question_embedding = self.encoder([question.instructions])[0]
        decision = self.raw_expert(state_embedding, question_embedding)
        if question.type in ("choice", "score"):
            criteria = question.normalized_criteria()
            descriptions = [criterion.description for criterion in criteria]
            criteria_embeddings = self.criterion_projection(self.encoder(descriptions))
            logits = self.candidate_scorer(decision, criteria_embeddings)
            probabilities = logits.softmax(dim=-1)
            prediction = logits.argmax(dim=-1)
            expected_score = None
            if question.type == "score":
                positions = torch.arange(logits.numel(), device=logits.device, dtype=logits.dtype)
                expected_score = (probabilities * positions).sum()
            return TaskOutput(question.id, question.type, logits, probabilities, prediction, expected_score)
        if question.type == "noul":
            logit = self.noul_head(decision)
            probability = logit.sigmoid()
            return TaskOutput(question.id, question.type, logit, probability, probability >= 0.5)
        raise ValueError(f"Unsupported question type: {question.type!r}")

    def forward(
        self, state: Mapping[str, object] | str, questions: Sequence[Question]
    ) -> list[TaskOutput]:
        state_embedding = self.encode_state(state)
        return [self.forward_question(state_embedding, question) for question in questions]

    @torch.inference_mode()
    def predict(
        self, state: Mapping[str, object] | str, questions: Sequence[Question]
    ) -> list[TaskOutput]:
        self.eval()
        return self(state, questions)