# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Question-local evidence routing for SPEV decision checkpoints."""

from dataclasses import dataclass

import torch
from torch import nn
from torch.nn import functional as F


class EvidenceBlock(nn.Module):
    def __init__(self, width: int, heads: int, ff: int):
        super().__init__()
        self.qnorm = nn.LayerNorm(width)
        self.mnorm = nn.LayerNorm(width)
        self.attn = nn.MultiheadAttention(width, heads, batch_first=True, dropout=0.0)
        self.fnorm = nn.LayerNorm(width)
        self.ff = nn.Sequential(nn.Linear(width, ff), nn.GELU(), nn.Linear(ff, width))

    def forward(self, query, memory, memory_pad):
        memory = self.mnorm(memory)
        value, _ = self.attn(self.qnorm(query), memory, memory, key_padding_mask=memory_pad, need_weights=False)
        query = query + value
        return query + self.ff(self.fnorm(query))


class IsolatedDecisionHead(nn.Module):
    def __init__(self, hidden: int, width: int, heads: int, ff: int, depth: int):
        super().__init__()
        self.norm = nn.LayerNorm(hidden)
        self.memory = nn.Linear(hidden, width, bias=False)
        self.qproj = nn.Linear(hidden, width, bias=False)
        self.oqproj = nn.Linear(hidden, width, bias=False)
        self.globalproj = nn.Linear(hidden, width, bias=False)
        self.context = nn.Linear(hidden, width, bias=False)
        self.lexical = nn.Linear(hidden, width, bias=False)
        self.types = nn.Embedding(3, width)
        self.evidence = nn.ModuleList([EvidenceBlock(width, heads, ff) for _ in range(depth)])
        self.summarynorm = nn.LayerNorm(width)
        self.joint = nn.ModuleList(
            [
                nn.TransformerDecoderLayer(
                    width, heads, ff, dropout=0.0, activation="gelu", batch_first=True, norm_first=True
                )
                for _ in range(depth)
            ]
        )
        self.fnorm = nn.LayerNorm(width)
        self.onorm = nn.LayerNorm(width)
        self.scorer = nn.Sequential(nn.Linear(width * 4, width), nn.GELU(), nn.Linear(width, 1))
        self.prior_scale = nn.Parameter(torch.zeros(()))
        self.joint_scale = nn.Parameter(torch.zeros(()))
        self.gate = nn.Parameter(torch.zeros(()))

    def forward(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
        memory = self.memory(self.norm(batch["memory"]))
        question = self.norm(batch["questions"])
        context = self.norm(batch["options"])
        global_feature = self.norm(batch["global"])
        options = self.context(context) + self.lexical(batch["lexical"]) + self.oqproj(question)
        for block in self.evidence:
            options = block(options, memory, batch["memory_pad"])
        fields = self.qproj(question)
        membership = ~batch["option_pad"][:, None, :]
        weights = torch.einsum("bqd,bcd->bqc", fields, options) / memory.shape[-1] ** 0.5
        weights = weights.masked_fill(~membership, -1e9).softmax(-1) * membership
        summary = torch.einsum("bqc,bcd->bqd", weights, options)
        fields = fields + self.summarynorm(summary) + self.globalproj(global_feature)[:, None, :]
        fields = fields + self.types(batch["types"])
        for block in self.joint:
            fields = block(fields, memory, memory_key_padding_mask=batch["memory_pad"])
        fields = self.fnorm(fields).expand(-1, options.shape[1], -1)
        options = self.onorm(options)
        prior = (
            F.normalize(question + global_feature[:, None, :], dim=-1) * F.normalize(batch["lexical"], dim=-1)
        ).sum(-1) * self.prior_scale.clamp(max=4.605170186).exp()
        joint = F.cosine_similarity(fields, options, dim=-1) * self.joint_scale.clamp(max=4.605170186).exp()
        residual = self.scorer(torch.cat([fields, options, fields * options, (fields - options).abs()], -1)).squeeze(-1)
        return prior + self.gate.sigmoid() * (joint + residual)


@dataclass(frozen=True)
class QuestionFeatures:
    memory: torch.Tensor
    question: torch.Tensor
    options: torch.Tensor
    lexical: torch.Tensor
    global_feature: torch.Tensor
    question_type: int


def score_questions(head: IsolatedDecisionHead, records: list[QuestionFeatures]) -> list[torch.Tensor]:
    """Score all questions in a request batch with one head invocation."""
    if not records:
        raise ValueError("At least one question is required")
    device = head.norm.weight.device
    hidden = head.norm.normalized_shape[0]
    memories = max(len(record.memory) for record in records)
    candidates = max(len(record.options) for record in records)
    batch = {
        name: torch.zeros(shape, device=device, dtype=torch.float32)
        for name, shape in {
            "memory": (len(records), memories, hidden),
            "questions": (len(records), 1, hidden),
            "options": (len(records), candidates, hidden),
            "lexical": (len(records), candidates, hidden),
            "global": (len(records), hidden),
        }.items()
    }
    batch["memory_pad"] = torch.ones(len(records), memories, device=device, dtype=torch.bool)
    batch["option_pad"] = torch.ones(len(records), candidates, device=device, dtype=torch.bool)
    batch["types"] = torch.tensor([[record.question_type] for record in records], device=device)
    for index, record in enumerate(records):
        if not len(record.memory) or not len(record.options) or record.question_type not in (0, 1, 2):
            raise ValueError("Invalid question features")
        if record.lexical.shape != record.options.shape:
            raise ValueError("Criterion embeddings must match candidate features")
        for name, value in (("memory", record.memory), ("options", record.options), ("lexical", record.lexical)):
            batch[name][index, : len(value)] = value.to(device=device, dtype=torch.float32)
        batch["questions"][index, 0] = record.question.to(device=device, dtype=torch.float32)
        batch["global"][index] = record.global_feature.to(device=device, dtype=torch.float32)
        batch["memory_pad"][index, : len(record.memory)] = False
        batch["option_pad"][index, : len(record.options)] = False
    values = head(batch)
    return [values[index, : len(record.options)] for index, record in enumerate(records)]