import os
from abc import ABC, abstractmethod
from functools import partial

import numpy as np
import torch
import torch.nn.functional as F
import torchaudio
from torchaudio.transforms import FrequencyMasking, TimeMasking
from torchvision.transforms import Compose
from transformers import ASTFeatureExtractor


class Transform(ABC):
    """Abstract base class for audio transformations."""

    @abstractmethod
    def __call__(self):
        """
        Abstract method to apply the transformation.

        :raises NotImplementedError: If the subclass does not implement this method.

        """
        pass


class Preprocess(ABC):
    """Abstract base class for preprocessing data.

    This class defines the interface for preprocessing data. Subclasses must implement the call method.

    """

    @abstractmethod
    def __call__(self):
        """Process the data.

        This method must be implemented by subclasses.

        :raises NotImplementedError: Subclasses must implement this method.

        """
        pass


class OneHotEncode(Transform):
    """Transform labels to one-hot encoded tensor.

    This class is a transform that takes a list of labels and returns a one-hot encoded tensor.
    The labels are converted to a tensor with one-hot encoding using the specified classes.

    :param c: A list of classes to be used for one-hot encoding.
    :type c: list
    :return: A one-hot encoded tensor.
    :rtype: torch.Tensor

    """

    def __init__(self, c: list):
        self.c = c

    def __call__(self, labels):
        """
        Transform labels to one-hot encoded tensor.

        :param labels: A list of labels to be encoded.
        :type labels: list
        :return: A one-hot encoded tensor.
        :rtype: torch.Tensor

        """

        target = torch.zeros(len(self.c), dtype=torch.float)
        for label in labels:
            idx = self.c.index(label)
            target[idx] = 1
        return target


class ParentMultilabel(Transform):
    """
    A transform that extracts a list of labels from the parent directory name of a file path.

    :param sep: The separator used to split the parent directory name into labels. Defaults to " ".
    :type sep: str
    """

    def __init__(self, sep=" "):
        self.sep = sep

    def __call__(self, path):
        """
        Extract a list of labels from the parent directory name of a file path.

        :param path: The file path from which to extract labels.
        :type path: str
        :return: A list of labels extracted from the parent directory name of the input file path.
        :rtype: List[str]
        """

        label = path.split(os.path.sep)[-2].split(self.sep)
        return label


class LabelsFromTxt(Transform):
    """
    Extract multilabel parent directory from file path.

    This class is a transform that extracts a multilabel parent directory from a file path.
    The directory names are split by a specified separator.

    :param sep: The separator used to split the directory names. Defaults to " ".
    :type sep: str

    """

    def __init__(self, delimiter=None):
        self.delimiter = delimiter

    def __call__(self, path):
        """
        Extract multilabel parent directory from file path.

        :param path: The path of the file to extract the multilabel directory from.
        :type path: str
        :return: A list of directory names representing the multilabel parent directory.
        :rtype: list

        """

        path = path.replace("wav", "txt")
        label = np.loadtxt(path, dtype=str, ndmin=1, delimiter=self.delimiter)
        return label


class PreprocessPipeline(Preprocess):
    """A preprocessing pipeline for audio data.

    This class is a preprocessing pipeline for audio data.
    The pipeline includes resampling to a target sampling rate, mixing down stereo to mono,
    and loading audio from a file.

    :param target_sr: The target sampling rate to resample to.
    :type target_sr: int
    """

    def __init__(self, target_sr):
        self.target_sr = target_sr

    def __call__(self, path):
        """
        Preprocess audio data using a pipeline.

        :param path: The path to the audio file to load.
        :type path: str
        :return: A NumPy array of preprocessed audio data.
        :rtype: numpy.ndarray

        """

        signal, sr = torchaudio.load(path)
        signal = self._resample(signal, sr)
        signal = self._mix_down(signal)
        return signal.numpy()

    def _mix_down(self, signal):
        """
        Mix down stereo to mono.

        :param signal: The audio signal to mix down.
        :type signal: torch.Tensor
        :return: The mixed down audio signal.
        :rtype: torch.Tensor

        """

        if signal.shape[0] > 1:
            signal = torch.mean(signal, dim=0, keepdim=True)
        return signal

    def _resample(self, signal, input_sr):
        """
        Resample audio signal to a target sampling rate.

        :param signal: The audio signal to resample.
        :type signal: torch.Tensor
        :param input_sr: The current sampling rate of the audio signal.
        :type input_sr: int
        :return: The resampled audio signal.
        :rtype: torch.Tensor

        """

        if input_sr != self.target_sr:
            resampler = torchaudio.transforms.Resample(input_sr, self.target_sr)
            signal = resampler(signal)
        return signal


class SpecToImage(Transform):
    def __init__(self, mean=None, std=None, eps=1e-6):
        self.mean = mean
        self.std = std
        self.eps = eps

    def __call__(self, spec):
        spec = torch.stack([spec, spec, spec], dim=-1)

        mean = torch.mean(spec) if self.mean is None else self.mean
        std = torch.std(spec) if self.std is None else self.std
        spec_norm = (spec - mean) / std

        spec_min, spec_max = torch.min(spec_norm), torch.max(spec_norm)
        spec_scaled = 255 * (spec_norm - spec_min) / (spec_max - spec_min)

        return spec_scaled.type(torch.uint8)


class MinMaxScale(Transform):
    def __call__(self, spec):
        spec_min, spec_max = torch.min(spec), torch.max(spec)

        return (spec - spec_min) / (spec_max - spec_min)


class Normalize(Transform):
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, spec):
        return (spec - self.mean) / self.std


class FeatureExtractor(Transform):
    """Extract features from audio signal using an AST feature extractor.

    This class is a transform that extracts features from an audio signal using an AST feature extractor.
    The features are returned as a PyTorch tensor.

    :param sr: The sampling rate of the audio signal.
    :type sr: int
    """

    def __init__(self, sr):
        self.transform = partial(ASTFeatureExtractor(), sampling_rate=sr, return_tensors="pt")

    def __call__(self, signal):
        """
        Extract features from audio signal using an AST feature extractor.

        :param signal: The audio signal to extract features from.
        :type signal: numpy.ndarray
        :return: A tensor of extracted audio features.
        :rtype: torch.Tensor

        """

        return self.transform(signal.squeeze()).input_values.mT


class Preemphasis(Transform):
    """perform preemphasis on the input signal.
    :param signal: The signal to filter.
    :param coeff: The preemphasis coefficient. 0 is none, default 0.97.
    :returns: the filtered signal.
    """

    def __init__(self, coeff: float = 0.97):
        self.coeff = coeff

    def __call__(self, signal):
        return torch.cat([signal[:, :1], signal[:, 1:] - self.coeff * signal[:, :-1]], dim=1)


class Spectrogram(Transform):
    def __init__(self, sample_rate, n_mels, hop_length, n_fft):
        self.transform = torchaudio.transforms.MelSpectrogram(
            sample_rate=sample_rate, n_mels=n_mels, hop_length=hop_length, n_fft=n_fft, f_min=20, center=False
        )

    def __call__(self, signal):
        return self.transform(signal)


class LogTransform(Transform):
    def __call__(self, signal):
        return torch.log(signal + 1e-8)


class PadCutToLength(Transform):
    def __init__(self, max_length):
        self.max_length = max_length

    def __call__(self, spec):
        seq_len = spec.shape[-1]

        if seq_len > self.max_length:
            return spec[..., : self.max_length]
        if seq_len < self.max_length:
            diff = self.max_length - seq_len
            return F.pad(spec, (0, diff), mode="constant", value=0)


class CustomFeatureExtractor(Transform):
    def __init__(self, sample_rate, n_mels, hop_length, n_fft, max_length, mean, std):
        self.extract = Compose(
            [
                Preemphasis(),
                Spectrogram(sample_rate=sample_rate, n_mels=n_mels, hop_length=hop_length, n_fft=n_fft),
                LogTransform(),
                PadCutToLength(max_length=max_length),
                Normalize(mean=mean, std=std),
            ]
        )

    def __call__(self, x):
        return self.extract(x)


class RepeatAudio(Transform):
    """A transform to repeat audio data.

    This class is a transform that repeats audio data a random number of times up to a maximum specified value.

    :param max_repeats: The maximum number of times to repeat the audio data.
    :type max_repeats: int
    """

    def __init__(self, max_repeats: int = 2):
        self.max_repeats = max_repeats

    def __call__(self, signal):
        """
        Repeat audio data a random number of times up to a maximum specified value.

        :param signal: The audio data to repeat.
        :type signal: numpy.ndarray
        :return: The repeated audio data.
        :rtype: numpy.ndarray

        """

        num_repeats = torch.randint(1, self.max_repeats, (1,)).item()
        return np.tile(signal, reps=num_repeats)


class MaskFrequency(Transform):
    """A transform to mask frequency of a spectrogram.

    This class is a transform that masks out a random number of consecutive frequencies from a spectrogram.

    :param max_mask_length: The maximum number of consecutive frequencies to mask out from the spectrogram.
    :type max_mask_length: int
    """

    def __init__(self, max_mask_length: int = 0):
        self.aug = FrequencyMasking(max_mask_length)

    def __call__(self, spec):
        """
        Mask out a random number of consecutive frequencies from a spectrogram.

        :param spec: The input spectrogram.
        :type spec: numpy.ndarray
        :return: The spectrogram with masked frequencies.
        :rtype: numpy.ndarray

        """

        return self.aug(spec)


class MaskTime(Transform):
    """A transform to mask time of a spectrogram.

    This class is a transform that masks out a random number of consecutive time steps from a spectrogram.

    :param max_mask_length: The maximum number of consecutive time steps to mask out from the spectrogram.
    :type max_mask_length: int
    """

    def __init__(self, max_mask_length: int = 0):
        self.aug = TimeMasking(max_mask_length)

    def __call__(self, spec):
        """
        Mask out a random number of consecutive time steps from a spectrogram.

        :param spec: The input spectrogram.
        :type spec: numpy.ndarray
        :return: The spectrogram with masked time steps.
        :rtype: numpy.ndarray

        """

        return self.aug(spec)