"""Текстовый путь UDM: замороженный Ettin + обучаемый выбор слоёв (по идее moss/adapter).

Почему не один слой. Скрытые состояния растут по норме с глубиной (в Qwen3.5-9B — в 40 раз,
13 у слоя 2 против 557 у слоя 27), поэтому «взять последний слой» теряет лексику и цифры,
а «взять ранний» — семантику. В moss/adapter это решили иначе: несколько срезов по K
нормируются ПОШТУЧНО, идут последовательностью длины K через 2 блока self-attention, затем
`Linear(K→1)` схлопывает ось срезов, а результат zero-init'ом добавляется к основному пути.
Тот же приём (плюс 0.8B вместо 9B) дал cos 0.924 к учителю.

Здесь то же самое, но цель — не регрессия к чужому пространству, а обусловливание ДиТ:
выход миксера идёт в `xkv` (кросс-внимание) и в `pmlp` (adaLN).

Срезы Ettin-150M: 22 слоя, global attention каждые 3 → берём выходы после глобальных слоёв.
"""
from __future__ import annotations

import math

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

# выходы после глобальных слоёв ModernBERT (22 слоя, global_attn_every_n_layers=3)
ETTIN_150M_SLICES = (2, 5, 8, 11, 14, 17, 20)


class AttnBlock(nn.Module):
    """Pre-norm self-attention + MLP (для последовательности длины K — оси срезов)."""

    def __init__(self, d: int, heads: int = 8, ffn_mult: float = 2.0):
        super().__init__()
        self.h = heads
        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv = nn.Linear(d, 3 * d)
        self.proj = nn.Linear(d, d)
        self.fc1 = nn.Linear(d, int(d * ffn_mult))
        self.fc2 = nn.Linear(int(d * ffn_mult), d)

    def forward(self, x):
        b, n, d = x.shape
        q, k, v = self.qkv(self.n1(x)).reshape(b, n, 3, self.h, d // self.h).permute(2, 0, 3, 1, 4)
        a = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(b, n, d)
        x = x + self.proj(a)
        return x + self.fc2(F.gelu(self.fc1(self.n2(x))))


class SliceMixer(nn.Module):
    """Выбор слоя: пер-срезный LayerNorm → attention по оси срезов → Linear(K→1) → zero-init out.

    blocks=0 оставляет только обученные веса слоёв (Linear(K→1) поверх нормированных срезов) —
    дешёвый вариант «какой слой важнее», без внимания между срезами.
    """

    def __init__(self, n_slices: int, d: int, blocks: int = 2, heads: int = 8):
        super().__init__()
        self.n, self.d = n_slices, d
        self.per = nn.LayerNorm(d, elementwise_affine=False)
        self.blocks = nn.ModuleList([AttnBlock(d, heads) for _ in range(blocks)])
        self.proj = nn.Linear(n_slices, 1, bias=False)
        self.out = nn.Linear(d, d)
        nn.init.zeros_(self.out.weight)
        nn.init.zeros_(self.out.bias)

    def forward(self, slices):                      # (B, L, K, D)
        b, l, k, d = slices.shape
        normed = self.per(slices)                   # пер-срезная нормировка ОБЯЗАТЕЛЬНА:
        h = normed.reshape(b * l, k, d)             # сырые нормы срезов растут с глубиной (52…548)
        for blk in self.blocks:
            h = blk(h)
        mixed = self.out(self.proj(h.permute(0, 2, 1)).squeeze(-1)).reshape(b, l, d)
        return normed.mean(2) + mixed               # среднее по нормированным срезам + zero-init добавка

    def layer_weights(self) -> torch.Tensor:
        """Что микшер выучил: ненормированные веса слоёв (для логов)."""
        return self.proj.weight.detach().squeeze()


class EttinText(nn.Module):
    """Замороженный энкодер Ettin-150M + обучаемый SliceMixer. Возвращает (B,L,768) и маску."""

    def __init__(self, path: str, device, max_len: int = 192, slices=ETTIN_150M_SLICES,
                 blocks: int = 2, dtype=torch.bfloat16):
        super().__init__()
        from transformers import AutoModel, AutoTokenizer
        self.max_len, self.layer_idx = max_len, tuple(slices)
        self.tok = AutoTokenizer.from_pretrained(path)
        self.model = AutoModel.from_pretrained(path, dtype=dtype).to(device).eval().requires_grad_(False)
        self.dim = self.model.config.hidden_size
        self.mixer = SliceMixer(len(self.layer_idx), self.dim, blocks=blocks).to(device)
        self.device = torch.device(device)

    def tokenize(self, caps: list[str], pad_to: int = 8, pad_to_len: int | None = None):
        # padding="longest" (а не до max_len): короткие капшены стоят дешевле
        t = self.tok(caps, padding="longest", truncation=True, max_length=self.max_len,
                     return_tensors="pt")
        ids, mask = t["input_ids"], t["attention_mask"].bool()
        into = pad_to_len or ids.shape[1]
        extra = into - ids.shape[1]
        if extra > 0:
            ids = F.pad(ids, (0, extra), value=self.tok.pad_token_id)
            mask = F.pad(mask, (0, extra))
        elif extra < 0:
            ids, mask = ids[:, :into], mask[:, :into]
        else:
            extra = (-ids.shape[1]) % pad_to
            if extra:
                ids = F.pad(ids, (0, extra), value=self.tok.pad_token_id)
                mask = F.pad(mask, (0, extra))
        return ids.to(self.device), mask.to(self.device)

    @torch.no_grad()
    def slices(self, caps: list[str], pad_to_len: int | None = None):
        """Срезы слоёв — считаются ОДИН раз на капшен (энкодер заморожен)."""
        ids, mask = self.tokenize(caps, pad_to_len=pad_to_len)
        out = self.model(input_ids=ids, attention_mask=mask, output_hidden_states=True)
        return torch.stack([out.hidden_states[i] for i in self.layer_idx], dim=2), mask   # (B,L,K,D)

    def mix(self, hs):
        """Обучаемый путь: среднее по нормированным срезам + zero-init добавка микшера."""
        return self.mixer(hs.float())

    def forward(self, caps: list[str]):
        hs, mask = self.slices(caps)
        return self.mix(hs), mask, hs

    def trainable(self):
        return list(self.mixer.parameters())


def empty_text(enc: EttinText):
    """Пустой промпт = [CLS][SEP] — безусловный контекст для CFG."""
    return enc([""])