"""SPEV backbone, residual adapters and contextual decision head."""

import math
from dataclasses import dataclass
from typing import ClassVar

import torch
from torch import nn
from transformers import PreTrainedModel, Qwen3OmniMoeThinkerForConditionalGeneration
from transformers.utils import ModelOutput

from .configuration_spev import SPEVConfig


class ResidualAdapter(nn.Module):
    def __init__(self, base, rank, alpha):
        super().__init__()
        self.base = base
        self.scale = alpha / rank
        self.a = nn.Parameter(torch.empty(rank, base.in_features, dtype=torch.float32))
        self.b = nn.Parameter(torch.zeros(base.out_features, rank, dtype=torch.float32))
        nn.init.kaiming_uniform_(self.a, a=math.sqrt(5))

    def forward(self, x):
        residual = nn.functional.linear(nn.functional.linear(x.float(), self.a), self.b)
        return self.base(x) + (residual * self.scale).to(x.dtype)


class OptionHead(nn.Module):
    def __init__(self, hidden_size, width):
        super().__init__()
        self.scorer = nn.Sequential(
            nn.LayerNorm(hidden_size),
            nn.Linear(hidden_size, width),
            nn.GELU(),
            nn.Linear(width, 1),
        )

    def forward(self, features):
        return self.scorer(features.float()).squeeze(-1)


@dataclass
class SPEVOutput(ModelOutput):
    logits: torch.Tensor | None = None
    input_tokens: torch.Tensor | None = None
    temperature: float | None = None


class SPEVModel(PreTrainedModel):
    config_class = SPEVConfig
    base_model_prefix = "backbone"
    main_input_name = "input_ids"
    _supports_sdpa = True
    _supports_flash_attn = True
    _keep_in_fp32_modules_strict: ClassVar[list[str]] = [r"\.(a|b)$", r"^head\."]
    _no_split_modules: ClassVar[list[str]] = [
        "Qwen3OmniMoeThinkerTextDecoderLayer",
        "ResidualAdapter",
        "OptionHead",
    ]

    def __init__(self, config):
        super().__init__(config)
        self.backbone = Qwen3OmniMoeThinkerForConditionalGeneration(
            config.thinker_config
        )
        self.backbone.lm_head = nn.Identity()
        for name, module in list(self.backbone.named_modules()):
            if name.startswith("model.layers.") and name.endswith(
                (".self_attn.q_proj", ".self_attn.v_proj")
            ):
                parent, leaf = name.rsplit(".", 1)
                setattr(
                    self.backbone.get_submodule(parent),
                    leaf,
                    ResidualAdapter(module, config.adapter_rank, config.adapter_alpha),
                )
        self.head = OptionHead(
            config.thinker_config.text_config.hidden_size, config.head_width
        )
        self.post_init()

    def get_input_embeddings(self):
        return self.backbone.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.backbone.set_input_embeddings(value)

    def forward(
        self, input_ids=None, attention_mask=None, scoring_positions=None, **kwargs
    ):
        if input_ids is None or input_ids.ndim != 2 or input_ids.shape[0] != 1:
            raise ValueError("SPEV accepts one evidence sequence per forward pass")
        if input_ids.shape[1] > self.config.max_tokens:
            raise ValueError("Input exceeds the configured context limit")
        if (
            scoring_positions is None
            or scoring_positions.ndim != 2
            or scoring_positions.shape[0] != 1
        ):
            raise ValueError("scoring_positions must have shape (1, candidates)")
        if (
            scoring_positions.numel() == 0
            or torch.any(scoring_positions < 0)
            or torch.any(scoring_positions >= input_ids.shape[1])
        ):
            raise ValueError("Scoring positions are outside the input sequence")
        kwargs.pop("use_cache", None)
        kwargs.pop("return_dict", None)
        hidden = self.backbone(
            input_ids=input_ids,
            attention_mask=attention_mask,
            use_cache=False,
            return_dict=True,
            **kwargs,
        ).logits
        positions = scoring_positions.to(hidden.device)
        features = hidden.gather(
            1, positions.unsqueeze(-1).expand(-1, -1, hidden.shape[-1])
        )
        head_device = self.head.scorer[0].weight.device
        logits = self.head(features.to(device=head_device, dtype=torch.float32))
        count = (
            attention_mask.sum(-1)
            if attention_mask is not None
            else input_ids.new_tensor([input_ids.shape[1]])
        )
        return SPEVOutput(
            logits=logits, input_tokens=count, temperature=self.config.temperature
        )


SPEVModel.register_for_auto_class("AutoModel")