# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""HuggingFace-compatible PanoLM-KDA model (Kimi Delta Attention).

Standalone, KDA-locked variant of the PanoLM HF wrapper.

Layer tree (matches the torchtitan ``PanoLMStateDictAdapter`` output verbatim):

    PanoLMForCausalLM
    └── model
        └── text                           PanoLMTextModel
            ├── tok_embeddings             nn.Embedding
            ├── lower_bounds               nn.Parameter (n_layers, hidden_size) — only when use_lower_bound
            ├── layers.{i}                 PanoLMDecoderLayer
            │   ├── attn_norm              RMSNorm
            │   ├── attn                   FLAKimiDeltaAttention w/ optional BitLinear
            │   ├── mlp_norm               RMSNorm
            │   └── mlp                    PanoLMMLP (BitLinear or nn.Linear)
            ├── norm                       RMSNorm
            └── output                     (Fused)BitLinear or nn.Linear   # causal LM head
"""
from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
from fla.layers.kda import KimiDeltaAttention as FLAKimiDeltaAttention
from fla.modules import RMSNorm as FLARMSNorm
from fla.modules.fused_bitlinear import BitLinear, FusedBitLinear
from torch.nn import RMSNorm as TorchRMSNorm
from transformers import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_panolm import PanoLMConfig


def _replace_linears(module: nn.Module, fuse_bitlinear: bool) -> None:
    """Recursively replace ``nn.Linear`` with (Fused)BitLinear in-place."""
    for name, child in module.named_children():
        if isinstance(child, nn.Linear):
            cls = FusedBitLinear if fuse_bitlinear else BitLinear
            setattr(
                module,
                name,
                cls(
                    in_features=child.in_features,
                    out_features=child.out_features,
                    bias=child.bias is not None,
                ),
            )
        else:
            _replace_linears(child, fuse_bitlinear)


class _TorchRMSNormGatedSigmoid(TorchRMSNorm):
    """Sigmoid-gated RMSNorm — matches KDA's o_norm (FusedRMSNormGated activation='sigmoid')."""

    def __init__(self, normalized_shape, eps=None, elementwise_affine=True):
        super().__init__(
            normalized_shape, eps=eps, elementwise_affine=elementwise_affine
        )
        self.register_buffer("bias", None)

    def forward(self, x: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
        return super().forward(x) * g.sigmoid()


def _make_norm(hidden_size: int, eps: float, fuse_norm: bool) -> nn.Module:
    return (
        FLARMSNorm(hidden_size, eps=eps)
        if fuse_norm
        else TorchRMSNorm(hidden_size, eps=eps)
    )


def _linear_cls(use_bitlinear: bool, fuse_bitlinear: bool) -> type[nn.Linear]:
    """Pick the projection layer type used by the MLP and LM head."""
    if not use_bitlinear:
        return nn.Linear
    return FusedBitLinear if fuse_bitlinear else BitLinear


class PanoLMMLP(nn.Module):
    """MLP that fuses gate/up projections when ``fuse_bitlinear=True``."""

    def __init__(
        self,
        hidden_size: int,
        hidden_dim: int,
        use_bitlinear: bool,
        fuse_bitlinear: bool,
    ):
        super().__init__()
        cls = _linear_cls(use_bitlinear, fuse_bitlinear)
        if fuse_bitlinear:
            self.gate_proj = cls(hidden_size, 2 * hidden_dim, bias=False)
            self.up_proj = None
        else:
            self.gate_proj = cls(hidden_size, hidden_dim, bias=False)
            self.up_proj = cls(hidden_size, hidden_dim, bias=False)
        self.down_proj = cls(hidden_dim, hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.up_proj is not None:
            gate, y = self.gate_proj(x), self.up_proj(x)
        else:
            y = self.gate_proj(x)
            gate, y = torch.tensor_split(y, 2, -1)
        return self.down_proj(F.silu(gate) * y)


def _build_attn(config: PanoLMConfig, layer_idx: int) -> nn.Module:
    """Build the fla KimiDeltaAttention module.

    The submodule names produced by fla (``q/k/v_proj``, ``q/k/v_conv1d``,
    ``f_proj``, ``g_proj``, ``b_proj``, ``o_proj``, ``o_norm``, ``A_log``,
    ``dt_bias``) match the keys the PanoLM adapter expects under
    ``model.text.layers.{i}.attn.*``.
    """
    attn = FLAKimiDeltaAttention(
        hidden_size=config.hidden_size,
        expand_v=config.expand_v,
        head_dim=config.head_dim,
        num_heads=config.num_heads,
        num_v_heads=config.num_v_heads,
        mode=config.attn_mode,
        use_short_conv=config.use_short_conv,
        allow_neg_eigval=config.allow_neg_eigval,
        safe_gate=config.safe_gate,
        lower_bound=config.lower_bound,
        conv_size=config.conv_size,
        conv_bias=config.conv_bias,
        layer_idx=layer_idx,
        norm_eps=config.rms_norm_eps,
    )
    if not config.fuse_norm:
        attn.o_norm = _TorchRMSNormGatedSigmoid(
            attn.o_norm.hidden_size,
            eps=config.rms_norm_eps,
            elementwise_affine=attn.o_norm.elementwise_affine,
        )
    if config.use_bitlinear:
        _replace_linears(attn, config.fuse_bitlinear)
    return attn


class PanoLMDecoderLayer(nn.Module):
    def __init__(self, config: PanoLMConfig, layer_idx: int):
        super().__init__()
        self.attn_norm = _make_norm(
            config.hidden_size, config.rms_norm_eps, config.fuse_norm
        )
        self.attn = _build_attn(config, layer_idx)
        self.mlp_norm = _make_norm(
            config.hidden_size, config.rms_norm_eps, config.fuse_norm
        )
        hidden_dim = config.mlp_hidden_dim
        if hidden_dim is None:
            raise ValueError(
                "PanoLMConfig.mlp_hidden_dim must be set (computed at upload time)."
            )
        self.mlp = PanoLMMLP(
            config.hidden_size,
            hidden_dim,
            config.use_bitlinear,
            config.fuse_bitlinear,
        )

    def forward(
        self,
        x: torch.Tensor,
        past_key_values=None,
        use_cache: bool = False,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        h = self.attn_norm(x)
        # KDA does not consume a hierarchical lower-bound; its channel-wise
        # vector decay subsumes that role.
        h, _new_past, _ = self.attn(
            hidden_states=h,
            past_key_values=past_key_values,
            use_cache=use_cache,
            attention_mask=attention_mask,
        )
        x = x + h
        return x + self.mlp(self.mlp_norm(x))


class PanoLMTextModel(nn.Module):
    """The full text stack including the LM head. Matches the ``model.text.*`` HF prefix."""

    def __init__(self, config: PanoLMConfig):
        super().__init__()
        extended_vocab = config.vocab_size + config.num_reserved_token_slots
        self.tok_embeddings = nn.Embedding(extended_vocab, config.hidden_size)
        if config.use_lower_bound:
            # Carried for state-dict round-trip; KDA does not consume it.
            self.lower_bounds = nn.Parameter(
                torch.zeros(config.num_hidden_layers, config.hidden_size)
            )
        self.layers = nn.ModuleList(
            [PanoLMDecoderLayer(config, i) for i in range(config.num_hidden_layers)]
        )
        self.norm = _make_norm(
            config.hidden_size, config.rms_norm_eps, config.fuse_norm
        )
        out_cls = _linear_cls(config.use_bitlinear, config.fuse_bitlinear)
        self.output = out_cls(config.hidden_size, extended_vocab, bias=False)
        self.config = config

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        past_key_values=None,
        use_cache: bool = False,
    ) -> torch.Tensor:
        if inputs_embeds is None:
            inputs_embeds = self.tok_embeddings(input_ids)
        h = inputs_embeds
        for layer in self.layers:
            h = layer(
                h,
                past_key_values=past_key_values,
                use_cache=use_cache,
                attention_mask=attention_mask,
            )
        return self.output(self.norm(h))


class _ModelTextWrapper(nn.Module):
    """Container exposing ``.text`` so state-dict keys carry the ``model.text.*`` prefix."""

    def __init__(self, config: PanoLMConfig):
        super().__init__()
        self.text = PanoLMTextModel(config)

    def forward(self, *args, **kwargs):
        return self.text(*args, **kwargs)


class PanoLMForCausalLM(PreTrainedModel, GenerationMixin):
    """PanoLM-KDA with a causal LM head, HF-compatible."""

    config_class = PanoLMConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False

    def __init__(self, config: PanoLMConfig):
        super().__init__(config)
        self.model = _ModelTextWrapper(config)
        self.post_init()

    def get_input_embeddings(self) -> nn.Module:
        return self.model.text.tok_embeddings

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.model.text.tok_embeddings = value

    def get_output_embeddings(self) -> nn.Module:
        return self.model.text.output

    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
        self.model.text.output = new_embeddings

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        past_key_values=None,
        labels: torch.LongTensor | None = None,
        use_cache: bool | None = None,
        return_dict: bool | None = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        logits = self.model(
            input_ids=input_ids,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            use_cache=bool(use_cache),
        )
        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
            )
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=past_key_values,
        )

    def _init_weights(self, module: nn.Module) -> None:
        # Wrapper is loaded from converted weights; init is a no-op beyond
        # what each submodule does on construction.
        pass