"""
Zipformer-inspired encoder (arXiv:2310.11230): a U-Net-shaped stack that
downsamples the middle layers in time and upsamples back before the output
heads, with Conformer-style macaron-FF + self-attention + conv blocks.

Built from standard PyTorch primitives, not a byte-exact port of icefall's
custom ops. 
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional
from src.scaling import (
    ConvolutionModule,
    FeedForwardModule,
    PositionalEncoding,
    SelfAttentionModule,
    conv_subsample_length,
)


def make_pad_mask(lengths: torch.Tensor, max_len: int) -> torch.Tensor:
    """Returns (B, max_len) bool tensor, True at padded positions."""
    batch_size = lengths.size(0)
    seq_range = torch.arange(0, max_len, device=lengths.device).unsqueeze(0).expand(batch_size, -1)
    return seq_range >= lengths.unsqueeze(1)


class Conv2dSubsampling(nn.Module):
    """4x time subsampling: two stride-2 conv2d layers turn (T, n_mels) into a
    d_model sequence at 1/4 the frame rate (10ms -> 40ms per frame).
    """

    def __init__(self, n_mels: int, d_model: int, conv_channels: int = 128, dropout: float = 0.1):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, conv_channels, kernel_size=3, stride=2, padding=1),
            nn.SiLU(),
            nn.Conv2d(conv_channels, conv_channels, kernel_size=3, stride=2, padding=1),
            nn.SiLU(),
        )
        freq_out = n_mels
        for _ in range(2):
            freq_out = (freq_out + 2 * 1 - 3) // 2 + 1
        self.out_proj = nn.Linear(conv_channels * freq_out, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, feats: torch.Tensor, lengths: torch.Tensor):
        # feats: (B, T, n_mels)
        x = feats.unsqueeze(1)  # (B, 1, T, F)
        x = self.conv(x)  # (B, C, T', F')
        b, c, t, f = x.shape
        x = x.permute(0, 2, 1, 3).reshape(b, t, c * f)  # (B, T', C*F')
        x = self.out_proj(x)
        x = self.dropout(x)
        out_lengths = conv_subsample_length(lengths, num_layers=2, kernel_size=3, stride=2, padding=1)
        out_lengths = out_lengths.clamp(min=0, max=t)
        return x, out_lengths


class DownsampleConcat(nn.Module):
    """Groups every `factor` consecutive frames and projects down to d_model.
    Unlike a strided conv, no input frame's information is dropped.
    """

    def __init__(self, d_model: int, factor: int):
        super().__init__()
        self.factor = factor
        self.proj = nn.Linear(d_model * factor, d_model)

    def forward(self, x: torch.Tensor, lengths: torch.Tensor):
        b, t, c = x.shape
        pad_len = (self.factor - t % self.factor) % self.factor
        if pad_len > 0:
            x = F.pad(x, (0, 0, 0, pad_len))
        t_padded = t + pad_len
        x = x.reshape(b, t_padded // self.factor, self.factor * c)
        x = self.proj(x)
        new_lengths = torch.div(lengths + self.factor - 1, self.factor, rounding_mode="floor")
        return x, new_lengths


class UpsampleRepeat(nn.Module):
    """Inverse of DownsampleConcat: repeats each frame `factor` times, projects
    back to d_model, then trims/pads to match the skip connection's length.
    """

    def __init__(self, d_model: int, factor: int):
        super().__init__()
        self.factor = factor
        self.proj = nn.Linear(d_model, d_model)

    def forward(self, x: torch.Tensor, target_time_len: int):
        x = x.repeat_interleave(self.factor, dim=1)
        x = self.proj(x)
        cur_len = x.size(1)
        if cur_len > target_time_len:
            x = x[:, :target_time_len, :]
        elif cur_len < target_time_len:
            x = F.pad(x, (0, 0, 0, target_time_len - cur_len))
        return x


class ZipformerBlock(nn.Module):
    """One encoder block: macaron FF -> self-attn -> conv -> macaron FF -> norm."""

    def __init__(self, d_model: int, nhead: int, d_ff: int, conv_kernel: int, dropout: float):
        super().__init__()
        self.ff1 = FeedForwardModule(d_model, d_ff, dropout)
        self.self_attn = SelfAttentionModule(d_model, nhead, dropout)
        self.conv = ConvolutionModule(d_model, conv_kernel, dropout)
        self.ff2 = FeedForwardModule(d_model, d_ff, dropout)
        self.final_norm = nn.LayerNorm(d_model)

    def forward(self, x: torch.Tensor, pad_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        x = self.ff1(x)
        x = self.self_attn(x, pad_mask=pad_mask)
        x = self.conv(x, pad_mask=pad_mask)
        x = self.ff2(x)
        x = self.final_norm(x)
        return x


class Stage(nn.Module):
    def __init__(self, num_layers: int, d_model: int, nhead: int, d_ff: int, conv_kernel: int, dropout: float):
        super().__init__()
        self.blocks = nn.ModuleList(
            [ZipformerBlock(d_model, nhead, d_ff, conv_kernel, dropout) for _ in range(num_layers)]
        )

    def forward(self, x: torch.Tensor, pad_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        for block in self.blocks:
            x = block(x, pad_mask=pad_mask)
        return x


class ZipformerEncoder(nn.Module):
    """5-stage U-Net encoder: full -> half -> quarter (bottleneck) -> half ->
    full res, with additive skip connections bridging matching resolutions.
    """

    def __init__(
        self,
        n_mels: int = 80,
        d_model: int = 256,
        conv_channels: int = 128,
        nhead: int = 4,
        d_ff: int = 1024,
        conv_kernel: int = 15,
        stage_layers=(2, 3, 4, 3, 2),
        downsample_factors=(2, 2),
        dropout: float = 0.1,
    ):
        super().__init__()
        assert len(stage_layers) == 5, "expected 5 stages: full, half, bottleneck, half, full"
        assert len(downsample_factors) == 2

        self.frontend = Conv2dSubsampling(n_mels, d_model, conv_channels, dropout)
        self.pos_enc = PositionalEncoding(d_model, dropout=dropout)

        self.stage0 = Stage(stage_layers[0], d_model, nhead, d_ff, conv_kernel, dropout)
        self.down0 = DownsampleConcat(d_model, downsample_factors[0])
        self.stage1 = Stage(stage_layers[1], d_model, nhead, d_ff, conv_kernel, dropout)
        self.down1 = DownsampleConcat(d_model, downsample_factors[1])
        self.stage2 = Stage(stage_layers[2], d_model, nhead, d_ff, conv_kernel, dropout)
        self.up1 = UpsampleRepeat(d_model, downsample_factors[1])
        self.stage3 = Stage(stage_layers[3], d_model, nhead, d_ff, conv_kernel, dropout)
        self.up0 = UpsampleRepeat(d_model, downsample_factors[0])
        self.stage4 = Stage(stage_layers[4], d_model, nhead, d_ff, conv_kernel, dropout)

        self.final_norm = nn.LayerNorm(d_model)

    def forward(self, feats: torch.Tensor, lengths: torch.Tensor):
        """
        feats: (B, T, n_mels) log-mel features
        lengths: (B,) valid frame counts before padding
        returns: (encoder_out (B, T0, d_model), out_lengths (B,))
        """
        x, lengths0 = self.frontend(feats, lengths)
        x = self.pos_enc(x)
        t0 = x.size(1)
        pad_mask0 = make_pad_mask(lengths0, t0)

        x = self.stage0(x, pad_mask0)
        skip0 = x

        x, lengths1 = self.down0(x, lengths0)
        t1 = x.size(1)
        pad_mask1 = make_pad_mask(lengths1, t1)
        x = self.stage1(x, pad_mask1)
        skip1 = x

        x, lengths2 = self.down1(x, lengths1)
        t2 = x.size(1)
        pad_mask2 = make_pad_mask(lengths2, t2)
        x = self.stage2(x, pad_mask2)

        x = self.up1(x, target_time_len=t1)
        x = x + skip1
        x = self.stage3(x, pad_mask1)

        x = self.up0(x, target_time_len=t0)
        x = x + skip0
        x = self.stage4(x, pad_mask0)

        x = self.final_norm(x)
        return x, lengths0