# Copyright 2024 The YourMT3 Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Please see the details in the LICENSE file.
from typing import Literal
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange


def init_layer(layer: nn.Module) -> None:
    """Initialize a Linear or Convolutional layer."""
    nn.init.xavier_uniform_(layer.weight)
    if hasattr(layer, "bias") and layer.bias is not None:
        layer.bias.data.zero_()


def init_bn(bn: nn.Module) -> None:
    """Initialize a Batchnorm layer."""
    bn.bias.data.zero_()
    bn.weight.data.fill_(1.0)
    bn.running_mean.data.zero_()
    bn.running_var.data.fill_(1.0)


def act(x: torch.Tensor, activation: str) -> torch.Tensor:
    """Activation function."""
    funcs = {"relu": F.relu_, "leaky_relu": lambda x: F.leaky_relu_(x, 0.01), "swish": lambda x: x * torch.sigmoid(x)}
    return funcs.get(activation, lambda x: Exception("Incorrect activation!"))(x)


class Res2DAVPBlock(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size, avp_kernel_size, activation):
        """Convolutional residual block modified fromr bytedance/music_source_separation."""
        super().__init__()

        padding = kernel_size[0] // 2, kernel_size[1] // 2

        self.activation = activation
        self.bn1, self.bn2 = nn.BatchNorm2d(out_channels), nn.BatchNorm2d(out_channels)

        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding, bias=False)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size, padding=padding, bias=False)

        self.is_shortcut = in_channels != out_channels
        if self.is_shortcut:
            self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=(1, 1))

        self.avp = nn.AvgPool2d(avp_kernel_size)
        self.init_weights()

    def init_weights(self):
        for m in [self.conv1, self.conv2] + ([self.shortcut] if self.is_shortcut else []):
            init_layer(m)
        for m in [self.bn1, self.bn2]:
            init_bn(m)

    def forward(self, x):
        origin = x
        x = act(self.bn1(self.conv1(x)), self.activation)
        x = self.bn2(self.conv2(x))
        x += self.shortcut(origin) if self.is_shortcut else origin
        x = act(x, self.activation)
        return self.avp(x)


class PreEncoderBlockRes3B(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size=(3, 3), avp_kernerl_size=(1, 2), activation='relu'):
        """Pre-Encoder with 3 Res2DAVPBlocks."""
        super().__init__()

        self.blocks = nn.ModuleList([
            Res2DAVPBlock(in_channels if i == 0 else out_channels, out_channels, kernel_size, avp_kernerl_size,
                          activation) for i in range(3)
        ])

    def forward(self, x):  # (B, T, F)
        x = rearrange(x, 'b t f -> b 1 t f')
        for block in self.blocks:
            x = block(x)
        return rearrange(x, 'b c t f -> b t f c')


def test_res3b():
    # mel-spec input
    x = torch.randn(2, 256, 512)  # (B, T, F)
    pre = PreEncoderBlockRes3B(in_channels=1, out_channels=128)
    x = pre(x)  # (2, 256, 64, 128): B T,F,C

    x = torch.randn(2, 110, 1024)  # (B, T, F)
    pre = PreEncoderBlockRes3B(in_channels=1, out_channels=128)
    x = pre(x)  # (2, 110, 128, 128): B,T,F,C


# ====================================================================================================================
# PreEncoderBlockHFTT: hFT-Transformer-like Pre-encoder
# ====================================================================================================================
class PreEncoderBlockHFTT(nn.Module):

    def __init__(self, margin_pre=15, margin_post=16) -> None:
        """Pre-Encoder with hFT-Transformer-like convolutions."""
        super().__init__()

        self.margin_pre, self.margin_post = margin_pre, margin_post
        self.conv = nn.Conv2d(1, 4, kernel_size=(1, 5), padding='same', padding_mode='zeros')
        self.emb_freq = nn.Linear(128, 128)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (B, T, F)
        x = rearrange(x, 'b t f -> b 1 f t')  # (B, 1, F, T) or (2, 1, 128, 110)
        x = F.pad(x, (self.margin_pre, self.margin_post), value=1e-7)  # (B, 1, F, T+margin) or (2,1,128,141)
        x = self.conv(x)  # (B, C, F, T+margin) or (2, 4, 128, 141)
        x = x.unfold(dimension=3, size=32, step=1)  # (B, c1, T, F, c2) or (2, 4, 128, 110, 32)
        x = rearrange(x, 'b c1 f t c2 -> b t f (c1 c2)')  # (B, T, F, C) or (2, 110, 128, 128)
        return self.emb_freq(x)  # (B, T, F, C) or (2, 110, 128, 128)


def test_hftt():
    # from model.spectrogram import get_spectrogram_layer_from_audio_cfg
    # from config.config import audio_cfg as default_audio_cfg
    # audio_cfg = default_audio_cfg
    # audio_cfg['codec'] = 'melspec'
    # audio_cfg['hop_length'] = 300
    # audio_cfg['n_mels'] = 128
    # x = torch.randn(2, 1, 32767)
    # mspec, _ = get_spectrogram_layer_from_audio_cfg(audio_cfg)
    # x = mspec(x)
    x = torch.randn(2, 110, 128)  # (B, T, F)
    pre_enc_hftt = PreEncoderBlockHFTT()
    y = pre_enc_hftt(x)  # (2, 110, 128, 128): B, T, F, C


# ====================================================================================================================
# PreEncoderBlockRes3BHFTT: hFT-Transformer-like Pre-encoder with Res2DAVPBlock and spec input
# ====================================================================================================================
class PreEncoderBlockRes3BHFTT(nn.Module):

    def __init__(self, margin_pre: int = 15, margin_post: int = 16) -> None:
        """Pre-Encoder with hFT-Transformer-like convolutions.
        
        Args:
            margin_pre (int): padding before the input
            margin_post (int): padding after the input
            stack_dim (Literal['c', 'f']): stack dimension. channel or frequency

        """
        super().__init__()
        self.margin_pre, self.margin_post = margin_pre, margin_post
        self.res3b = PreEncoderBlockRes3B(in_channels=1, out_channels=4)
        self.emb_freq = nn.Linear(128, 128)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (B, T, F) or (2, 110, 1024), input spectrogram
        x = rearrange(x, 'b t f -> b f t')  # (2, 1024, 110): B,F,T
        x = F.pad(x, (self.margin_pre, self.margin_post), value=1e-7)  # (2, 1024, 141): B,F,T+margin
        x = rearrange(x, 'b f t -> b t f')  # (2, 141, 1024): B,T+margin,F
        x = self.res3b(x)  # (2, 141, 128, 4): B,T+margin,F,C
        x = x.unfold(dimension=1, size=32, step=1)  # (B, T, F, C1, C2) or (2, 110, 128, 4, 32)
        x = rearrange(x, 'b t f c1 c2 -> b t f (c1 c2)')  # (B, T, F, C) or (2, 110, 128, 128)
        return self.emb_freq(x)  # (B, T, F, C) or (2, 110, 128, 128)


def test_res3b_hftt():
    # from model.spectrogram import get_spectrogram_layer_from_audio_cfg
    # from config.config import audio_cfg as default_audio_cfg
    # audio_cfg = default_audio_cfg
    # audio_cfg['codec'] = 'spec'
    # audio_cfg['hop_length'] = 300
    # x = torch.randn(2, 1, 32767)
    # spec, _ = get_spectrogram_layer_from_audio_cfg(audio_cfg)
    # x = spec(x)  # (2, 110, 1024): B,T,F
    x = torch.randn(2, 110, 1024)  # (B, T, F)
    pre_enc_res3b_hftt = PreEncoderBlockRes3BHFTT()
    y = pre_enc_res3b_hftt(x)  # (2, 110, 128, 128): B, T, F, C


# # ====================================================================================================================
# # PreEncoderBlockConv1D: Pre-encoder without activation, with Melspec input
# # ====================================================================================================================
# class PreEncoderBlockConv1D(nn.Module):

#     def __init__(self,
#                  in_channels,
#                  out_channels,
#                  kernel_size=3) -> None:
#         """Pre-Encoder with 1D convolution."""
#         super().__init__()
#         self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=1)
#         self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size, stride=1)

#     def forward(self, x: torch.Tensor) -> torch.Tensor:
#         # x: (B, T, F) or (2, 128, 256), input melspec
#         x = rearrange(x, 'b t f -> b f t')  # (2, 256, 128): B,F,T
#         x = self.conv1(x)  # (2, 128, 128): B,F,T
#         return rearrange(x, 'b f t -> b t f')  # (2, 110, 128): B,T,F

# def test_conv1d():
#     # from model.spectrogram import get_spectrogram_layer_from_audio_cfg
#     # from config.config import audio_cfg as default_audio_cfg
#     # audio_cfg = default_audio_cfg
#     # audio_cfg['codec'] = 'melspec'
#     # audio_cfg['hop_length'] = 256
#     # audio_cfg['n_mels'] = 512
#     # x = torch.randn(2, 1, 32767)
#     # mspec, _ = get_spectrogram_layer_from_audio_cfg(audio_cfg)
#     # x = mspec(x)
#     x = torch.randn(2, 128, 128)  # (B, T, F)
#     pre_enc_conv1d = PreEncoderBlockConv1D(in_channels=1, out_channels=128)
#     y = pre_enc_conv1d(x)  # (2, 110, 128, 128): B, T, F, C