import torch
import torchaudio
from torch import nn
from torch.nn import Conv1d, ConvTranspose1d
from torch.nn import functional as F
from torch.nn.utils.parametrizations import weight_norm
from torch.nn.utils.parametrize import remove_parametrizations

from .......common.utilities import load_fsspec

LRELU_SLOPE = 0.1


def get_padding(k, d):
    """Calculate padding size for a convolutional layer.

    Args:
        k (int): Kernel size.
        d (int): Dilation rate.

    Returns:
        int: Required padding size.
    """
    return int((k * d - d) / 2)


class ResBlock1(torch.nn.Module):
    """Residual Block Type 1. It has 3 convolutional layers in each convolutional block.

    Network::

        x -> lrelu -> conv1_1 -> conv1_2 -> conv1_3 -> z -> lrelu -> conv2_1 -> conv2_2 -> conv2_3 -> o -> + -> o
        |--------------------------------------------------------------------------------------------------|


    Args:
        channels (int): number of hidden channels for the convolutional layers.
        kernel_size (int): size of the convolution filter in each layer.
        dilations (list): list of dilation value for each conv layer in a block.
    """

    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
        super().__init__()
        self.convs1 = nn.ModuleList(
            [
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[i],
                        padding=get_padding(kernel_size, dilation[i]),
                    )
                )
                for i in range(3)
            ]
        )

        self.convs2 = nn.ModuleList(
            [
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=1,
                        padding=get_padding(kernel_size, 1),
                    )
                )
                for _ in range(3)
            ]
        )

    def forward(self, x):
        """
        Args:
            x (Tensor): input tensor.
        Returns:
            Tensor: output tensor.
        Shapes:
            x: [B, C, T]
        """
        for c1, c2 in zip(self.convs1, self.convs2):
            xt = F.leaky_relu(x, LRELU_SLOPE)
            xt = c1(xt)
            xt = F.leaky_relu(xt, LRELU_SLOPE)
            xt = c2(xt)
            x = xt + x
        return x

    def remove_weight_norm(self):
        for l in self.convs1:
            remove_parametrizations(l, "weight")
        for l in self.convs2:
            remove_parametrizations(l, "weight")


class ResBlock2(torch.nn.Module):
    """Residual Block Type 2. It has 1 convolutional layers in each convolutional block.

    Network::

        x -> lrelu -> conv1-> -> z -> lrelu -> conv2-> o -> + -> o
        |---------------------------------------------------|


    Args:
        channels (int): number of hidden channels for the convolutional layers.
        kernel_size (int): size of the convolution filter in each layer.
        dilations (list): list of dilation value for each conv layer in a block.
    """

    def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
        super().__init__()
        self.convs = nn.ModuleList(
            [
                weight_norm(
                    Conv1d(
                        channels,
                        channels,
                        kernel_size,
                        1,
                        dilation=dilation[i],
                        padding=get_padding(kernel_size, dilation[i]),
                    )
                )
                for i in range(2)
            ]
        )

    def forward(self, x):
        for c in self.convs:
            xt = F.leaky_relu(x, LRELU_SLOPE)
            xt = c(xt)
            x = xt + x
        return x

    def remove_weight_norm(self):
        for l in self.convs:
            remove_parametrizations(l, "weight")


class HifiganGenerator(torch.nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        resblock_type,
        resblock_dilation_sizes,
        resblock_kernel_sizes,
        upsample_kernel_sizes,
        upsample_initial_channel,
        upsample_factors,
        inference_padding=5,
        cond_channels=0,
        conv_pre_weight_norm=True,
        conv_post_weight_norm=True,
        conv_post_bias=True,
        cond_in_each_up_layer=False,
    ):
        r"""HiFiGAN Generator with Multi-Receptive Field Fusion (MRF)

        Network:
            x -> lrelu -> upsampling_layer -> resblock1_k1x1 -> z1 -> + -> z_sum / #resblocks -> lrelu -> conv_post_7x1 -> tanh -> o
                                                 ..          -> zI ---|
                                              resblockN_kNx1 -> zN ---'

        Args:
            in_channels (int): number of input tensor channels.
            out_channels (int): number of output tensor channels.
            resblock_type (str): type of the `ResBlock`. '1' or '2'.
            resblock_dilation_sizes (List[List[int]]): list of dilation values in each layer of a `ResBlock`.
            resblock_kernel_sizes (List[int]): list of kernel sizes for each `ResBlock`.
            upsample_kernel_sizes (List[int]): list of kernel sizes for each transposed convolution.
            upsample_initial_channel (int): number of channels for the first upsampling layer. This is divided by 2
                for each consecutive upsampling layer.
            upsample_factors (List[int]): upsampling factors (stride) for each upsampling layer.
            inference_padding (int): constant padding applied to the input at inference time. Defaults to 5.
        """
        super().__init__()
        self.inference_padding = inference_padding
        self.num_kernels = len(resblock_kernel_sizes)
        self.num_upsamples = len(upsample_factors)
        self.cond_in_each_up_layer = cond_in_each_up_layer

        # initial upsampling layers
        self.conv_pre = weight_norm(Conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3))
        resblock = ResBlock1 if resblock_type == "1" else ResBlock2
        # upsampling layers
        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(upsample_factors, upsample_kernel_sizes)):
            self.ups.append(
                weight_norm(
                    ConvTranspose1d(
                        upsample_initial_channel // (2**i),
                        upsample_initial_channel // (2 ** (i + 1)),
                        k,
                        u,
                        padding=(k - u) // 2,
                    )
                )
            )
        # MRF blocks
        self.resblocks = nn.ModuleList()
        for i in range(len(self.ups)):
            ch = upsample_initial_channel // (2 ** (i + 1))
            for k, d in zip(resblock_kernel_sizes, resblock_dilation_sizes):
                self.resblocks.append(resblock(ch, k, d))
        # post convolution layer
        self.conv_post = weight_norm(Conv1d(ch, out_channels, 7, 1, padding=3, bias=conv_post_bias))
        if cond_channels > 0:
            self.cond_layer = nn.Conv1d(cond_channels, upsample_initial_channel, 1)

        if not conv_pre_weight_norm:
            remove_parametrizations(self.conv_pre, "weight")

        if not conv_post_weight_norm:
            remove_parametrizations(self.conv_post, "weight")

        if self.cond_in_each_up_layer:
            self.conds = nn.ModuleList()
            for i in range(len(self.ups)):
                ch = upsample_initial_channel // (2 ** (i + 1))
                self.conds.append(nn.Conv1d(cond_channels, ch, 1))

    def forward(self, x, g=None):
        """
        Args:
            x (Tensor): feature input tensor.
            g (Tensor): global conditioning input tensor.

        Returns:
            Tensor: output waveform.

        Shapes:
            x: [B, C, T]
            Tensor: [B, 1, T]
        """
        with torch.no_grad():
            with torch.amp.autocast('cuda'):
                x = self.conv_pre(x).unsqueeze(0)
                if hasattr(self, "cond_layer"):
                    x.add_(self.cond_layer(g))
                for i in range(self.num_upsamples):
                    x = F.leaky_relu(x, LRELU_SLOPE, inplace=True)
                    x = self.ups[i](x)

                    if self.cond_in_each_up_layer:
                        x.add_(self.conds[i](g))

                    z_sum = 0
                    for j in range(self.num_kernels):
                        z_sum += (self.resblocks[i * self.num_kernels + j](x)).float()
                    x = z_sum / self.num_kernels
                x = F.leaky_relu(x, inplace=True)
                x = self.conv_post(x)
                x = torch.tanh(x)
                return x

    @torch.no_grad()
    def inference(self, c):
        """
        Args:
            x (Tensor): conditioning input tensor.

        Returns:
            Tensor: output waveform.

        Shapes:
            x: [B, C, T]
            Tensor: [B, 1, T]
        """
        c = c.to(self.conv_pre.weight.device)
        c = torch.nn.functional.pad(c, (self.inference_padding, self.inference_padding), "replicate")
        return self.forward(c)

    def remove_weight_norm(self):
        print("Removing weight norm...")
        for l in self.ups:
            remove_parametrizations(l, "weight")
        for l in self.resblocks:
            l.remove_weight_norm()
        remove_parametrizations(self.conv_pre, "weight")
        remove_parametrizations(self.conv_post, "weight")

    def load_checkpoint(self, config, checkpoint_path, eval=False, cache=False):
        state = torch.load(checkpoint_path, map_location=torch.device("cpu"))
        self.load_state_dict(state["model"])
        if eval:
            self.eval()
            assert not self.training
            self.remove_weight_norm()

    def estimate_receptive_field(self):
        """
        Estimate the receptive field of the model based on its configuration.

        Steps:
        1. Start from the initial conv (conv_pre) with kernel=7 (no dilation):
           receptive_field = 7
        2. For each upsampling stage:
           - Multiply the current receptive field by the upsampling factor (since time length is scaled).
           - Add the receptive field contribution from the associated MRF blocks at this scale.

        The MRF block receptive field is calculated by summing the receptive fields of all resblocks at that scale.
        Each resblock adds its own receptive field based on the dilations and kernel sizes.

        Note: This is a heuristic estimation assuming that the resblocks are sequentially affecting the receptive field.
        """

        # Start from conv_pre: kernel_size=7, dilation=1
        # Receptive field increment = (7 - 1) * 1 = 6
        # Since we talk about total receptive field size, let's consider receptive_field as the number of frames.
        # Here we can say receptive_field = 7 (the number of frames covered by kernel=7)
        receptive_field = 7

        idx = 0
        for i, up_factor in enumerate(self.upsample_factors):
            # After upsampling, the receptive field scales
            receptive_field = receptive_field * up_factor

            # Now add the contribution of the MRF blocks at this scale
            # We have num_kernels blocks per scale
            scale_rf = 0
            for j in range(self.num_kernels):
                block = self.resblocks[idx]
                idx += 1
                scale_rf = max(scale_rf, block.receptive_field())  # Take max since they are parallel and merged

            # The MRF blocks process after upsampling, so we add that receptive field increment.
            # Since these blocks are in series at the same scale (averaged), we approximate by adding them.
            # Actually, they are parallel and then averaged. The effective RF should consider the largest block.
            # We'll consider just the largest one for a safer overestimate.
            receptive_field += scale_rf

        return receptive_field


class SELayer(nn.Module):
    """Squeeze-and-Excitation layer for channel-wise attention.
    
    This layer implements the Squeeze-and-Excitation mechanism that adaptively
    recalibrates channel-wise feature responses by explicitly modeling
    interdependencies between channels.
    """

    def __init__(self, channel, reduction=8):
        """Initialize SE layer.

        Args:
            channel (int): Number of input channels.
            reduction (int, optional): Channel reduction factor. Defaults to 8.
        """
        super(SELayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel),
            nn.Sigmoid(),
        )

    def forward(self, x):
        """Apply channel-wise attention.

        Args:
            x (torch.Tensor): Input tensor of shape [B, C, T].

        Returns:
            torch.Tensor: Channel-wise scaled tensor.
        """
        y = self.avg_pool(x).view(x.size(0), x.size(1))
        y = self.fc(y).view(x.size(0), x.size(1), 1, 1)
        return x * y


class SEBasicBlock(nn.Module):
    """Basic ResNet block with Squeeze-and-Excitation.
    
    This block combines residual connections with SE attention for improved
    feature extraction.
    """

    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None, reduction=8):
        """Initialize SE-ResNet block.

        Args:
            inplanes (int): Number of input channels.
            planes (int): Number of output channels.
            stride (int, optional): Stride for convolution. Defaults to 1.
            downsample (nn.Module, optional): Downsampling layer. Defaults to None.
            reduction (int, optional): SE reduction factor. Defaults to 8.
        """
        super(SEBasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.se = SELayer(planes, reduction)
        self.downsample = downsample

    def forward(self, x):
        """Process input through SE-ResNet block.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Processed tensor.
        """
        residual = x

        x = self.conv1(x)
        x = self.relu(x)
        x = self.bn1(x)

        x = self.conv2(x)
        x = self.bn2(x)
        x = self.se(x)

        if self.downsample is not None:
            residual = self.downsample(residual)

        x += residual
        x = self.relu(x)
        return x


def set_init_dict(model_dict, checkpoint_state, c):
    # Partial initialization: if there is a mismatch with new and old layer, it is skipped.
    for k, v in checkpoint_state.items():
        if k not in model_dict:
            print(f" | > Layer missing in the model definition: {k}")
    pretrained_dict = {k: v for k, v in checkpoint_state.items() if k in model_dict}
    # 2. filter out different size layers
    pretrained_dict = {k: v for k, v in pretrained_dict.items() if v.numel() == model_dict[k].numel()}
    # 3. skip reinit layers
    if c.has("reinit_layers") and c.reinit_layers is not None:
        for reinit_layer_name in c.reinit_layers:
            pretrained_dict = {k: v for k, v in pretrained_dict.items() if reinit_layer_name not in k}
    # 4. overwrite entries in the existing state dict
    model_dict.update(pretrained_dict)
    print(f" | > {len(pretrained_dict)} / {len(model_dict)} layers are restored.")
    return model_dict


class PreEmphasis(nn.Module):
    """Pre-emphasis filter for audio processing.
    
    Applies pre-emphasis filtering to the input audio signal to enhance
    high-frequency components.
    """

    def __init__(self, coefficient=0.97):
        """Initialize pre-emphasis filter.

        Args:
            coefficient (float, optional): Pre-emphasis coefficient. Defaults to 0.97.
        """
        super().__init__()
        self.coefficient = coefficient
        self.register_buffer("filter", torch.tensor([-self.coefficient, 1.0], dtype=torch.float32).view(1, 1, -1))

    def forward(self, x):
        """Apply pre-emphasis filtering.

        Args:
            x (torch.Tensor): Input audio tensor.

        Returns:
            torch.Tensor: Pre-emphasized audio.
        """
        assert len(x.size()) == 2

        x = torch.nn.functional.pad(x.unsqueeze(1), (1, 0), "reflect")
        x = torch.nn.functional.conv1d(x, self.filter).squeeze(1)
        return x


class ResNetSpeakerEncoder(nn.Module):
    """ResNet-based speaker encoder for voice conversion.
    
    This module extracts speaker embeddings from audio using a modified ResNet
    architecture with optional attentive statistical pooling.
    """

    # pylint: disable=W0102
    def __init__(
        self,
        input_dim=64,
        proj_dim=512,
        layers=[3, 4, 6, 3],
        num_filters=[32, 64, 128, 256],
        encoder_type="ASP",
        log_input=False,
        use_torch_spec=False,
        audio_config=None,
    ):
        """Initialize speaker encoder.

        Args:
            input_dim (int, optional): Input feature dimension. Defaults to 64.
            proj_dim (int, optional): Projection dimension. Defaults to 512.
            layers (List[int], optional): Number of layers in each block. Defaults to [3,4,6,3].
            num_filters (List[int], optional): Number of filters in each block. Defaults to [32,64,128,256].
            encoder_type (str, optional): Type of encoder ("ASP" or "SAP"). Defaults to "ASP".
            log_input (bool, optional): Whether to apply log to input. Defaults to False.
            use_torch_spec (bool, optional): Whether to use torch spectrogram. Defaults to False.
            audio_config (dict, optional): Audio processing configuration. Defaults to None.
        """
        super(ResNetSpeakerEncoder, self).__init__()

        self.encoder_type = encoder_type
        self.input_dim = input_dim
        self.log_input = log_input
        self.use_torch_spec = use_torch_spec
        self.audio_config = audio_config
        self.proj_dim = proj_dim

        self.conv1 = nn.Conv2d(1, num_filters[0], kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU(inplace=True)
        self.bn1 = nn.BatchNorm2d(num_filters[0])

        self.inplanes = num_filters[0]
        self.layer1 = self.create_layer(SEBasicBlock, num_filters[0], layers[0])
        self.layer2 = self.create_layer(SEBasicBlock, num_filters[1], layers[1], stride=(2, 2))
        self.layer3 = self.create_layer(SEBasicBlock, num_filters[2], layers[2], stride=(2, 2))
        self.layer4 = self.create_layer(SEBasicBlock, num_filters[3], layers[3], stride=(2, 2))

        self.instancenorm = nn.InstanceNorm1d(input_dim)

        if self.use_torch_spec:
            self.torch_spec = torch.nn.Sequential(
                PreEmphasis(audio_config["preemphasis"]),
                torchaudio.transforms.MelSpectrogram(
                    sample_rate=audio_config["sample_rate"],
                    n_fft=audio_config["fft_size"],
                    win_length=audio_config["win_length"],
                    hop_length=audio_config["hop_length"],
                    window_fn=torch.hamming_window,
                    n_mels=audio_config["num_mels"],
                ),
            )

        else:
            self.torch_spec = None

        outmap_size = int(self.input_dim / 8)

        self.attention = nn.Sequential(
            nn.Conv1d(num_filters[3] * outmap_size, 128, kernel_size=1),
            nn.ReLU(inplace=True),
            nn.BatchNorm1d(128),
            nn.Conv1d(128, num_filters[3] * outmap_size, kernel_size=1),
            nn.Softmax(dim=2),
        )

        if self.encoder_type == "SAP":
            out_dim = num_filters[3] * outmap_size
        elif self.encoder_type == "ASP":
            out_dim = num_filters[3] * outmap_size * 2
        else:
            raise ValueError("Undefined encoder")

        self.fc = nn.Linear(out_dim, proj_dim)

        self._init_layers()

    def _init_layers(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def create_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = [block(self.inplanes, planes, stride, downsample)]
        self.inplanes = planes * block.expansion
        layers.extend(block(self.inplanes, planes) for _ in range(1, blocks))

        return nn.Sequential(*layers)

    # pylint: disable=R0201
    def new_parameter(self, *size):
        out = nn.Parameter(torch.FloatTensor(*size))
        nn.init.xavier_normal_(out)
        return out

    def forward(self, x, l2_norm=False):
        """Extract speaker embeddings from input features.

        Args:
            x (torch.Tensor): Input features.
            l2_norm (bool, optional): Whether to apply L2 normalization. Defaults to False.

        Returns:
            torch.Tensor: Speaker embeddings.
        """
        x.squeeze_(1)
        # if you torch spec compute it otherwise use the mel spec computed by the AP
        if self.use_torch_spec:
            x = self.torch_spec(x)

        if self.log_input:
            x.add_(1e-6).log_()
        x = self.instancenorm(x).unsqueeze(1)

        x = self.conv1(x)
        x = self.relu(x)
        x = self.bn1(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = x.reshape(x.size(0), -1, x.size(-1))

        w = self.attention(x)

        if self.encoder_type == "SAP":
            x = torch.sum(x * w, dim=2)
        elif self.encoder_type == "ASP":
            mu = torch.sum(x * w, dim=2)
            sg = torch.sqrt((torch.sum((x ** 2) * w, dim=2) - mu ** 2).clamp(min=1e-5))
            x = torch.cat((mu, sg), 1)

        x = x.view(x.size()[0], -1)
        x = self.fc(x)

        if l2_norm:
            x = torch.nn.functional.normalize(x, p=2, dim=1)
        return x

    def load_checkpoint(
        self,
        checkpoint_path: str,
        eval: bool = False,
        use_cuda: bool = False,
        criterion=None,
        cache=False,
    ):
        state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"), cache=cache)
        try:
            self.load_state_dict(state["model"])
            print(" > Model fully restored. ")
        except (KeyError, RuntimeError) as error:
            # If eval raise the error
            if eval:
                raise error

            print(" > Partial model initialization.")
            model_dict = self.state_dict()
            model_dict = set_init_dict(model_dict, state["model"])
            self.load_state_dict(model_dict)
            del model_dict

        # load the criterion for restore_path
        if criterion is not None and "criterion" in state:
            try:
                criterion.load_state_dict(state["criterion"])
            except (KeyError, RuntimeError) as error:
                print(" > Criterion load ignored because of:", error)

        if use_cuda:
            self.cuda()
            if criterion is not None:
                criterion = criterion.cuda()

        if eval:
            self.eval()
            assert not self.training

        if not eval:
            return criterion, state["step"]
        return criterion


class HifiDecoder(torch.nn.Module):
    """HiFi-GAN based decoder for high-quality speech synthesis.
    
    This module converts mel-spectrograms or other acoustic features into
    high-fidelity waveforms using a HiFi-GAN architecture with optional
    speaker conditioning.
    """

    def __init__(
        self,
        input_sample_rate=22050,
        output_sample_rate=24000,
        output_hop_length=256,
        ar_mel_length_compression=1024,
        decoder_input_dim=1024,
        resblock_type_decoder="1",
        resblock_dilation_sizes_decoder=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
        resblock_kernel_sizes_decoder=[3, 7, 11],
        upsample_rates_decoder=[8, 8, 2, 2],
        upsample_initial_channel_decoder=512,
        upsample_kernel_sizes_decoder=[16, 16, 4, 4],
        d_vector_dim=512,
        cond_d_vector_in_each_upsampling_layer=True,
        speaker_encoder_audio_config={
            "fft_size": 512,
            "win_length": 400,
            "hop_length": 160,
            "sample_rate": 16000,
            "preemphasis": 0.97,
            "num_mels": 64,
        },
    ):
        """Initialize HiFi decoder.

        Args:
            input_sample_rate (int, optional): Input sampling rate. Defaults to 22050.
            output_sample_rate (int, optional): Output sampling rate. Defaults to 24000.
            output_hop_length (int, optional): Output hop length. Defaults to 256.
            ar_mel_length_compression (int, optional): Autoregressive compression factor. Defaults to 1024.
            decoder_input_dim (int, optional): Input dimension for decoder. Defaults to 1024.
            resblock_type_decoder (str, optional): Type of residual blocks. Defaults to "1".
            resblock_dilation_sizes_decoder (List[List[int]], optional): Dilation sizes for residual blocks.
            resblock_kernel_sizes_decoder (List[int], optional): Kernel sizes for residual blocks.
            upsample_rates_decoder (List[int], optional): Upsampling rates.
            upsample_initial_channel_decoder (int, optional): Initial number of channels.
            upsample_kernel_sizes_decoder (List[int], optional): Kernel sizes for upsampling.
            d_vector_dim (int, optional): Speaker embedding dimension. Defaults to 512.
            cond_d_vector_in_each_upsampling_layer (bool, optional): Whether to condition each layer.
            speaker_encoder_audio_config (dict, optional): Speaker encoder configuration.
        """
        super().__init__()
        self.input_sample_rate = input_sample_rate
        self.output_sample_rate = output_sample_rate
        self.output_hop_length = output_hop_length
        self.ar_mel_length_compression = ar_mel_length_compression
        self.speaker_encoder_audio_config = speaker_encoder_audio_config
        self.waveform_decoder = HifiganGenerator(
            decoder_input_dim,
            1,
            resblock_type_decoder,
            resblock_dilation_sizes_decoder,
            resblock_kernel_sizes_decoder,
            upsample_kernel_sizes_decoder,
            upsample_initial_channel_decoder,
            upsample_rates_decoder,
            inference_padding=0,
            cond_channels=d_vector_dim,
            conv_pre_weight_norm=False,
            conv_post_weight_norm=False,
            conv_post_bias=False,
            cond_in_each_up_layer=cond_d_vector_in_each_upsampling_layer,
        )
        self.speaker_encoder = ResNetSpeakerEncoder(
            input_dim=64,
            proj_dim=512,
            log_input=True,
            use_torch_spec=True,
            audio_config=speaker_encoder_audio_config,
        )

    @property
    def device(self):
        return next(self.parameters()).device

    def forward(self, latents, g=None):
        """Generate waveform from latent features.

        Args:
            latents (torch.Tensor): Input latent features.
            g (torch.Tensor, optional): Speaker embedding. Defaults to None.

        Returns:
            torch.Tensor: Generated waveform.
        """

        z = torch.nn.functional.interpolate(
            latents.transpose(1, 2),
            scale_factor=self.ar_mel_length_compression / self.output_hop_length,
            mode="linear",
            align_corners=False,
        ).squeeze(1)
        # upsample to the right sr
        if self.output_sample_rate != self.input_sample_rate:
            z = torch.nn.functional.interpolate(
                z,
                scale_factor=self.output_sample_rate / self.input_sample_rate,
                mode="linear",
                align_corners=False,
            ).squeeze(0)
        o = self.waveform_decoder(z, g=g)
        return o

    @torch.no_grad()
    def inference(self, c, g):
        """Generate waveform in inference mode.

        Args:
            c (torch.Tensor): Input features.
            g (torch.Tensor): Speaker embedding.

        Returns:
            torch.Tensor: Generated waveform.
        """
        return self.forward(c, g=g)

    def load_checkpoint(self, checkpoint_path, eval=False):  # pylint: disable=unused-argument, redefined-builtin
        """Load model checkpoint.

        Args:
            checkpoint_path (str): Path to checkpoint file.
            eval (bool, optional): Whether to set model to eval mode. Defaults to False.

        Returns:
            dict: Loaded checkpoint state.
        """
        state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"))
        # remove unused keys
        state = state["model"]
        for key in list(state.keys()):
            if "waveform_decoder." not in key and "speaker_encoder." not in key:
                del state[key]

        self.load_state_dict(state)
        if eval:
            self.eval()
            assert not self.training
            self.waveform_decoder.remove_weight_norm()
