import os
from pathlib import Path
from typing import List, Tuple, Union

import torch
from torch.utils.data import Dataset
from torchaudio.datasets.utils import _load_waveform

_TASKS_TO_MIXTURE = {
    "sep_clean": "mix_clean",
    "enh_single": "mix_single",
    "enh_both": "mix_both",
    "sep_noisy": "mix_both",
}


class LibriMix(Dataset):
    r"""*LibriMix* :cite:`cosentino2020librimix` dataset.

    Args:
        root (str or Path): The path where the directory ``Libri2Mix`` or
            ``Libri3Mix`` is stored. Not the path of those directories.
        subset (str, optional): The subset to use. Options: [``"train-360"``, ``"train-100"``,
            ``"dev"``, and ``"test"``] (Default: ``"train-360"``).
        num_speakers (int, optional): The number of speakers, which determines the directories
            to traverse. The Dataset will traverse ``s1`` to ``sN`` directories to collect
            N source audios. (Default: 2)
        sample_rate (int, optional): Sample rate of audio files. The ``sample_rate`` determines
            which subdirectory the audio are fetched. If any of the audio has a different sample
            rate, raises ``ValueError``. Options: [8000, 16000] (Default: 8000)
        task (str, optional): The task of LibriMix.
            Options: [``"enh_single"``, ``"enh_both"``, ``"sep_clean"``, ``"sep_noisy"``]
            (Default: ``"sep_clean"``)
        mode (str, optional): The mode when creating the mixture. If set to ``"min"``, the lengths of mixture
            and sources are the minimum length of all sources. If set to ``"max"``, the lengths of mixture and
            sources are zero padded to the maximum length of all sources.
            Options: [``"min"``, ``"max"``]
            (Default: ``"min"``)

    Note:
        The LibriMix dataset needs to be manually generated. Please check https://github.com/JorisCos/LibriMix
    """

    def __init__(
        self,
        root: Union[str, Path],
        subset: str = "train-360",
        num_speakers: int = 2,
        sample_rate: int = 8000,
        task: str = "sep_clean",
        mode: str = "min",
    ):
        self.root = Path(root) / f"Libri{num_speakers}Mix"
        if not os.path.exists(self.root):
            raise RuntimeError(
                f"The path {self.root} doesn't exist. "
                "Please check the ``root`` path and ``num_speakers`` or download the dataset manually."
            )
        if mode not in ["max", "min"]:
            raise ValueError(f'Expect ``mode`` to be one in ["min", "max"]. Found {mode}.')
        if sample_rate == 8000:
            mix_dir = self.root / "wav8k" / mode / subset
        elif sample_rate == 16000:
            mix_dir = self.root / "wav16k" / mode / subset
        else:
            raise ValueError(f"Unsupported sample rate. Found {sample_rate}.")
        self.sample_rate = sample_rate
        self.task = task

        self.mix_dir = mix_dir / _TASKS_TO_MIXTURE[task]
        if task == "enh_both":
            self.src_dirs = [(mix_dir / "mix_clean")]
        else:
            self.src_dirs = [(mix_dir / f"s{i+1}") for i in range(num_speakers)]

        self.files = [p.name for p in self.mix_dir.glob("*.wav")]
        self.files.sort()

    def _load_sample(self, key) -> Tuple[int, torch.Tensor, List[torch.Tensor]]:
        metadata = self.get_metadata(key)
        mixed = _load_waveform(self.root, metadata[1], metadata[0])
        srcs = []
        for i, path_ in enumerate(metadata[2]):
            src = _load_waveform(self.root, path_, metadata[0])
            if mixed.shape != src.shape:
                raise ValueError(f"Different waveform shapes. mixed: {mixed.shape}, src[{i}]: {src.shape}")
            srcs.append(src)
        return self.sample_rate, mixed, srcs

    def get_metadata(self, key: int) -> Tuple[int, str, List[str]]:
        """Get metadata for the n-th sample from the dataset.

        Args:
            key (int): The index of the sample to be loaded

        Returns:
            Tuple of the following items;

            int:
                Sample rate
            str:
                Path to mixed audio
            List of str:
                List of paths to source audios
        """
        filename = self.files[key]
        mixed_path = os.path.relpath(self.mix_dir / filename, self.root)
        srcs_paths = []
        for dir_ in self.src_dirs:
            src = os.path.relpath(dir_ / filename, self.root)
            srcs_paths.append(src)
        return self.sample_rate, mixed_path, srcs_paths

    def __len__(self) -> int:
        return len(self.files)

    def __getitem__(self, key: int) -> Tuple[int, torch.Tensor, List[torch.Tensor]]:
        """Load the n-th sample from the dataset.

        Args:
            key (int): The index of the sample to be loaded

        Returns:
            Tuple of the following items;

            int:
                Sample rate
            Tensor:
                Mixture waveform
            List of Tensors:
                List of source waveforms
        """
        return self._load_sample(key)