from pathlib import Path
import unittest

from torch import Tensor
import torchaudio
from voicefixer import Vocoder

from training.datasets.hifi_libri_dataset import HifiLibriDataset, HifiLibriItem


class TestHifiLibriDataset(unittest.TestCase):
    def setUp(self):
        self.cache_dir = "datasets_cache"
        self.dataset = HifiLibriDataset(cache_dir=self.cache_dir, cache=True)
        self.vocoder_vf = Vocoder(44100)

    def test_init(self):
        self.assertEqual(len(self.dataset.cutset), 129751)

    def test_get_cache_subdir_path(self):
        idx = 1234
        expected_path = Path(self.cache_dir) / "cache-hifitts-librittsr" / "2000"
        self.assertEqual(self.dataset.get_cache_subdir_path(idx), expected_path)

    def test_get_cache_file_path(self):
        idx = 1234
        expected_path = (
            Path(self.cache_dir) / "cache-hifitts-librittsr" / "2000" / f"{idx}.pt"
        )
        self.assertEqual(self.dataset.get_cache_file_path(idx), expected_path)

    def test_getitem(self):
        # Take the hifi items from the beginning of the dataset
        item = self.dataset[0]
        self.assertIsInstance(item, HifiLibriItem)
        self.assertEqual(item.dataset_type, "hifitts")

        # Convert mel spectrogram to waveform and save it to a file
        # NOTE: Vocoder expects the mel spectrogram to be prepared in a specific way
        # wav = self.vocoder_vf.forward(item.mel.permute((1, 0)).unsqueeze(0))
        # wav_path = Path(f"results/{item.id}.wav")

        # torchaudio.save(str(wav_path), wav, 44100)

        # Check that the cache file is created
        cache_file = self.dataset.get_cache_file_path(0)
        self.assertTrue(cache_file.exists())
        # Take the same id again to check if the cache is used
        item = self.dataset[0]
        self.assertIsInstance(item, HifiLibriItem)
        self.assertEqual(item.dataset_type, "hifitts")

        item = self.dataset[10]
        self.assertIsInstance(item, HifiLibriItem)
        self.assertEqual(item.dataset_type, "hifitts")
        # Check that the cache file is created
        cache_file = self.dataset.get_cache_file_path(10)
        self.assertTrue(cache_file.exists())

        item = self.dataset[20]
        self.assertIsInstance(item, HifiLibriItem)
        self.assertEqual(item.dataset_type, "hifitts")

        # Take the libri items from the end of the dataset
        item = self.dataset[len(self.dataset) - 20]
        self.assertIsInstance(item, HifiLibriItem)
        self.assertEqual(item.dataset_type, "libritts")
        # Check that the cache file is created
        cache_file = self.dataset.get_cache_file_path(len(self.dataset) - 20)
        self.assertTrue(cache_file.exists())

        item = self.dataset[len(self.dataset) - 10]
        self.assertIsInstance(item, HifiLibriItem)
        self.assertEqual(item.dataset_type, "libritts")
        item = self.dataset[len(self.dataset) - 5]
        self.assertIsInstance(item, HifiLibriItem)
        self.assertEqual(item.dataset_type, "libritts")

    def test_collate_fn(self):
        data = [self.dataset[0] for _ in range(10)]
        collated = self.dataset.collate_fn(data)
        self.assertIsInstance(collated, list)
        self.assertIsInstance(collated[0], list)  # ids
        self.assertIsInstance(collated[1], list)  # raw_texts
        self.assertIsInstance(collated[2], Tensor)  # speakers
        self.assertIsInstance(collated[3], Tensor)  # texts
        self.assertIsInstance(collated[4], Tensor)  # src_lens
        self.assertIsInstance(collated[5], Tensor)  # mels
        self.assertIsInstance(collated[6], Tensor)  # pitches
        self.assertIsInstance(collated[7], list)  # pitches_stat
        self.assertIsInstance(collated[8], Tensor)  # mel_lens
        self.assertIsInstance(collated[9], Tensor)  # langs
        self.assertIsInstance(collated[10], Tensor)  # attn_priors
        self.assertIsInstance(collated[11], Tensor)  # wavs
        self.assertIsInstance(collated[12], Tensor)  # energy

    def test_include_libri(self):
        dataset_with_libri = HifiLibriDataset(
            cache_dir="datasets_cache",
            include_libri=True,
        )
        dataset_without_libri = HifiLibriDataset(
            cache_dir="datasets_cache",
            include_libri=False,
        )

        # Check that the dataset with LibriTTS is larger than the dataset without LibriTTS
        self.assertTrue(len(dataset_with_libri) > len(dataset_without_libri))

        # Check that the dataset with LibriTTS includes items of type 'libritts'
        libri_item = dataset_with_libri[len(dataset_with_libri) - 10]
        self.assertIsInstance(libri_item, HifiLibriItem)
        self.assertEqual(libri_item.dataset_type, "libritts")

        # Check that the dataset without LibriTTS does not include items of type 'libritts'
        hifi_item = dataset_without_libri[len(dataset_without_libri) - 10]
        self.assertIsInstance(hifi_item, HifiLibriItem)
        self.assertEqual(hifi_item.dataset_type, "hifitts")

    def test_dur_filter(self):
        # Test with a duration of 0.2
        self.assertFalse(self.dataset.dur_filter(0.2))
        # Test with a duration of 1.0
        self.assertTrue(self.dataset.dur_filter(1.0))
        # Test with a duration of 2.0
        self.assertTrue(self.dataset.dur_filter(2.0))
        # Test with a duration of 30.0
        self.assertFalse(self.dataset.dur_filter(30.0))


if __name__ == "__main__":
    unittest.main()