"""Pure PyTorch implementation of the XEUS E-Branchformer acoustic foundation model.

Converted to SafeTensors format (bfloat16) with zero external dependencies (no ESPnet installation required).
Original paper: XEUS - A Cross-lingual Encoder for Universal Speech (CMU WAVLab).
Weights repository: https://huggingface.co/giangndm/xeus
"""

import warnings
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file


class TransposeLayerNorm(nn.LayerNorm):
    """LayerNorm operating on channel dimension of (B, C, T) tensors."""

    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        transposed = input_tensor.transpose(-2, -1)
        normalized = F.layer_norm(
            transposed, self.normalized_shape, self.weight, self.bias, self.eps
        )
        return normalized.transpose(-2, -1)


class ConvLayerBlock(nn.Module):
    """Convolutional downsampling block used in CNNFrontend."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        stride: int,
        bias: bool = True,
    ) -> None:
        super().__init__()
        self.kernel_size = kernel_size
        self.stride = stride
        self.conv = nn.Conv1d(
            in_channels, out_channels, kernel_size=kernel_size, stride=stride, bias=bias
        )
        self.layer_norm = TransposeLayerNorm(out_channels)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.gelu(self.layer_norm(self.conv(x)))


class CNNFrontend(nn.Module):
    """Wav2Vec-style 7-layer convolutional feature extractor for 16kHz audio."""

    def __init__(self) -> None:
        super().__init__()
        shapes = [
            (512, 10, 5),
            (512, 3, 2),
            (512, 3, 2),
            (512, 3, 2),
            (512, 3, 2),
            (512, 2, 2),
            (512, 2, 2),
        ]
        in_channels = 1
        blocks = []
        for out_channels, kernel_size, stride in shapes:
            blocks.append(
                ConvLayerBlock(
                    in_channels, out_channels, kernel_size=kernel_size, stride=stride
                )
            )
            in_channels = out_channels
        self.layers = nn.Sequential(*blocks)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.ndim == 2:
            x = F.layer_norm(x, (x.shape[-1],))
            x = x.unsqueeze(1)
        first_weight = self.layers[0].conv.weight
        if x.dtype != first_weight.dtype:
            x = x.to(dtype=first_weight.dtype)
        x = self.layers(x)
        return x.transpose(1, 2)


class Preencoder(nn.Module):
    """Linear projection from frontend 512 dim to encoder 1024 dim."""

    def __init__(self, in_dim: int = 512, out_dim: int = 1024) -> None:
        super().__init__()
        self.linear_out = nn.Linear(in_dim, out_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear_out(x)


class ConvPositionalEmbedding(nn.Module):
    """Depthwise 1D convolutional relative positional embedding."""

    def __init__(self, hidden_dim: int = 1024, kernel_size: int = 31) -> None:
        super().__init__()
        conv = nn.Conv1d(
            hidden_dim,
            hidden_dim,
            kernel_size=kernel_size,
            padding=(kernel_size - 1) // 2,
            groups=16,
        )
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            conv = nn.utils.weight_norm(conv, name="weight", dim=2)
        self.convs = nn.ModuleList([conv])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = x.transpose(1, 2)
        for conv in self.convs:
            x = conv(x)
            x = F.gelu(x)
        x = x.transpose(1, 2)
        return residual + x


class CSGU(nn.Module):
    """Convolutional Spatial Gating Unit."""

    def __init__(self, size: int = 3072, kernel_size: int = 31) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(size // 2)
        self.conv = nn.Conv1d(
            size // 2,
            size // 2,
            kernel_size=kernel_size,
            padding=(kernel_size - 1) // 2,
            groups=size // 2,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        u, v = x.chunk(2, dim=-1)
        v = self.norm(v)
        v = v.transpose(1, 2)
        v = self.conv(v)
        v = v.transpose(1, 2)
        return u * v


class ConvolutionalGatingMLP(nn.Module):
    """Branch containing CSGU gated linear unit and depthwise conv."""

    def __init__(
        self, size: int = 1024, linear_units: int = 3072, kernel_size: int = 31
    ) -> None:
        super().__init__()
        self.channel_proj1 = nn.Linear(size, linear_units)
        self.csgu = CSGU(linear_units, kernel_size)
        self.channel_proj2 = nn.Linear(linear_units // 2, size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.channel_proj1(x)
        x = F.gelu(x)
        x = self.csgu(x)
        return self.channel_proj2(x)


class MultiHeadedAttention(nn.Module):
    """Multi-Head Self-Attention with projection matrices."""

    def __init__(self, n_head: int = 8, n_feat: int = 1024) -> None:
        super().__init__()
        self.d_k = n_feat // n_head
        self.h = n_head
        self.linear_q = nn.Linear(n_feat, n_feat)
        self.linear_k = nn.Linear(n_feat, n_feat)
        self.linear_v = nn.Linear(n_feat, n_feat)
        self.linear_out = nn.Linear(n_feat, n_feat)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, t, d = x.shape
        q = self.linear_q(x).view(b, t, self.h, self.d_k).transpose(1, 2)
        k = self.linear_k(x).view(b, t, self.h, self.d_k).transpose(1, 2)
        v = self.linear_v(x).view(b, t, self.h, self.d_k).transpose(1, 2)
        out = F.scaled_dot_product_attention(q, k, v)
        out = out.transpose(1, 2).contiguous().view(b, t, d)
        return self.linear_out(out)


class EBranchformerEncoderLayer(nn.Module):
    """Single E-Branchformer layer combining Attn, CGMLP and FFN branches."""

    def __init__(
        self,
        size: int = 1024,
        linear_units: int = 4096,
        cgmlp_units: int = 3072,
        heads: int = 8,
        kernel: int = 31,
    ) -> None:
        super().__init__()
        self.norm_mha = nn.LayerNorm(size)
        self.norm_mlp = nn.LayerNorm(size)
        self.mha = MultiHeadedAttention(heads, size)
        self.cgmlp = ConvolutionalGatingMLP(size, cgmlp_units, kernel)
        self.norm_merge = nn.LayerNorm(size)
        self.merge_proj = nn.Linear(size + size, size)
        self.norm_ffn = nn.LayerNorm(size)
        self.ffn = nn.Sequential(
            nn.Linear(size, linear_units),
            nn.GELU(),
            nn.Linear(linear_units, size),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Branch 1: MHA
        branch1 = self.mha(self.norm_mha(x))
        # Branch 2: CGMLP
        branch2 = self.cgmlp(self.norm_mlp(x))
        # Merge branches
        merged = torch.cat([branch1, branch2], dim=-1)
        merged = self.merge_proj(merged)
        x = x + self.norm_merge(merged)
        # FFN sublayer
        x = x + self.ffn(self.norm_ffn(x))
        return x


class EBranchformerEncoder(nn.Module):
    """Stack of E-Branchformer blocks with relative positional embeddings."""

    def __init__(
        self,
        num_blocks: int = 19,
        size: int = 1024,
        linear_units: int = 4096,
        heads: int = 8,
        kernel: int = 31,
    ) -> None:
        super().__init__()
        self.embed = nn.Sequential(ConvPositionalEmbedding(size, kernel))
        self.encoders = nn.ModuleList([
            EBranchformerEncoderLayer(size, linear_units, 3072, heads, kernel)
            for _ in range(num_blocks)
        ])
        self.after_norm = nn.LayerNorm(size)
        self.gradient_checkpointing = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.embed(x)
        for layer in self.encoders:
            if self.gradient_checkpointing and self.training:
                x = torch.utils.checkpoint.checkpoint(layer, x, use_reentrant=False)
            else:
                x = layer(x)
        return self.after_norm(x)


class XeusTrunk(nn.Module):
    """Complete XEUS acoustic foundation trunk (Frontend + Preencoder + EBranchformer)."""

    DEFAULT_REPO_ID = "giangndm/xeus"
    DEFAULT_FILENAME = "model.safetensors"

    def __init__(
        self,
        num_blocks: int = 19,
        hidden_dim: int = 1024,
        linear_units: int = 4096,
        heads: int = 8,
        kernel: int = 31,
    ) -> None:
        super().__init__()
        self.frontend = CNNFrontend()
        self.preencoder = Preencoder(512, hidden_dim)
        self.encoder = EBranchformerEncoder(
            num_blocks, hidden_dim, linear_units, heads, kernel
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.frontend(x)
        x = self.preencoder(x)
        return self.encoder(x)

    @classmethod
    def from_pretrained(
        cls,
        repo_id: str = DEFAULT_REPO_ID,
        revision: Optional[str] = None,
        filename: str = DEFAULT_FILENAME,
        dtype: Optional[torch.dtype] = torch.bfloat16,
    ) -> "XeusTrunk":
        model = cls()
        if dtype is not None:
            model = model.to(dtype=dtype)
        ckpt_path = hf_hub_download(repo_id=repo_id, revision=revision, filename=filename)
        if filename.endswith(".safetensors"):
            state_dict = load_file(ckpt_path)
        else:
            state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=True)
        filtered_state = {
            k: v
            for k, v in state_dict.items()
            if k.startswith(("frontend.", "preencoder.", "encoder."))
        }
        model.load_state_dict(filtered_state, strict=True)
        return model