import copy
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from funasr.models.transformer.utils.nets_utils import make_pad_mask


def sense_voice_encode_forward(
    self,
    x: torch.Tensor,
    ilens: torch.Tensor = None,
    **kwargs,
):
    use_padmask = self.use_padmask
    x = F.gelu(self.conv1(x))
    x = F.gelu(self.conv2(x))
    x = x.permute(0, 2, 1)

    n_frames = x.size(1)
    max_pos = self.positional_embedding.size(0)
    max_pos = n_frames if n_frames < max_pos else max_pos
    x = (x[:, :max_pos, :] + self.positional_embedding[None, :max_pos, :]).to(x.dtype)

    if ilens is not None:
        if self.downsample_rate == 4:
            olens = (
                1
                + (ilens - self.conv1.kernel_size[0] + 2 * self.conv1.padding[0])
                // self.conv1.stride[0]
            )
        else:
            olens = ilens
        olens = (
            1
            + (olens - self.conv2.kernel_size[0] + 2 * self.conv2.padding[0])
            // self.conv2.stride[0]
        )
        olens = torch.clamp(olens, max=max_pos)
    else:
        olens = None

    if use_padmask and olens is not None:
        padding_mask = (~make_pad_mask(olens)[:, None, :]).to(x.device)
    else:
        padding_mask = None

    for layer, block in enumerate(self.blocks):
        x = block(x, mask=padding_mask, is_pad_mask=True)

    x = self.ln_post(x)

    if ilens is None:
        return x
    else:
        return x, olens