from typing import Tuple

import torch
import torch.nn as nn
import torchaudio


class AttPool(nn.Module):
    """Attention-Pooling module that estimates the attention score.

    Args:
        input_dim (int): Input feature dimension.
        att_dim (int): Attention Tensor dimension.
    """

    def __init__(self, input_dim: int, att_dim: int):
        super(AttPool, self).__init__()

        self.linear1 = nn.Linear(input_dim, 1)
        self.linear2 = nn.Linear(input_dim, att_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply attention and pooling.

        Args:
            x (torch.Tensor): Input Tensor with dimensions `(batch, time, feature_dim)`.

        Returns:
            (torch.Tensor): Attention score with dimensions `(batch, att_dim)`.
        """

        att = self.linear1(x)  # (batch, time, 1)
        att = att.transpose(2, 1)  # (batch, 1, time)
        att = nn.functional.softmax(att, dim=2)
        x = torch.matmul(att, x).squeeze(1)  # (batch, input_dim)
        x = self.linear2(x)  # (batch, att_dim)
        return x


class Predictor(nn.Module):
    """Prediction module that apply pooling and attention, then predict subjective metric scores.

    Args:
        input_dim (int): Input feature dimension.
        att_dim (int): Attention Tensor dimension.
    """

    def __init__(self, input_dim: int, att_dim: int):
        super(Predictor, self).__init__()
        self.att_pool_layer = AttPool(input_dim, att_dim)
        self.att_dim = att_dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Predict subjective evaluation metric score.

        Args:
            x (torch.Tensor): Input Tensor with dimensions `(batch, time, feature_dim)`.

        Returns:
            (torch.Tensor): Subjective metric score. Tensor with dimensions `(batch,)`.
        """
        x = self.att_pool_layer(x)
        x = nn.functional.softmax(x, dim=1)
        B = torch.linspace(0, 4, steps=self.att_dim, device=x.device)
        x = (x * B).sum(dim=1)
        return x


class SquimSubjective(nn.Module):
    """Speech Quality and Intelligibility Measures (SQUIM) model that predicts **subjective** metric scores
    for speech enhancement (e.g., Mean Opinion Score (MOS)). The model is adopted from *NORESQA-MOS*
    :cite:`manocha2022speech` which predicts MOS scores given the input speech and a non-matching reference.

    Args:
        ssl_model (torch.nn.Module): The self-supervised learning model for feature extraction.
        projector (torch.nn.Module): Projection layer that projects SSL feature to a lower dimension.
        predictor (torch.nn.Module): Predict the subjective scores.
    """

    def __init__(self, ssl_model: nn.Module, projector: nn.Module, predictor: nn.Module):
        super(SquimSubjective, self).__init__()
        self.ssl_model = ssl_model
        self.projector = projector
        self.predictor = predictor

    def _align_shapes(self, waveform: torch.Tensor, reference: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Cut or pad the reference Tensor to make it aligned with waveform Tensor.

        Args:
            waveform (torch.Tensor): Input waveform for evaluation. Tensor with dimensions `(batch, time)`.
            reference (torch.Tensor): Non-matching clean reference. Tensor with dimensions `(batch, time_ref)`.

        Returns:
            (torch.Tensor, torch.Tensor): The aligned waveform and reference Tensors
                with same dimensions `(batch, time)`.
        """
        T_waveform = waveform.shape[-1]
        T_reference = reference.shape[-1]
        if T_reference < T_waveform:
            num_padding = T_waveform // T_reference + 1
            reference = torch.cat([reference for _ in range(num_padding)], dim=1)
        return waveform, reference[:, :T_waveform]

    def forward(self, waveform: torch.Tensor, reference: torch.Tensor):
        """Predict subjective evaluation metric score.

        Args:
            waveform (torch.Tensor): Input waveform for evaluation. Tensor with dimensions `(batch, time)`.
            reference (torch.Tensor): Non-matching clean reference. Tensor with dimensions `(batch, time_ref)`.

        Returns:
            (torch.Tensor): Subjective metric score. Tensor with dimensions `(batch,)`.
        """
        waveform, reference = self._align_shapes(waveform, reference)
        waveform = self.projector(self.ssl_model.extract_features(waveform)[0][-1])
        reference = self.projector(self.ssl_model.extract_features(reference)[0][-1])
        concat = torch.cat((reference, waveform), dim=2)
        score_diff = self.predictor(concat)  # Score difference compared to the reference
        return 5 - score_diff


def squim_subjective_model(
    ssl_type: str,
    feat_dim: int,
    proj_dim: int,
    att_dim: int,
) -> SquimSubjective:
    """Build a custome :class:`torchaudio.prototype.models.SquimSubjective` model.

    Args:
        ssl_type (str): Type of self-supervised learning (SSL) models.
            Must be one of ["wav2vec2_base", "wav2vec2_large"].
        feat_dim (int): Feature dimension of the SSL feature representation.
        proj_dim (int): Output dimension of projection layer.
        att_dim (int): Dimension of attention scores.
    """
    ssl_model = getattr(torchaudio.models, ssl_type)()
    projector = nn.Linear(feat_dim, proj_dim)
    predictor = Predictor(proj_dim * 2, att_dim)
    return SquimSubjective(ssl_model, projector, predictor)


def squim_subjective_base() -> SquimSubjective:
    """Build :class:`torchaudio.prototype.models.SquimSubjective` model with default arguments."""
    return squim_subjective_model(
        ssl_type="wav2vec2_base",
        feat_dim=768,
        proj_dim=32,
        att_dim=5,
    )