"""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
from .branch_attention import install as install_attention, request_context
from .decision_head import IsolatedDecisionHead, QuestionFeatures, score_questions
from .decision_layout import DecisionLayout
from .token_operations import install as install_token_operations


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)


@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",
        "IsolatedDecisionHead",
    ]

    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 = IsolatedDecisionHead(
            config.thinker_config.text_config.hidden_size, config.head_width,
            config.head_heads, config.head_ff, config.head_depth,
        )
        self.post_init()
        install_attention(self.backbone)
        install_token_operations(self, linears=True, moe=False)

    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, decision_layout=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 not isinstance(decision_layout, dict):
            raise ValueError("A processor-generated decision_layout is required")
        decision_layout = DecisionLayout.from_dict(decision_layout)
        decision_layout.validate(input_ids.shape[1])
        tree = decision_layout.tree
        if attention_mask is not None and not torch.all(attention_mask == 1):
            raise ValueError("SPEV accepts unpadded evidence sequences")
        kwargs.pop("use_cache", None)
        kwargs.pop("return_dict", None)
        kwargs.pop("position_ids", None)
        prefix = input_ids[:, :tree.prefix_length]
        rope_kwargs = {key: kwargs[key] for key in (
            "image_grid_thw", "video_grid_thw", "use_audio_in_video", "second_per_grids"
        ) if key in kwargs}
        if "video_second_per_grid" in kwargs:
            rope_kwargs["second_per_grids"] = kwargs["video_second_per_grid"]
        if "feature_attention_mask" in kwargs:
            rope_kwargs["audio_seqlens"] = kwargs["feature_attention_mask"].sum(1)
        positions, _ = self.backbone.get_rope_index(
            prefix, attention_mask=torch.ones_like(prefix), **rope_kwargs
        )
        position_ids = tree.position_ids(positions).to(input_ids.device)
        with request_context(tree):
            hidden = self.backbone(
                input_ids=input_ids,
                attention_mask=tree.additive_mask(dtype=self.backbone.dtype, device=input_ids.device),
                position_ids=position_ids,
                use_cache=False,
                return_dict=True,
                **kwargs,
            ).logits[0]
        lexical = []
        for tokens in decision_layout.lexical_ids:
            tokens = input_ids.new_tensor(tokens)
            lexical.append(self.get_input_embeddings()(tokens).mean(0))
        records = []
        starts = [tree.prefix_length]
        for length in tree.branch_lengths:
            starts.append(starts[-1] + length)
        for question, kind in enumerate(decision_layout.question_types):
            indices = [i for i, owner in enumerate(decision_layout.candidate_questions) if owner == question]
            memory = torch.cat([hidden[:tree.prefix_length]] + [hidden[starts[i]:starts[i + 1]] for i in indices])
            records.append(QuestionFeatures(
                memory=memory,
                question=hidden[list(decision_layout.question_positions[question])].mean(0),
                options=hidden[[decision_layout.anchors[i] for i in indices]],
                lexical=torch.stack([lexical[i] for i in indices]),
                global_feature=hidden[tree.prefix_length - 1],
                question_type=kind,
            ))
        grouped = score_questions(self.head, records)
        logits = torch.empty(len(lexical), device=grouped[0].device, dtype=grouped[0].dtype)
        for question, values in enumerate(grouped):
            indices = [i for i, owner in enumerate(decision_layout.candidate_questions) if owner == question]
            logits[indices] = values
        return SPEVOutput(
            logits=logits[None], input_tokens=input_ids.new_tensor([input_ids.shape[1]]),
            temperature=self.config.temperature
        )


SPEVModel.register_for_auto_class("AutoModel")