# coding=utf-8
"""KenT-20M: mạng hồi quy một-token-một-bước với bộ nhớ EMA học được ở MỌI độ sâu.

Mỗi bước nhận ĐÚNG MỘT token x[t] (vector 1 x D); lịch sử nằm hoàn toàn trong trạng thái H (EMA). Sơ đồ:

    h_0 = x[t]                                   # embedding (Tokenizer)
    h_i = h_{i-1} + SwiGLU( RMSNorm(h_{i-1}) )   # i = 1..n   (khối xanh, residual)

    u_i = RMSNorm(h_i) · D^T · W_i               # i = 0..n   (khối cam "~DCT-II Proj (D^T)W": DCT-II cố định rồi W học được)
    H_i[t+1] = (1 - A_i) ⊙ H_i[t] + A_i ⊙ u_i    # "Learnable EMA Matrix": A = sigmoid(ema_logit), shape (n+1, D)

    S   = sum_i H_i[t+1]                         # Add
    y   = S · D                                  # "DCT-II Proj (D)": DCT ngược (D trực chuẩn, D^T D = I)
    out = y + SwiGLU( RMSNorm(y) )               # RMSNorm -> SwiGLU -> Add (skip lấy sau phép chiếu D)
    logits = LM_head( RMSNorm_f(out) )           # head buộc chung trọng số với embedding

Quy ước hàng: x là vector hàng nên DCT-II là x·D^T (D[k, i] = α_k cos(π(i+½)k/D), hàng = tần số k, cột = vị trí i)
và phép chiếu ngược là y·D. W khởi tạo ≈ ma trận đơn vị nên lúc đầu khối cam đúng bằng DCT-II và khối ra đúng
bằng DCT ngược.

Độ sâu (n khối xanh) cho n+1 "điểm lấy mẫu" (tap), mỗi tap có bộ nhớ EMA riêng H_i shape (D,). Trạng thái giải mã
toàn mô hình chỉ là H shape (n+1, D): O(1) theo độ dài chuỗi.

EMA là hồi quy tuyến tính theo từng kênh nên song song hóa được theo thời gian bằng affine-scan log-depth
(Hillis-Steele): kết quả bằng đúng hồi quy tuần tự. EMA và RMSNorm chạy ở FP32; trọng số lưu ở FP16.

RMSNorm cuối (`norm_f`) trước LM head không có trong sơ đồ; thêm vào để logits ổn định.
"""

import math
from dataclasses import dataclass
from functools import lru_cache
from typing import Optional, Tuple, Union

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

from transformers import GenerationMixin, PretrainedConfig, PreTrainedModel
from transformers.utils import ModelOutput


# --------------------------------------------------------------------------- #
# Config
# --------------------------------------------------------------------------- #
class KentConfig(PretrainedConfig):
    r"""Cấu hình KenT-20M.

    Args:
        vocab_size: kích thước từ vựng.
        hidden_size: D, độ rộng của vector token x[t] (cũng là độ rộng của mọi tap, của DCT và của EMA).
        num_hidden_layers: n, số khối xanh (SwiGLU residual) nối tiếp. Có n + 1 tap EMA (tap 0 lấy từ embedding).
        intermediate_size: kích thước hidden của mọi SwiGLU. Nên >= 2D để SwiGLU không thành nút cổ chai.
        rms_norm_eps: epsilon của RMSNorm.
        initializer_range: độ lệch chuẩn khởi tạo trọng số.
        ema_alpha_first: hệ số a khởi tạo của tap đầu, feature đầu (quên nhanh).
        ema_alpha_last: hệ số a khởi tạo của tap cuối, feature cuối (nhớ lâu); phần còn lại nội suy theo cấp số nhân
            theo trung bình chỉ số (tap, feature). a vẫn là tham số học được.
        use_cache: trả về KentCache khi suy luận.
        use_gradient_checkpointing: bật gradient checkpointing theo từng khối xanh khi khởi tạo model.
    """

    model_type = "kent"

    def __init__(
        self,
        vocab_size: int = 50257,
        hidden_size: int = 240,
        num_hidden_layers: int = 12,
        intermediate_size: int = 720,
        rms_norm_eps: float = 1e-5,
        initializer_range: float = 0.02,
        ema_alpha_first: float = 0.5,
        ema_alpha_last: float = 0.01,
        use_cache: bool = True,
        use_gradient_checkpointing: bool = False,
        pad_token_id: int = 50256,
        bos_token_id: int = 50256,
        eos_token_id: int = 50256,
        tie_word_embeddings: bool = True,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.intermediate_size = intermediate_size
        self.rms_norm_eps = rms_norm_eps
        self.initializer_range = initializer_range
        self.ema_alpha_first = ema_alpha_first
        self.ema_alpha_last = ema_alpha_last
        self.use_cache = use_cache
        self.use_gradient_checkpointing = use_gradient_checkpointing
        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )

    @property
    def num_taps(self) -> int:
        return self.num_hidden_layers + 1

    def alpha_init_matrix(self):
        """Ma trận a ((n+1) x D) khởi tạo: giảm theo cấp số nhân theo độ sâu của tap và theo chỉ số feature."""
        n_t, n_d = self.num_taps, self.hidden_size
        ratio = self.ema_alpha_last / self.ema_alpha_first
        return [[float(self.ema_alpha_first * ratio ** (0.5 * ((i / (n_t - 1) if n_t > 1 else 0.0)
                                                              + (j / (n_d - 1) if n_d > 1 else 0.0))))
                 for j in range(n_d)] for i in range(n_t)]

    def alpha_init(self, tap_idx: int) -> float:
        """Trung bình (theo feature) của a khởi tạo của tap `tap_idx` (để log)."""
        row = self.alpha_init_matrix()[tap_idx]
        return float(sum(row) / len(row))


# --------------------------------------------------------------------------- #
# Fixed DCT-II operator (trục cuối, vector hàng)
# --------------------------------------------------------------------------- #
@lru_cache(maxsize=32)
def _dct2_matrix(n: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    """Ma trận DCT-II trực chuẩn D (cố định, không phải tham số học): D[k, i] = α_k cos(π(i+½)k/n)."""
    k = torch.arange(n, dtype=torch.float64).unsqueeze(1)
    i = torch.arange(n, dtype=torch.float64).unsqueeze(0)
    a = torch.cos(math.pi / n * (i + 0.5) * k)
    alpha = torch.full((n, 1), math.sqrt(2.0 / n), dtype=torch.float64)
    alpha[0, 0] = math.sqrt(1.0 / n)
    return (alpha * a).to(device=device, dtype=dtype)


def dct2(x: torch.Tensor) -> torch.Tensor:
    """DCT-II của vector hàng trên trục cuối: x · D^T (tương đương D x cho vector cột)."""
    d = _dct2_matrix(x.shape[-1], x.device, x.dtype)
    return torch.matmul(x, d.t())


def idct2(x: torch.Tensor) -> torch.Tensor:
    """Phép chiếu ngược: x · D (D trực chuẩn nên idct2(dct2(x)) = x)."""
    d = _dct2_matrix(x.shape[-1], x.device, x.dtype)
    return torch.matmul(x, d)


# --------------------------------------------------------------------------- #
# EMA song song theo thời gian (affine-scan Hillis-Steele, hệ số riêng từng kênh)
# --------------------------------------------------------------------------- #
def ema_scan(u: torch.Tensor, mask: torch.Tensor, h0: torch.Tensor, a: torch.Tensor):
    """H_t = (1 - a m_t) ⊙ H_{t-1} + a m_t ⊙ u_t với a: (N, D) khác nhau theo từng kênh; luôn chạy fp32.

    u: (B,T,N,D) | mask: (B,T) {0,1} (0 = padding: H giữ nguyên) | h0: (B,N,D).
    Mỗi bước là ánh xạ affine h -> A_t h + B_t (A_t = 1 - a m_t, B_t = a m_t u_t); inclusive scan log-depth cho
    trạng thái sau mọi t. Mọi A_t nằm trong (0, 1] nên ổn định. Trả về H (B,T,N,D) và h_cuối (B,N,D).
    """
    with torch.autocast(device_type=u.device.type, enabled=False):
        u, h0, a = u.float(), h0.float(), a.float()
        bsz, t, n, d = u.shape
        am = a.view(1, 1, n, d) * mask.to(u.dtype).view(bsz, t, 1, 1)
        dec = 1.0 - am
        acc = am * u
        off = 1
        while off < t:
            acc = torch.cat([acc[:, :off], dec[:, off:] * acc[:, :-off] + acc[:, off:]], dim=1)
            dec = torch.cat([dec[:, :off], dec[:, off:] * dec[:, :-off]], dim=1)
            off *= 2
        h = dec * h0.unsqueeze(1) + acc
        return h, h[:, -1]


# --------------------------------------------------------------------------- #
# Cache (trạng thái H của toàn mô hình)
# --------------------------------------------------------------------------- #
class KentCache:
    """Trạng thái hồi quy của KenT: một tensor H (EMA, fp32) shape (B, n+1, D), giữ trong danh sách `h` (1 phần tử)."""

    def __init__(self, config: KentConfig, batch_size: int = 1, device=None):
        shape = (batch_size, config.num_taps, config.hidden_size)
        self.h = [torch.zeros(shape, dtype=torch.float32, device=device)]
        self.steps = 0  # số vị trí đã nạp (đếm ở phía Python, không cần sync GPU)

    @property
    def batch_size(self) -> int:
        return self.h[0].shape[0]

    @property
    def seen_tokens(self) -> int:
        return self.steps

    def get_seq_length(self, layer_idx: int = 0) -> int:
        return self.steps

    def reset(self):
        for h in self.h:
            h.zero_()
        self.steps = 0

    def clone(self) -> "KentCache":
        new = KentCache.__new__(KentCache)
        new.h = [h.clone() for h in self.h]
        new.steps = self.steps
        return new

    def to(self, device):
        self.h = [h.to(device) for h in self.h]
        return self

    def reorder_cache(self, beam_idx: torch.LongTensor):
        self.h = [h.index_select(0, beam_idx.to(h.device)) for h in self.h]


# --------------------------------------------------------------------------- #
# Building blocks
# --------------------------------------------------------------------------- #
class KentRMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dtype = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return (x * self.weight.float()).to(dtype)


class KentSwiGLU(nn.Module):
    """SwiGLU: down( silu(x W_g) ⊙ (x W_u) ), không bias."""

    def __init__(self, dim: int, hidden: int):
        super().__init__()
        self.gate_up = nn.Linear(dim, 2 * hidden, bias=False)
        self.down = nn.Linear(hidden, dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        g, u = self.gate_up(x).chunk(2, dim=-1)
        return self.down(F.silu(g) * u)


class KentBlock(nn.Module):
    """Khối xanh: h' = h + SwiGLU(RMSNorm(h)) (skip bắt đầu trước RMSNorm, như sơ đồ)."""

    def __init__(self, config: KentConfig):
        super().__init__()
        self.norm = KentRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mlp = KentSwiGLU(config.hidden_size, config.intermediate_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.mlp(self.norm(x))


class KentTap(nn.Module):
    """Khối cam "~DCT-II Proj (D^T)W": u = RMSNorm(h) · D^T · W. D cố định, W (D x D) học được, khởi tạo ≈ I."""

    def __init__(self, config: KentConfig):
        super().__init__()
        self.norm = KentRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)  # W
        self.proj._identity_init = True

    def forward(self, h: torch.Tensor) -> torch.Tensor:
        return self.proj(dct2(self.norm(h)))


# --------------------------------------------------------------------------- #
# Outputs
# --------------------------------------------------------------------------- #
@dataclass
class KentOutput(ModelOutput):
    last_hidden_state: Optional[torch.FloatTensor] = None
    cache_params: Optional[KentCache] = None
    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None


@dataclass
class KentCausalLMOutput(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    logits: Optional[torch.FloatTensor] = None
    cache_params: Optional[KentCache] = None
    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None


# --------------------------------------------------------------------------- #
# Models
# --------------------------------------------------------------------------- #
class KentPreTrainedModel(PreTrainedModel):
    config_class = KentConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["KentBlock"]
    _is_stateful = True

    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
        """Bật checkpoint theo khối xanh; mặc định use_reentrant=False."""
        if gradient_checkpointing_kwargs is None:
            gradient_checkpointing_kwargs = {"use_reentrant": False}
        super().gradient_checkpointing_enable(gradient_checkpointing_kwargs=gradient_checkpointing_kwargs)

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=std)
        elif isinstance(module, KentRMSNorm):
            nn.init.ones_(module.weight)
        elif isinstance(module, KentModel):
            with torch.no_grad():
                a0 = torch.tensor(self.config.alpha_init_matrix(), dtype=torch.float32)
                module.ema_logit.copy_(torch.log(a0 / (1.0 - a0)).to(module.ema_logit.device))
        if isinstance(module, nn.Linear) and getattr(module, "_identity_init", False):
            # W của khối cam ≈ I: khối cam khởi đầu đúng bằng DCT-II (làm ở mức Linear vì thứ tự duyệt module của HF)
            with torch.no_grad():
                w = module.weight
                w.copy_(torch.eye(w.shape[0], dtype=w.dtype, device=w.device) + std * torch.randn_like(w))
        if isinstance(module, nn.Linear) and getattr(module, "_residual_out", False):
            # đầu ra ghi vào residual stream => thu nhỏ theo (2 (n+1))^(-1/2) (kiểu GPT-2)
            with torch.no_grad():
                module.weight.mul_((2 * (self.config.num_hidden_layers + 1)) ** -0.5)


class KentModel(KentPreTrainedModel):
    def __init__(self, config: KentConfig):
        super().__init__(config)
        d = config.hidden_size
        self.embeddings = nn.Embedding(config.vocab_size, d)
        self.blocks = nn.ModuleList([KentBlock(config) for _ in range(config.num_hidden_layers)])
        self.taps = nn.ModuleList([KentTap(config) for _ in range(config.num_taps)])
        self.ema_logit = nn.Parameter(torch.zeros(config.num_taps, d))   # Learnable EMA Matrix: a = sigmoid(logit)
        self.out_norm = KentRMSNorm(d, config.rms_norm_eps)              # RMSNorm sau phép chiếu D
        self.out_mlp = KentSwiGLU(d, config.intermediate_size)           # SwiGLU cuối
        self.norm_f = KentRMSNorm(d, config.rms_norm_eps)
        for blk in self.blocks:
            blk.mlp.down._residual_out = True
        self.out_mlp.down._residual_out = True
        self.gradient_checkpointing = False
        self.post_init()
        if getattr(config, "use_gradient_checkpointing", False):
            self.gradient_checkpointing_enable()

    def get_input_embeddings(self):
        return self.embeddings

    def set_input_embeddings(self, value):
        self.embeddings = value

    def alpha(self) -> torch.Tensor:
        """EMA Learner Matrix a ((n+1) x D) trong (0,1), tính ở fp32."""
        return torch.sigmoid(self.ema_logit.float()).clamp_max(0.999)

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        cache_params: Optional[KentCache] = None,
        use_cache: Optional[bool] = None,
        chunked: bool = False,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, KentOutput]:
        """
        chunked=False: trục batch là các chuỗi độc lập; (B, T).
        chunked=True : trục batch là K chunk liên tiếp của MỘT tài liệu; (K, C). Cache vào/ra có batch size 1.
        attention_mask: 1 = token thật, 0 = padding (H giữ nguyên tại vị trí đó).
        """
        use_cache = use_cache if use_cache is not None else (self.config.use_cache and not self.training)
        output_hidden_states = bool(output_hidden_states)
        return_dict = return_dict if return_dict is not None else True

        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Cần đúng một trong input_ids hoặc inputs_embeds.")
        x = self.embeddings(input_ids) if inputs_embeds is None else inputs_embeds
        b, t, d = x.shape

        if attention_mask is None:
            mask = torch.ones(b, t, dtype=torch.long, device=x.device)
        else:
            mask = attention_mask.to(device=x.device, dtype=torch.long)
            if mask.shape[1] != t:  # generate() truyền mask toàn chuỗi
                mask = mask[:, -t:]

        if cache_params is not None and cache_params.batch_size != (1 if chunked else b):
            raise ValueError("Batch size của cache không khớp với đầu vào.")
        if use_cache and cache_params is None:
            cache_params = KentCache(self.config, 1 if chunked else b, x.device)

        # --- chuỗi khối xanh (mọi token độc lập nhau) -> n+1 điểm lấy mẫu h_0..h_n
        all_hidden = (x,) if output_hidden_states else None
        feats = [x]
        for blk in self.blocks:
            if self.gradient_checkpointing and self.training:
                x = self._gradient_checkpointing_func(blk.__call__, x)
            else:
                x = blk(x)
            feats.append(x)
            if output_hidden_states:
                all_hidden += (x,)

        # --- khối cam: u_i = RMSNorm(h_i) D^T W_i, rồi EMA riêng từng tap
        u = torch.stack([tap(h) for tap, h in zip(self.taps, feats)], dim=2).float()   # (B,T,n+1,D)
        if chunked:  # K chunk liên tiếp của một tài liệu: nối thành một chuỗi cho EMA
            u = u.reshape(1, b * t, self.config.num_taps, d)
            m = mask.reshape(1, b * t)
        else:
            m = mask
        h0 = cache_params.h[0] if cache_params is not None else u.new_zeros((u.shape[0], self.config.num_taps, d))
        h, h_last = ema_scan(u, m, h0, self.alpha())                                    # (.,.,n+1,D), (.,n+1,D)
        if use_cache:
            cache_params.h[0] = h_last
            cache_params.steps += b * t if chunked else t

        # --- Add -> DCT-II Proj (D) -> RMSNorm -> SwiGLU -> Add
        s = h.sum(dim=2).to(feats[0].dtype)
        y = idct2(s)
        y = y + self.out_mlp(self.out_norm(y))
        y = self.norm_f(y)                                                              # (1, K*C, D) nếu chunked
        if chunked:
            y = y.reshape(1, b * t, d)
        if output_hidden_states:
            all_hidden += (y,)

        if not return_dict:
            return tuple(v for v in (y, cache_params if use_cache else None, all_hidden) if v is not None)
        return KentOutput(
            last_hidden_state=y,
            cache_params=cache_params if use_cache else None,
            hidden_states=all_hidden,
        )


class KentForCausalLM(KentPreTrainedModel, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "model.embeddings.weight"}

    def __init__(self, config: KentConfig):
        super().__init__(config)
        self.model = KentModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

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

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

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def new_cache(self, batch_size: int = 1) -> KentCache:
        p = next(self.parameters())
        return KentCache(self.config, batch_size, p.device)

    # ---- tích hợp generate() ------------------------------------------------
    def prepare_inputs_for_generation(
        self,
        input_ids,
        inputs_embeds=None,
        attention_mask=None,
        use_cache=None,
        cache_params=None,
        past_key_values=None,
        **kwargs,
    ):
        if cache_params is None and isinstance(past_key_values, KentCache):
            cache_params = past_key_values
        if cache_params is not None and cache_params.steps > 0:
            input_ids = input_ids[:, -1:]  # trạng thái đã chứa phần còn lại
            inputs_embeds = None
        inputs = {"inputs_embeds": inputs_embeds} if inputs_embeds is not None else {"input_ids": input_ids}
        inputs.update(
            attention_mask=attention_mask,
            cache_params=cache_params,
            use_cache=True if use_cache is None else use_cache,
        )
        return inputs

    def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder=False, **kwargs):
        model_kwargs["cache_params"] = outputs.get("cache_params", None)
        model_kwargs.pop("past_key_values", None)
        if model_kwargs.get("attention_mask", None) is not None:
            am = model_kwargs["attention_mask"]
            model_kwargs["attention_mask"] = torch.cat([am, am.new_ones((am.shape[0], 1))], dim=-1)
        if model_kwargs.get("cache_position", None) is not None:
            model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + 1
        return model_kwargs

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        cache_params: Optional[KentCache] = None,
        past_key_values: Optional[KentCache] = None,
        use_cache: Optional[bool] = None,
        chunked: bool = False,
        labels: Optional[torch.LongTensor] = None,
        logits_to_keep: int = 0,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, KentCausalLMOutput]:
        if cache_params is None and isinstance(past_key_values, KentCache):
            cache_params = past_key_values
        return_dict = return_dict if return_dict is not None else True
        out = self.model(
            input_ids=input_ids,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            cache_params=cache_params,
            use_cache=use_cache,
            chunked=chunked,
            output_hidden_states=output_hidden_states,
            return_dict=True,
        )
        hidden = out.last_hidden_state
        if labels is None and logits_to_keep:
            hidden = hidden[:, -logits_to_keep:]
        logits = self.lm_head(hidden)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous().float()
            shift_labels = labels[..., 1:].contiguous().to(shift_logits.device)
            loss = CrossEntropyLoss()(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

        if not return_dict:
            return tuple(v for v in (loss, logits, out.cache_params, out.hidden_states) if v is not None)
        return KentCausalLMOutput(
            loss=loss,
            logits=logits,
            cache_params=out.cache_params,
            hidden_states=out.hidden_states,
        )


__all__ = [
    "KentConfig",
    "KentCache",
    "KentModel",
    "KentForCausalLM",
    "KentPreTrainedModel",
    "dct2",
    "idct2",
    "ema_scan",
]