from __future__ import annotations

from typing import Any, Dict, Tuple, Union, Optional

import torch
import yaml
from torch import nn
from .heads import ISTFTHead
from .models import VocosBackbone


class Vocos(nn.Module):
    """
    The Vocos class represents a Fourier-based neural vocoder for audio synthesis.
    This class is primarily designed for inference, with support for loading from pretrained
    model checkpoints. It consists of three main components: a feature extractor,
    a backbone, and a head.
    """

    def __init__(
        self, args,
    ):
        super().__init__()
        self.backbone = VocosBackbone(
            input_channels=args.vocos.backbone.input_channels,
            dim=args.vocos.backbone.dim,
            intermediate_dim=args.vocos.backbone.intermediate_dim,
            num_layers=args.vocos.backbone.num_layers,
        )
        self.head = ISTFTHead(
            dim=args.vocos.head.dim,
            n_fft=args.vocos.head.n_fft,
            hop_length=args.vocos.head.hop_length,
            padding=args.vocos.head.padding,
        )

    def forward(self, features_input: torch.Tensor, **kwargs: Any) -> torch.Tensor:
        """
        Method to decode audio waveform from already calculated features. The features input is passed through
        the backbone and the head to reconstruct the audio output.

        Args:
            features_input (Tensor): The input tensor of features of shape (B, C, L), where B is the batch size,
                                     C denotes the feature dimension, and L is the sequence length.

        Returns:
            Tensor: The output tensor representing the reconstructed audio waveform of shape (B, T).
        """
        x = self.backbone(features_input, **kwargs)
        audio_output = self.head(x)
        return audio_output