import os
import random

from torchaudio.datasets import iemocap
from torchaudio_unittest.common_utils import get_whitenoise, save_wav, TempDirMixin, TorchaudioTestCase

LABELS = ["neu", "hap", "ang", "sad", "exc", "xxx"]
SAMPLE_RATE = 16000


def _save_wav(filepath: str, seed: int):
    wav = get_whitenoise(
        sample_rate=SAMPLE_RATE,
        duration=0.01,
        n_channels=1,
        seed=seed,
    )
    save_wav(filepath, wav, SAMPLE_RATE)
    return wav


def _save_label(label_folder: str, filename: str, wav_stem: str):
    label = random.choice(LABELS)
    if label == "exc":
        label = "hap"
    line = f"[xxx]\t{wav_stem}\t{label}\t[yyy]"
    filepath = os.path.join(label_folder, filename)

    with open(filepath, "a") as f:
        f.write(line + "\n")
    return label


def _get_samples(dataset_dir: str, session: int):
    session_folder = os.path.join(dataset_dir, f"Session{session}")
    os.makedirs(session_folder, exist_ok=True)

    wav_folder = os.path.join(session_folder, "sentences", "wav")
    label_folder = os.path.join(session_folder, "dialog", "EmoEvaluation")
    os.makedirs(wav_folder, exist_ok=True)
    os.makedirs(label_folder, exist_ok=True)

    samples = []
    wav_stems = []
    for i in range(5):
        for g in ["F", "M"]:
            speaker = f"Ses0{session}{g}"
            subfolder = f"{speaker}_impro0{i}"
            subfolder_path = os.path.join(wav_folder, subfolder)
            os.makedirs(subfolder_path, exist_ok=True)

            for j in range(5):
                wav_stem = f"{subfolder}_F00{j}"
                wav_stems.append(wav_stem)

    wav_stems = sorted(wav_stems)
    for wav_stem in wav_stems:
        subfolder = wav_stem[:-5]
        speaker = subfolder.split("_")[0]

        wav_file = os.path.join(wav_folder, subfolder, wav_stem + ".wav")
        wav = _save_wav(wav_file, seed=0)
        label = _save_label(label_folder, subfolder + ".txt", wav_stem)
        if label == "xxx":
            continue
        sample = (wav, SAMPLE_RATE, wav_stem, label, speaker)
        samples.append(sample)

    return samples


def get_mock_dataset(dataset_dir: str):
    os.makedirs(dataset_dir, exist_ok=True)

    samples = []
    for session in range(1, 4):
        samples += _get_samples(dataset_dir, session)
    return samples


class TestIemocap(TempDirMixin, TorchaudioTestCase):
    root_dir = None
    backend = "default"

    samples = []

    @classmethod
    def setUpClass(cls):
        cls.root_dir = cls.get_base_temp_dir()
        dataset_dir = os.path.join(cls.root_dir, "IEMOCAP")
        cls.samples = get_mock_dataset(dataset_dir)

    def _testIEMOCAP(self, dataset, samples):
        num_samples = 0
        for i, data in enumerate(dataset):
            self.assertEqual(data, samples[i])
            num_samples += 1

        assert num_samples == len(samples)

    def testIEMOCAPDataset(self):
        dataset = iemocap.IEMOCAP(self.root_dir)
        self._testIEMOCAP(dataset, self.samples)
