# coding=utf-8
"""KenT: kiến trúc hồi quy hai luồng (representation X_t / temporal-memory H_t).

Mỗi KenT layer thực hiện một bước, với ĐÚNG MỘT WNN Block:

    V   = B_in(Norm_X(X))                      # (…, N, d_h)
    Z   = A_N V                                # DCT-II cố định trên trục N
    u   = A_N [ Z + SwiGLU(Norm(Z)) ]          # một WNN Block
    H_t = (1 - a) H_{t-1} + a u_t              # EMA, a = a_layer học được, riêng từng layer
    X'  = X + B_out(Norm_H(H_t))

Độ sâu đến từ việc xếp chồng nhiều layer (mỗi layer một WNN), không còn chuỗi L WNN trong
một layer. Bộ nhớ H có cơ chế quên: a_layer = sigmoid(ema_logit) là một số vô hướng học được
cho mỗi layer; khởi tạo giảm dần theo độ sâu (layer nông quên nhanh, layer sâu nhớ lâu).

EMA là một phép hồi quy tuyến tính nên vẫn song song hóa được theo thời gian: trong từng khối
độ dài `ema_block_size` dùng ma trận phân rã tam giác dưới, các khối được ghép bằng
affine-scan log-depth. Kết quả bằng đúng hồi quy tuần tự. Trạng thái giải mã chỉ là H cho mỗi
layer: O(1) theo độ dài.

Trọng số được lưu ở FP16; EMA và normalization chạy ở FP32.
"""

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.

    Args:
        vocab_size: kích thước từ vựng.
        hidden_size: số chiều của representation stream X_t (d_x).
        num_hidden_layers: số KenT layer xếp chồng (mỗi layer có đúng một WNN Block).
        resolution_size: N, độ dài trục mà DCT-II tác động lên.
        feature_dim: d_h, số chiều đặc trưng trên mỗi vị trí resolution.
        intermediate_size: kích thước hidden của SwiGLU (trên trục d_h).
        rms_norm_eps: epsilon của RMSNorm.
        use_bias: dùng bias trong Bilinear / SwiGLU.
        initializer_range: độ lệch chuẩn khởi tạo trọng số.
        ema_alpha_first: a_layer khởi tạo của layer đầu (quên nhanh).
        ema_alpha_last: a_layer khởi tạo của layer cuối (nhớ lâu); các layer ở giữa nội suy
            theo cấp số nhân. a_layer vẫn là tham số học được.
        ema_block_size: độ dài khối của EMA song song (chỉ ảnh hưởng tốc độ/bộ nhớ).
        use_cache: trả về KentCache khi suy luận.
    """

    model_type = "kent"

    def __init__(
        self,
        vocab_size: int = 50257,
        hidden_size: int = 512,
        num_hidden_layers: int = 24,
        resolution_size: int = 8,
        feature_dim: int = 512,
        intermediate_size: int = 2048,
        rms_norm_eps: float = 1e-5,
        use_bias: bool = True,
        initializer_range: float = 0.02,
        ema_alpha_first: float = 0.5,
        ema_alpha_last: float = 0.01,
        ema_block_size: int = 128,
        use_cache: bool = True,
        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.resolution_size = resolution_size
        self.feature_dim = feature_dim
        self.intermediate_size = intermediate_size
        self.rms_norm_eps = rms_norm_eps
        self.use_bias = use_bias
        self.initializer_range = initializer_range
        self.ema_alpha_first = ema_alpha_first
        self.ema_alpha_last = ema_alpha_last
        self.ema_block_size = ema_block_size
        self.use_cache = use_cache
        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,
        )

    def alpha_init(self, layer_idx: int) -> float:
        """a_layer khởi tạo của layer `layer_idx` (giảm theo cấp số nhân theo độ sâu)."""
        n = self.num_hidden_layers
        if n <= 1:
            return float(self.ema_alpha_first)
        r = layer_idx / (n - 1)
        return float(self.ema_alpha_first * (self.ema_alpha_last / self.ema_alpha_first) ** r)


# --------------------------------------------------------------------------- #
# Fixed DCT-II operator
# --------------------------------------------------------------------------- #
@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 A_N (cố định, không phải tham số học)."""
    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(z: torch.Tensor) -> torch.Tensor:
    """D_N(Z) = A_N Z, tác động lên trục -2 (độ dài N) của tensor (..., N, d_h)."""
    a = _dct2_matrix(z.shape[-2], z.device, z.dtype)
    return torch.matmul(a, z)


# --------------------------------------------------------------------------- #
# EMA: song song theo thời gian (khối tam giác + affine-scan giữa các khối)
# --------------------------------------------------------------------------- #
def _affine_scan(a: torch.Tensor, b: torch.Tensor):
    """Inclusive scan dọc dim 1 của các ánh xạ affine h -> a*h + b (Hillis-Steele, log K bước).

    a: (B,K) | b: (B,K,F). Phần tử k được áp dụng SAU các phần tử < k.
    Trả về (A, B) sao cho trạng thái sau khối k là A[k]*h_vào + B[k].
    """
    k = a.shape[1]
    off = 1
    while off < k:
        b = torch.cat([b[:, :off], a[:, off:, None] * b[:, :-off] + b[:, off:]], dim=1)
        a = torch.cat([a[:, :off], a[:, off:] * a[:, :-off]], dim=1)
        off *= 2
    return a, b


def ema_scan(u: torch.Tensor, mask: torch.Tensor, h0: torch.Tensor, a: torch.Tensor, block: int = 128):
    """EMA song song, luôn chạy fp32 (tắt autocast để matmul không bị ép về fp16)."""
    with torch.autocast(device_type=u.device.type, enabled=False):
        return _ema_scan_fp32(u.float(), mask, h0.float(), a.float(), block)


def _ema_scan_fp32(u: torch.Tensor, mask: torch.Tensor, h0: torch.Tensor, a: torch.Tensor, block: int = 128):
    """H_t = (1 - a m_t) H_{t-1} + a m_t u_t, tính song song cho mọi t.

    u: (B,T,N,D) fp32 | mask: (B,T) {0,1} (0 = padding: H giữ nguyên) | h0: (B,N,D) fp32
    a: số vô hướng fp32 trong (0,1).
    Trả về H (B,T,N,D) và h_cuối (B,N,D).

    Trong một khối: H_t = sum_{s<=t} a m_s [prod_{s<r<=t}(1 - a m_r)] u_s + [prod_{r<=t}(1 - a m_r)] h_vào,
    với tích tính qua hiệu của cumsum log (luôn <= 0 nên ổn định). Các khối được nối bằng affine-scan.
    """
    bsz, t, n, d = u.shape
    f = n * d
    blk = max(1, min(block, t))
    pad = (-t) % blk
    uf = u.reshape(bsz, t, f)
    m = mask.to(u.dtype)
    if pad:  # padding ở cuối: mask = 0 nên trạng thái không đổi
        uf = F.pad(uf, (0, 0, 0, pad))
        m = F.pad(m, (0, pad))
    k = (t + pad) // blk
    uk = uf.reshape(bsz, k, blk, f)
    am = a * m.reshape(bsz, k, blk)                              # (B,K,blk)

    cum = torch.cumsum(torch.log1p(-am), dim=2)                  # log prod_{r<=t}(1 - a m_r)
    tril = torch.ones(blk, blk, dtype=torch.bool, device=u.device).tril()
    decay = torch.exp((cum.unsqueeze(-1) - cum.unsqueeze(-2)).masked_fill(~tril, float("-inf")))
    local = torch.matmul(decay * am.unsqueeze(-2), uk)           # (B,K,blk,F): khối bắt đầu từ H = 0

    a_blk = torch.exp(cum[..., -1])                              # (B,K): hệ số phân rã cả khối
    a_p, b_p = _affine_scan(a_blk, local[:, :, -1])
    h0f = h0.reshape(bsz, 1, f)
    s = a_p.unsqueeze(-1) * h0f + b_p                            # trạng thái sau mỗi khối
    h_in = torch.cat([h0f, s[:, :-1]], dim=1)                    # trạng thái vào mỗi khối
    h = local + torch.exp(cum).unsqueeze(-1) * h_in.unsqueeze(2)
    h = h.reshape(bsz, k * blk, f)[:, :t].reshape(bsz, t, n, d)
    return h, s[:, -1].reshape(bsz, n, d)


# --------------------------------------------------------------------------- #
# Cache (trạng thái H cho từng layer)
# --------------------------------------------------------------------------- #
class KentCache:
    """Trạng thái hồi quy của KenT: H (EMA, fp32) theo layer."""

    def __init__(self, config: KentConfig, batch_size: int = 1, device=None):
        shape = (batch_size, config.resolution_size, config.feature_dim)
        self.h = [
            torch.zeros(shape, dtype=torch.float32, device=device)
            for _ in range(config.num_hidden_layers)
        ]
        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 KentBilinear(nn.Module):
    """Bilinear projection: y = (x W_a + b_a) ⊙ (x W_b + b_b)."""

    def __init__(self, in_features: int, out_features: int, bias: bool = True):
        super().__init__()
        self.proj = nn.Linear(in_features, 2 * out_features, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        a, b = self.proj(x).chunk(2, dim=-1)
        return a * b


class KentSwiGLU(nn.Module):
    """SwiGLU(x) = (SiLU(x W_g + b_g) ⊙ (x W_v + b_v)) W_o + b_o."""

    def __init__(self, dim: int, hidden: int, bias: bool = True):
        super().__init__()
        self.gate_value = nn.Linear(dim, 2 * hidden, bias=bias)
        self.out = nn.Linear(hidden, dim, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        g, v = self.gate_value(x).chunk(2, dim=-1)
        return self.out(F.silu(g) * v)


class KentWNNBlock(nn.Module):
    """Z' = A_N [ Z + SwiGLU(Norm(Z)) ]  (residual trước DCT-II)."""

    def __init__(self, config: KentConfig):
        super().__init__()
        self.norm = KentRMSNorm(config.feature_dim, config.rms_norm_eps)
        self.swiglu = KentSwiGLU(config.feature_dim, config.intermediate_size, config.use_bias)

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        return dct2(z + self.swiglu(self.norm(z)))


class KentLayer(nn.Module):
    """Một bước KenT: (X_t, H_{t-1}) -> (X_{t+1}, H_t), với một WNN Block và EMA học được."""

    def __init__(self, config: KentConfig, layer_idx: int = 0):
        super().__init__()
        self.n = config.resolution_size
        self.d = config.feature_dim
        self.block_size = config.ema_block_size
        self.norm_x = KentRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.b_in = KentBilinear(config.hidden_size, self.n * self.d, config.use_bias)
        self.wnn = KentWNNBlock(config)
        self.norm_h = KentRMSNorm(config.feature_dim, config.rms_norm_eps)
        self.b_out = KentBilinear(self.n * self.d, config.hidden_size, config.use_bias)
        a0 = config.alpha_init(layer_idx)
        self.ema_logit_init = math.log(a0 / (1.0 - a0))
        self.ema_logit = nn.Parameter(torch.tensor(self.ema_logit_init))

    def alpha(self) -> torch.Tensor:
        """a_layer in (0,1), tính ở fp32."""
        return torch.sigmoid(self.ema_logit.float()).clamp_max(0.999)

    def forward(self, x, mask, h0, chunked: bool = False):
        b, t, _ = x.shape
        v = self.b_in(self.norm_x(x)).view(b, t, self.n, self.d)
        u = self.wnn(dct2(v)).float()                 # Z^(0) = A_N V -> một WNN Block; EMA ở fp32
        if chunked:  # trục batch là các chunk liên tiếp: nối thành một chuỗi cho EMA
            u = u.reshape(1, b * t, self.n, self.d)
            m = mask.reshape(1, b * t)
        else:
            m = mask
        if h0 is None:
            h0 = u.new_zeros((u.shape[0], self.n, self.d))
        h, h_last = ema_scan(u, m, h0, self.alpha(), self.block_size)
        h = h.reshape(b, t, self.n, self.d)
        dx = self.b_out(self.norm_h(h).to(x.dtype).reshape(b, t, self.n * self.d))
        return x + dx, h_last


# --------------------------------------------------------------------------- #
# 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 = ["KentLayer"]
    _is_stateful = True

    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, KentLayer):
            with torch.no_grad():
                module.ema_logit.fill_(module.ema_logit_init)
        if isinstance(module, KentBilinear) and getattr(module, "_residual_out", False):
            # tích của hai nhánh tuyến tính => thu nhỏ mỗi nhánh (2L)^(-1/4)
            with torch.no_grad():
                module.proj.weight.mul_((2 * self.config.num_hidden_layers) ** -0.25)


class KentModel(KentPreTrainedModel):
    def __init__(self, config: KentConfig):
        super().__init__(config)
        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([KentLayer(config, i) for i in range(config.num_hidden_layers)])
        for layer in self.layers:
            layer.b_out._residual_out = True
        self.norm_f = KentRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.gradient_checkpointing = False
        self.post_init()

    def get_input_embeddings(self):
        return self.embeddings

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

    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 (trạng thái sau chunk cuối).
        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, _ = 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)

        all_hidden = () if output_hidden_states else None
        for i, layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden += (x,)
            h0 = cache_params.h[i] if cache_params is not None else None
            if self.gradient_checkpointing and self.training:
                x, h_last = self._gradient_checkpointing_func(layer.__call__, x, mask, h0, chunked)
            else:
                x, h_last = layer(x, mask, h0, chunked)
            if use_cache:
                cache_params.h[i] = h_last
        if use_cache:
            cache_params.steps += b * t if chunked else t

        x = self.norm_f(x)
        if chunked:  # nối các chunk lại thành một chuỗi (1, K*C, d)
            x = x.reshape(1, b * t, -1)
            if all_hidden is not None:
                all_hidden = tuple(h.reshape(1, b * t, -1) for h in all_hidden)
        if all_hidden is not None:
            all_hidden += (x,)

        if not return_dict:
            return tuple(v for v in (x, cache_params, all_hidden) if v is not None)
        return KentOutput(
            last_hidden_state=x,
            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",
    "ema_scan",
]