"""Fixed-shape Qwen3 candidate encoder and complete NanoJev decision head."""

from __future__ import annotations

import math

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

MASK = -1e4


def rotate_half(value: torch.Tensor) -> torch.Tensor:
    first, second = value.chunk(2, dim=-1)
    return torch.cat((-second, first), dim=-1)


def rms_norm(hidden: torch.Tensor, norm: nn.Module, width: int) -> torch.Tensor:
    doubled = torch.cat((hidden, -hidden), dim=-1)
    normalized = F.layer_norm(doubled, (width * 2,), eps=norm.variance_epsilon)
    return normalized[..., :width] * norm.weight


class NanoEncoder(nn.Module):
    """Encode the complete set of candidate paths with the trained Qwen3 body."""

    def __init__(self, decision_model: nn.Module, length: int, candidates: int):
        super().__init__()
        body = decision_model.backbone
        config = body.config
        self.embed_tokens = body.embed_tokens
        self.layers = body.layers
        self.norm = body.norm
        self.length = length
        self.candidates = candidates
        self.heads = config.num_attention_heads
        self.kv_heads = config.num_key_value_heads
        self.groups = self.heads // self.kv_heads
        self.head_dim = config.head_dim
        self.hidden = config.hidden_size
        self.scale = self.head_dim**-0.5
        positions = torch.arange(length).unsqueeze(0)
        probe = torch.zeros(1, length, config.hidden_size)
        cos, sin = body.rotary_emb(probe, positions)
        self.register_buffer("cos", cos.unsqueeze(1).detach().clone())
        self.register_buffer("sin", sin.unsqueeze(1).detach().clone())
        causal = torch.triu(torch.full((length, length), MASK), diagonal=1)
        self.register_buffer("causal", causal.view(1, 1, length, length))

    def _layer(self, layer: nn.Module, hidden: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        attention = layer.self_attn
        normalized = rms_norm(hidden, layer.input_layernorm, self.hidden)
        q = attention.q_proj(normalized).view(self.candidates, self.length, self.heads, self.head_dim)
        k = attention.k_proj(normalized).view(self.candidates, self.length, self.kv_heads, self.head_dim)
        v = attention.v_proj(normalized).view(self.candidates, self.length, self.kv_heads, self.head_dim)
        q = attention.q_norm(q).transpose(1, 2)
        k = attention.k_norm(k).transpose(1, 2)
        v = v.transpose(1, 2)
        q = q * self.cos + rotate_half(q) * self.sin
        k = k * self.cos + rotate_half(k) * self.sin
        k = k.repeat_interleave(self.groups, dim=1)
        v = v.repeat_interleave(self.groups, dim=1)
        weights = torch.matmul(q, k.transpose(2, 3)) * self.scale + mask
        weights = torch.softmax(weights.float(), dim=-1).to(q.dtype)
        attended = torch.matmul(weights, v).transpose(1, 2).reshape(
            self.candidates, self.length, self.heads * self.head_dim
        )
        hidden = hidden + attention.o_proj(attended)
        return hidden + layer.mlp(rms_norm(hidden, layer.post_attention_layernorm, self.hidden))

    def forward(
        self, input_ids: torch.Tensor, attention_mask: torch.Tensor, eos_map: torch.Tensor
    ) -> torch.Tensor:
        padding = (1.0 - attention_mask.float()).view(self.candidates, 1, 1, self.length) * MASK
        mask = self.causal + padding
        hidden = self.embed_tokens(input_ids.long())
        for layer in self.layers:
            hidden = self._layer(layer, hidden, mask)
        hidden = rms_norm(hidden, self.norm, self.hidden).float()
        return torch.matmul(eos_map, hidden).transpose(0, 1)


class NanoHead(nn.Module):
    """The trained scalar and set-attention heads for Choice, Boolean, and Score."""

    def __init__(self, decision_model: nn.Module, candidates: int):
        super().__init__()
        self.norm = decision_model.norm
        self.scalar = decision_model.scalar
        self.set_project = decision_model.set_project
        self.set_attention = decision_model.set_attention
        self.set_output = decision_model.set_output
        self.candidates = candidates
        self.heads = 4
        self.head_dim = 32
        self.scale = 1.0 / math.sqrt(self.head_dim)
        boolean_mask = torch.zeros(1, candidates)
        boolean_mask[0, :2] = 1
        self.register_buffer("boolean_mask", boolean_mask)

    def forward(
        self,
        embeddings: torch.Tensor,
        candidate_mask: torch.Tensor,
        use_set_head: torch.Tensor,
        is_boolean: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        h = self.norm(embeddings)
        scalar = self.scalar(h).squeeze(-1).float()
        log_k = candidate_mask.sum(-1).clamp(min=1).float().log().view(1, 1, 1)
        log_k = log_k.expand(1, self.candidates, 1)
        u = self.set_project(torch.cat((h, log_k.to(h.dtype)), dim=-1))
        projected = F.linear(u, self.set_attention.in_proj_weight, self.set_attention.in_proj_bias)
        q, k, v = projected.chunk(3, dim=-1)
        q = q.view(1, self.candidates, self.heads, self.head_dim).transpose(1, 2)
        k = k.view(1, self.candidates, self.heads, self.head_dim).transpose(1, 2)
        v = v.view(1, self.candidates, self.heads, self.head_dim).transpose(1, 2)
        scores = torch.matmul(q, k.transpose(2, 3)) * self.scale
        scores = scores + (1.0 - candidate_mask).view(1, 1, 1, self.candidates) * MASK
        mixed = torch.matmul(torch.softmax(scores.float(), dim=-1).to(v.dtype), v)
        mixed = mixed.transpose(1, 2).reshape(1, self.candidates, -1)
        mixed = self.set_attention.out_proj(mixed)
        delta = self.set_output(torch.tanh(u + mixed)).squeeze(-1).float()
        normal_logits = scalar + use_set_head * delta
        boolean_logits = torch.cat((scalar[:, :1] * 0, scalar[:, :1], scalar[:, 2:] * 0), dim=-1)
        logits = normal_logits * (1.0 - is_boolean) + boolean_logits * is_boolean
        output_mask = candidate_mask * (1.0 - is_boolean) + self.boolean_mask * is_boolean
        logits = logits * output_mask + (1.0 - output_mask) * MASK
        return logits, torch.softmax(logits, dim=-1)