import os
from typing import Tuple

import torchaudio
from torch import Tensor
from torch.utils.data import Dataset
from torchaudio.datasets.utils import (
    download_url,
    extract_archive,
    walk_files,
)

URL = "train-clean-100"
FOLDER_IN_ARCHIVE = "LibriTTS"
_CHECKSUMS = {
    "http://www.openslr.org/60/dev-clean.tar.gz": "0c3076c1e5245bb3f0af7d82087ee207",
    "http://www.openslr.org/60/dev-other.tar.gz": "815555d8d75995782ac3ccd7f047213d",
    "http://www.openslr.org/60/test-clean.tar.gz": "7bed3bdb047c4c197f1ad3bc412db59f",
    "http://www.openslr.org/60/test-other.tar.gz": "ae3258249472a13b5abef2a816f733e4",
    "http://www.openslr.org/60/train-clean-100.tar.gz": "4a8c202b78fe1bc0c47916a98f3a2ea8",
    "http://www.openslr.org/60/train-clean-360.tar.gz": "a84ef10ddade5fd25df69596a2767b2d",
    "http://www.openslr.org/60/train-other-500.tar.gz": "7b181dd5ace343a5f38427999684aa6f",
}


def load_libritts_item(
    fileid: str,
    path: str,
    ext_audio: str,
    ext_original_txt: str,
    ext_normalized_txt: str,
) -> Tuple[Tensor, int, str, str, int, int, str]:
    speaker_id, chapter_id, segment_id, utterance_id = fileid.split("_")
    utterance_id = fileid

    normalized_text = utterance_id + ext_normalized_txt
    normalized_text = os.path.join(path, speaker_id, chapter_id, normalized_text)

    original_text = utterance_id + ext_original_txt
    original_text = os.path.join(path, speaker_id, chapter_id, original_text)

    file_audio = utterance_id + ext_audio
    file_audio = os.path.join(path, speaker_id, chapter_id, file_audio)

    # Load audio
    waveform, sample_rate = torchaudio.load(file_audio)

    # Load original text
    with open(original_text) as ft:
        original_text = ft.readline()

    # Load normalized text
    with open(normalized_text, "r") as ft:
        normalized_text = ft.readline()

    return (
        waveform,
        sample_rate,
        original_text,
        normalized_text,
        int(speaker_id),
        int(chapter_id),
        utterance_id,
    )


class LIBRITTS(Dataset):
    """Create a Dataset for LibriTTS.

    Args:
        root (str): Path to the directory where the dataset is found or downloaded.
        url (str, optional): The URL to download the dataset from,
            or the type of the dataset to dowload.
            Allowed type values are ``"dev-clean"``, ``"dev-other"``, ``"test-clean"``,
            ``"test-other"``, ``"train-clean-100"``, ``"train-clean-360"`` and
            ``"train-other-500"``. (default: ``"train-clean-100"``)
        folder_in_archive (str, optional):
            The top-level directory of the dataset. (default: ``"LibriTTS"``)
        download (bool, optional):
            Whether to download the dataset if it is not found at root path. (default: ``False``).
    """

    _ext_original_txt = ".original.txt"
    _ext_normalized_txt = ".normalized.txt"
    _ext_audio = ".wav"

    def __init__(
        self,
        root: str,
        url: str = URL,
        folder_in_archive: str = FOLDER_IN_ARCHIVE,
        download: bool = False,
    ) -> None:

        if url in [
            "dev-clean",
            "dev-other",
            "test-clean",
            "test-other",
            "train-clean-100",
            "train-clean-360",
            "train-other-500",
        ]:

            ext_archive = ".tar.gz"
            base_url = "http://www.openslr.org/resources/60/"

            url = os.path.join(base_url, url + ext_archive)

        basename = os.path.basename(url)
        archive = os.path.join(root, basename)

        basename = basename.split(".")[0]
        folder_in_archive = os.path.join(folder_in_archive, basename)

        self._path = os.path.join(root, folder_in_archive)

        if download:
            if not os.path.isdir(self._path):
                if not os.path.isfile(archive):
                    checksum = _CHECKSUMS.get(url, None)
                    download_url(url, root, hash_value=checksum)
                extract_archive(archive)

        walker = walk_files(
            self._path, suffix=self._ext_audio, prefix=False, remove_suffix=True
        )
        self._walker = list(walker)

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

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

        Returns:
            tuple: ``(waveform, sample_rate, original_text, normalized_text, speaker_id,
            chapter_id, utterance_id)``
        """
        fileid = self._walker[n]
        return load_libritts_item(
            fileid,
            self._path,
            self._ext_audio,
            self._ext_original_txt,
            self._ext_normalized_txt,
        )

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

[((37, 22, 37, 81), 'os.path.join', 'os.path.join', ({(37, 35, 37, 39): 'path', (37, 41, 37, 51): 'speaker_id', (37, 53, 37, 63): 'chapter_id', (37, 65, 37, 80): 'normalized_text'}, {}), '(path, speaker_id, chapter_id, normalized_text)', False, 'import os\n'), ((40, 20, 40, 77), 'os.path.join', 'os.path.join', ({(40, 33, 40, 37): 'path', (40, 39, 40, 49): 'speaker_id', (40, 51, 40, 61): 'chapter_id', (40, 63, 40, 76): 'original_text'}, {}), '(path, speaker_id, chapter_id, original_text)', False, 'import os\n'), ((43, 17, 43, 71), 'os.path.join', 'os.path.join', ({(43, 30, 43, 34): 'path', (43, 36, 43, 46): 'speaker_id', (43, 48, 43, 58): 'chapter_id', (43, 60, 43, 70): 'file_audio'}, {}), '(path, speaker_id, chapter_id, file_audio)', False, 'import os\n'), ((46, 28, 46, 55), 'torchaudio.load', 'torchaudio.load', ({(46, 44, 46, 54): 'file_audio'}, {}), '(file_audio)', False, 'import torchaudio\n'), ((110, 19, 110, 40), 'os.path.basename', 'os.path.basename', ({(110, 36, 110, 39): 'url'}, {}), '(url)', False, 'import os\n'), ((111, 18, 111, 46), 'os.path.join', 'os.path.join', ({(111, 31, 111, 35): 'root', (111, 37, 111, 45): 'basename'}, {}), '(root, basename)', False, 'import os\n'), ((114, 28, 114, 69), 'os.path.join', 'os.path.join', ({(114, 41, 114, 58): 'folder_in_archive', (114, 60, 114, 68): 'basename'}, {}), '(folder_in_archive, basename)', False, 'import os\n'), ((116, 21, 116, 58), 'os.path.join', 'os.path.join', ({(116, 34, 116, 38): 'root', (116, 40, 116, 57): 'folder_in_archive'}, {}), '(root, folder_in_archive)', False, 'import os\n'), ((125, 17, 127, 9), 'torchaudio.datasets.utils.walk_files', 'walk_files', (), '', False, 'from torchaudio.datasets.utils import download_url, extract_archive, walk_files\n'), ((108, 18, 108, 59), 'os.path.join', 'os.path.join', ({(108, 31, 108, 39): 'base_url', (108, 41, 108, 58): 'url + ext_archive'}, {}), '(base_url, url + ext_archive)', False, 'import os\n'), ((119, 19, 119, 44), 'os.path.isdir', 'os.path.isdir', ({(119, 33, 119, 43): 'self._path'}, {}), '(self._path)', False, 'import os\n'), ((123, 16, 123, 40), 'torchaudio.datasets.utils.extract_archive', 'extract_archive', ({(123, 32, 123, 39): 'archive'}, {}), '(archive)', False, 'from torchaudio.datasets.utils import download_url, extract_archive, walk_files\n'), ((120, 23, 120, 46), 'os.path.isfile', 'os.path.isfile', ({(120, 38, 120, 45): 'archive'}, {}), '(archive)', False, 'import os\n'), ((122, 20, 122, 64), 'torchaudio.datasets.utils.download_url', 'download_url', (), '', False, 'from torchaudio.datasets.utils import download_url, extract_archive, walk_files\n')]