import importlib
import numpy as np
import io
import os
import posixpath
import random
import re
import subprocess
import time
import torch
import torchaudio
import webdataset as wds
import pytorch_lightning as pl

from aeiou.core import is_silence
from os import path
from pedalboard.io import AudioFile
from torchaudio import transforms as T
from typing import Optional, Callable, List

from .utils import Stereo, Mono, PhaseFlipper, PadCrop_Normalized_T
from suno_utils.utils.text import read_jsonl

AUDIO_KEYS = ("flac", "wav", "mp3", "m4a", "ogg", "opus")

# fast_scandir implementation by Scott Hawley originally in https://github.com/zqevans/audio-diffusion/blob/main/dataset/dataset.py


def fast_scandir(
    dir: str,  # top-level directory at which to begin scanning
    ext: list,  # list of allowed file extensions,
    # max_size = 1 * 1000 * 1000 * 1000 # Only files < 1 GB
):
    "very fast `glob` alternative. from https://stackoverflow.com/a/59803793/4259243"
    subfolders, files = [], []
    ext = [
        "." + x if x[0] != "." else x for x in ext
    ]  # add starting period to extensions if needed
    try:  # hope to avoid 'permission denied' by this try
        for f in os.scandir(dir):
            try:  # 'hope to avoid too many levels of symbolic links' error
                if f.is_dir():
                    subfolders.append(f.path)
                elif f.is_file():
                    file_ext = os.path.splitext(f.name)[1].lower()
                    is_hidden = os.path.basename(f.path).startswith(".")

                    if file_ext in ext and not is_hidden:
                        files.append(f.path)
            except:
                pass
    except:
        pass

    for dir in list(subfolders):
        sf, f = fast_scandir(dir, ext)
        subfolders.extend(sf)
        files.extend(f)
    return subfolders, files


def keyword_scandir(
    dir: str,  # top-level directory at which to begin scanning
    ext: list,  # list of allowed file extensions
    keywords: list,  # list of keywords to search for in the file name
):
    "very fast `glob` alternative. from https://stackoverflow.com/a/59803793/4259243"
    subfolders, files = [], []
    # make keywords case insensitive
    keywords = [keyword.lower() for keyword in keywords]
    # add starting period to extensions if needed
    ext = ["." + x if x[0] != "." else x for x in ext]
    banned_words = ["paxheader", "__macosx"]
    try:  # hope to avoid 'permission denied' by this try
        for f in os.scandir(dir):
            try:  # 'hope to avoid too many levels of symbolic links' error
                if f.is_dir():
                    subfolders.append(f.path)
                elif f.is_file():
                    is_hidden = f.name.split("/")[-1][0] == "."
                    has_ext = os.path.splitext(f.name)[1].lower() in ext
                    name_lower = f.name.lower()
                    has_keyword = any([keyword in name_lower for keyword in keywords])
                    has_banned = any(
                        [banned_word in name_lower for banned_word in banned_words]
                    )
                    if (
                        has_ext
                        and has_keyword
                        and not has_banned
                        and not is_hidden
                        and not os.path.basename(f.path).startswith("._")
                    ):
                        files.append(f.path)
            except:
                pass
    except:
        pass

    for dir in list(subfolders):
        sf, f = keyword_scandir(dir, ext, keywords)
        subfolders.extend(sf)
        files.extend(f)
    return subfolders, files


def get_audio_filenames(
    paths: list,  # directories in which to search
    keywords=None,
    exts=[".wav", ".mp3", ".flac", ".ogg", ".aif", ".opus"],
):
    "recursively get a list of audio filenames"
    filenames = []
    if type(paths) is str:
        paths = [paths]
    for path in paths:  # get a list of relevant filenames
        if keywords is not None:
            subfolders, files = keyword_scandir(path, exts, keywords)
        else:
            subfolders, files = fast_scandir(path, exts)
        filenames.extend(files)
    return filenames


# globals
RATE_HZ = 25
SEMANTIC_CODEBOOK_SIZE = 4000
SEMANTIC_N_CODEBOOKS = 1
CODEC_CODEBOOK_SIZE = 2048
CODEC_N_CODEBOOKS = 12
VAE_DIM = 128
N_TOKENS_MEMMAP = 250


class DenoisingMemmapDataset(torch.utils.data.IterableDataset):
    def __init__(
        self,
        input_memmap_path: str,
        corrupt_memmap_path: str,
        rate_hz: int = 100,
        codebook_size: int = 2048,
        n_codebooks: int = 12,
        t_memmap: int = 1000,
        embed_dim: int = 128,
        custom_metadata_fn: Optional[Callable[[str], str]] = None,
    ):
        super().__init__()
        self.filenames = []
        self.rate_hz = rate_hz
        self.codebook_size = codebook_size
        self.n_codebooks = n_codebooks
        self.t_memmap = t_memmap
        self.custom_metadata_fn = custom_metadata_fn

        # load the input and corrupt memmaps
        input_data = np.memmap(input_memmap_path, dtype=np.uint16, mode="r")
        input_data = input_data.reshape(-1, t_memmap, n_codebooks)
        self.input_data = input_data

        corrupt_data = np.memmap(corrupt_memmap_path, dtype=np.uint16, mode="r")
        corrupt_data = corrupt_data.reshape(-1, t_memmap, n_codebooks)
        self.corrupt_data = corrupt_data

        print(
            "input_data", self.input_data.shape, "corrupt_data", self.corrupt_data.shape
        )
        assert self.corrupt_data.shape[0] == self.input_data.shape[0]

    def __len__(self):
        return self.input_data.shape[0]

    def __iter__(self):
        while True:
            rand_idx = np.random.randint(0, self.input_data.shape[0] - 1)
            start_time = time.time()

            info = {}

            t_start = 0
            t_end = 10.0
            seconds_start = 0
            seconds_total = t_end

            # input_codes are the targets and corrupt_codes are the conditioning
            input_codes = torch.from_numpy(self.input_data[rand_idx, ...].copy()).long()
            corrupt_codes = torch.from_numpy(
                self.corrupt_data[rand_idx, ...].copy()
            ).long()

            info["codec_codes"] = corrupt_codes  # these are used as conditioning
            info["idx"] = rand_idx
            info["timestamps"] = (t_start, t_end)
            info["seconds_start"] = seconds_start
            info["seconds_total"] = seconds_total

            end_time = time.time()

            info["load_time"] = end_time - start_time

            # disabling this interface for now and hardcoding conditioning
            # this means we do not do any conditioning dropout rn
            # if self.custom_metadata_fn is not None:
            #    custom_metadata = self.custom_metadata_fn(info, semantic_codes)
            #    info.update(custom_metadata)
            #
            #    if "__reject__" in info and info["__reject__"]:
            #        return self[random.randrange(len(self))]

            # vae_embed = vae_embed.permute(1, 0)  # channels, seq_len

            yield (input_codes.permute(1, 0), info)


class SemanticToCodecMemmapDataset(torch.utils.data.IterableDataset):
    def __init__(
        self,
        semantic_memmap_path: str,
        codec_memmap_path: str,
        custom_metadata_fn: Optional[Callable[[str], str]] = None,
    ):
        super().__init__()
        self.custom_metadata_fn = custom_metadata_fn

        semantic_data = np.memmap(semantic_memmap_path, dtype=np.uint16, mode="r")
        semantic_data = semantic_data.reshape(-1, N_TOKENS_MEMMAP)
        self.semantic_data = semantic_data
        print("semantic_data", self.semantic_data.shape)

        codec_data = np.memmap(codec_memmap_path, dtype=np.uint16, mode="r")
        codec_data = codec_data.reshape(-1, N_TOKENS_MEMMAP, CODEC_N_CODEBOOKS)
        self.codec_data = codec_data
        print("codec_data", self.codec_data.shape)

        assert (
            codec_data.shape[0] == semantic_data.shape[0]
        )  # must have same number of rows

        print(f"Found {semantic_data.shape[0]} examples.")

    def __len__(self):
        return self.semantic_data.shape[0]

    def __iter__(self):
        while True:
            rand_idx = np.random.randint(0, self.semantic_data.shape[0] - 1)
            start_time = time.time()

            info = {}
            codec_codes = torch.from_numpy(self.codec_data[rand_idx, ...].copy()).long()

            semantic_codes = torch.from_numpy(
                self.semantic_data[rand_idx, ...].copy()
            ).long()
            info["semantic_codes"] = semantic_codes

            info["idx"] = rand_idx
            end_time = time.time()
            info["load_time"] = end_time - start_time

            codec_codes = codec_codes.permute(1, 0)  # n_codebooks, seq_len

            yield (codec_codes, info)


class SemanticToAutoencoderMemmapDataset(torch.utils.data.IterableDataset):
    def __init__(
        self,
        semantic_memmap_path: str,
        vae_memmap_path: str,
        custom_metadata_fn: Optional[Callable[[str], str]] = None,
    ):
        super().__init__()
        self.custom_metadata_fn = custom_metadata_fn

        semantic_data = np.memmap(semantic_memmap_path, dtype=np.uint16, mode="r")
        semantic_data = semantic_data.reshape(-1, N_TOKENS_MEMMAP)
        self.semantic_data = semantic_data
        print("semantic_data", self.semantic_data.shape)

        vae_data = np.memmap(vae_memmap_path, dtype=np.float32, mode="r")
        vae_data = vae_data.reshape(-1, N_TOKENS_MEMMAP, VAE_DIM)
        self.vae_data = vae_data
        print("vae_data", self.vae_data.shape)

        assert (
            vae_data.shape[0] == semantic_data.shape[0]
        )  # must have same number of rows

        print(f"Found {semantic_data.shape[0]} examples.")

    def __len__(self):
        return self.semantic_data.shape[0]

    def __iter__(self):
        while True:
            rand_idx = np.random.randint(0, self.semantic_data.shape[0] - 1)
            start_time = time.time()

            info = {}
            vae_embeds = torch.from_numpy(self.vae_data[rand_idx, ...].copy()).float()

            semantic_codes = torch.from_numpy(
                self.semantic_data[rand_idx, ...].copy()
            ).long()
            info["semantic_codes"] = semantic_codes

            info["idx"] = rand_idx
            end_time = time.time()
            info["load_time"] = end_time - start_time

            vae_embeds = vae_embeds.permute(1, 0)  # channels, seq_len

            yield (vae_embeds, info)


class DiscreteVAEtoVAEMemmapDataset(torch.utils.data.IterableDataset):
    def __init__(
        self,
        discrete_memmap_path: str,
        vae_memmap_path: str,
        vae_dim: int = 128,
        n_tokens_memmap: int = 1000,
        custom_metadata_fn: Optional[Callable[[str], str]] = None,
    ):
        super().__init__()
        self.custom_metadata_fn = custom_metadata_fn

        discrete_data = np.memmap(discrete_memmap_path, dtype=np.uint16, mode="r")
        discrete_data = discrete_data.reshape(-1, n_tokens_memmap)
        self.discrete_data = discrete_data
        print("discrete_data", self.discrete_data.shape)

        vae_data = np.memmap(vae_memmap_path, dtype=np.float32, mode="r")
        vae_data = vae_data.reshape(-1, vae_dim, n_tokens_memmap)
        self.vae_data = vae_data
        print("vae_data", self.vae_data.shape)

        assert (
            vae_data.shape[0] == discrete_data.shape[0]
        )  # must have same number of rows

        print(f"Found {discrete_data.shape[0]} examples.")

    def __len__(self):
        return self.discrete_data.shape[0]

    def __iter__(self):
        while True:
            rand_idx = np.random.randint(0, self.discrete_data.shape[0] - 1)
            start_time = time.time()

            info = {}
            vae_embeds = torch.from_numpy(self.vae_data[rand_idx, ...].copy()).float()

            discrete_codes = torch.from_numpy(
                self.discrete_data[rand_idx, ...].copy()
            ).long()
            info["discrete_codes"] = discrete_codes

            info["idx"] = rand_idx
            end_time = time.time()
            info["load_time"] = end_time - start_time

            # vae_embeds = vae_embeds.permute(1, 0)  # channels, seq_len

            yield (vae_embeds, info)


class VAEMemmapDataset(torch.utils.data.IterableDataset):
    def __init__(
        self,
        vae_memmap_path: str,
        vae_metas_path: Optional[str] = None,
        vae_dim: int = 128,
        n_tokens_memmap: int = 1000,
        custom_metadata_fn: Optional[Callable[[str], str]] = None,
    ):
        """For use in training unconditional diffusion model.

        When a metas file is provided, the metadata is loaded and returned with the data.
        This can be used for lyric conditioning, etc.

        """
        super().__init__()
        self.custom_metadata_fn = custom_metadata_fn
        self.vae_metas_path = vae_metas_path
        self.vae_dim = vae_dim
        self.n_tokens_memmap = n_tokens_memmap

        vae_data = np.memmap(vae_memmap_path, dtype=np.float32, mode="r")
        vae_data = vae_data.reshape(-1, vae_dim, n_tokens_memmap)
        self.vae_data = vae_data
        print(f"Found {vae_data.shape[0]} examples.")

        if vae_metas_path is not None:
            self.metas = read_jsonl(vae_metas_path)
            assert len(self.metas) == self.vae_data.shape[0]
            print("Loaded metadata for", len(self.metas), "examples.")
        else:
            self.metas = None

    def __len__(self):
        return self.vae_data.shape[0]

    def __iter__(self):
        while True:
            rand_idx = np.random.randint(0, self.vae_data.shape[0] - 1)
            start_time = time.time()

            info = {}
            vae_embeds = torch.from_numpy(self.vae_data[rand_idx, ...].copy()).float()

            info["idx"] = rand_idx
            end_time = time.time()
            info["load_time"] = end_time - start_time
            info["seconds_start"] = 0
            info["seconds_total"] = 10.0

            if self.metas is not None:
                info["lyrics"] = self.metas[rand_idx]["lyrics"]
                # info["title"] = self.metas[rand_idx]["title"]

            # vae_embeds = vae_embeds.permute(1, 0)  # channels, seq_len

            yield (vae_embeds, info)


class VAEMemmapMapDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        vae_memmap_path: str,
        vae_metas_path: Optional[str] = None,
        vae_dim: int = 128,
        n_tokens_memmap: int = 1000,
        custom_metadata_fn: Optional[Callable[[str], str]] = None,
    ):
        """For use in training unconditional diffusion model.

        When a metas file is provided, the metadata is loaded and returned with the data.
        This can be used for lyric conditioning, etc.

        """
        super().__init__()
        self.custom_metadata_fn = custom_metadata_fn
        self.vae_metas_path = vae_metas_path
        self.vae_dim = vae_dim
        self.n_tokens_memmap = n_tokens_memmap

        vae_data = np.memmap(vae_memmap_path, dtype=np.float32, mode="r")
        vae_data = vae_data.reshape(-1, vae_dim, n_tokens_memmap)
        self.vae_data = vae_data
        print(f"Found {vae_data.shape[0]} examples.")

        if vae_metas_path is not None:
            self.metas = read_jsonl(vae_metas_path)
            assert len(self.metas) == self.vae_data.shape[0]
            print("Loaded metadata for", len(self.metas), "examples.")
        else:
            self.metas = None

    def __len__(self):
        return self.vae_data.shape[0]

    def __getitem__(self, idx):
        start_time = time.time()

        info = {}
        vae_embeds = torch.from_numpy(self.vae_data[idx, ...].copy()).float()

        info["idx"] = idx
        end_time = time.time()
        info["load_time"] = end_time - start_time
        info["seconds_start"] = 0
        info["seconds_total"] = 10.0

        if self.metas is not None:
            info["lyrics"] = self.metas[idx]["lyrics"]
            # info["title"] = self.metas[rand_idx]["title"]

        # vae_embeds = vae_embeds.permute(1, 0)  # channels, seq_len

        return (vae_embeds, info)


class DummyMemmapMapDataset(torch.utils.data.Dataset):
    """Dummy dataset for testing throughput."""

    def __init__(
        self,
        num_examples: int = 1000000,
        vae_dim: int = 128,
        vae_n_tokens: int = 3000,
        semantic_n_tokens: int = 750,
    ):
        super().__init__()
        self.num_examples = num_examples
        self.vae_dim = vae_dim
        self.vae_n_tokens = vae_n_tokens
        self.semantic_n_tokens = semantic_n_tokens

    def __len__(self):
        return self.num_examples

    def __getitem__(self, idx):
        # channels, seq_len
        vae_embeds = torch.randn(self.vae_dim, self.vae_n_tokens)
        semantic_codes = torch.randint(0, 4000, (self.semantic_n_tokens,)).long()

        info = {}
        info["idx"] = idx
        info["semantic_codes"] = semantic_codes
        info["lyrics"] = "Country roads, take me home, to the place, where I belong"
        info["tags"] = "country, folk, bluegrass"
        info["tags_and_lyrics"] = (info["tags"], info["lyrics"])

        print("vae_embeds", vae_embeds.shape, "semantic_codes", semantic_codes.shape)

        return (vae_embeds, info)


class ContextMemmapMapDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        dataset_dir: str,
        vae_memmap_filename: str = "data_vae_val.bin",
        semantic_memmap_filename: str = "data_semantic_val.bin",
        metas_filename: str = "metas_prev_val.jsonl",
        vae_dim: int = 128,
        vae_rate_hz: int = 100,
        vae_n_tokens: int = 3000,
        vae_use_float16: bool = False,
        vae_pad_token: torch.Tensor = torch.zeros(128),
        semantic_rate_hz: int = 25,
        semantic_n_tokens: int = 750,
    ):
        """For use in training conditional diffusion model.

        When a metas file is provided, the metadata is loaded and returned with the data.
        This can be used for lyric and tags conditioning.

        """
        super().__init__()
        self.dataset_dir = dataset_dir
        self.vae_dim = vae_dim
        self.vae_memmap_filename = vae_memmap_filename
        self.semantic_memmap_filename = semantic_memmap_filename
        self.vae_n_tokens = vae_n_tokens
        self.semantic_n_tokens = semantic_n_tokens
        self.metas_filename = metas_filename
        self.vae_use_float16 = vae_use_float16
        self.vae_pad_token = vae_pad_token
        self.vae_rate_hz = vae_rate_hz
        self.semantic_rate_hz = semantic_rate_hz

        self.metas = read_jsonl(
            os.path.join(dataset_dir, metas_filename), progress=False
        )
        print(
            "Loaded metadata for",
            len(self.metas),
            "examples from",
            os.path.join(dataset_dir, metas_filename),
        )

        assert (
            "prev_context_id" in self.metas[0]
        )  # check that the metas contain the field "prev_context_id"

        # open vae memmap
        if vae_use_float16:
            vae_data = np.memmap(
                os.path.join(dataset_dir, vae_memmap_filename),
                dtype=np.float16,
                mode="r",
            )
        else:
            vae_data = np.memmap(
                os.path.join(dataset_dir, vae_memmap_filename),
                dtype=np.float32,
                mode="r",
            )

        vae_data = vae_data.reshape(-1, vae_n_tokens, vae_dim)
        self.vae_data = vae_data
        print(
            "vae_data",
            os.path.join(dataset_dir, vae_memmap_filename),
            self.vae_data.shape,
        )

        # open semantic memmap
        semantic_data = np.memmap(
            os.path.join(dataset_dir, semantic_memmap_filename),
            dtype=np.uint16,
            mode="r",
        )
        semantic_data = semantic_data.reshape(-1, semantic_n_tokens, 1)
        self.semantic_data = semantic_data[:, :, 0]
        print(
            "semantic_data",
            os.path.join(dataset_dir, semantic_memmap_filename),
            self.semantic_data.shape,
        )

        assert (
            vae_data.shape[0] == semantic_data.shape[0]
        )  # must have same number of rows

        print(f"Found {vae_data.shape[0]} examples.")

        assert len(self.metas) == self.vae_data.shape[0]

    def __len__(self):
        return self.vae_data.shape[0]

    def __getitem__(self, idx):
        info = {}
        info["idx"] = idx
        info["seconds_start"] = 0
        info["seconds_total"] = 30.0

        # VAE embeds
        input_vae_embeds = torch.from_numpy(self.vae_data[idx, ...].copy()).float()
        input_vae_embeds = input_vae_embeds.permute(1, 0)  # channels, seq_len

        # semantic codes
        input_semantic_codes = torch.from_numpy(
            self.semantic_data[idx, ...].copy()
        ).long()
        info["semantic_codes"] = input_semantic_codes

        # setup latent context
        # in this case we just use the first prev context id
        # if len(self.metas[idx]["prev_context_id"]) > 0:
        #    context_vae_embeds = torch.from_numpy(
        #        self.vae_data[self.metas[idx]["prev_context_id"][0], ...].copy()
        #    ).float()
        #    context_vae_embeds = context_vae_embeds.permute(1, 0)  # channels, seq_len
        # else:
        #    context_vae_embeds = None

        # check the start time of the this example
        start_s = self.metas[idx]["start_s"]

        # check if the start_s is 0
        # if not, then we need to find the previous meta
        context_vae_embeds = None
        if float(start_s) != 0.0:
            prev_meta_idx = idx - 1
            if prev_meta_idx >= 0:
                prev_meta = self.metas[prev_meta_idx]
                if prev_meta["end_s"] == self.metas[idx]["start_s"]:
                    # then we have previous context
                    context_vae_embeds = torch.from_numpy(
                        self.vae_data[prev_meta_idx, ...].copy()
                    ).float()
                    context_vae_embeds = context_vae_embeds.permute(
                        1, 0
                    )  # channels, seq_len

        # setup lyrics
        lyrics = self.metas[idx].get("text", "")
        lyrics_aligned = self.metas[idx].get("text_aligned", "")

        info["latent_context"] = context_vae_embeds
        info["tags"] = self.metas[idx].get("tags", [])
        info["phonemes"] = self.metas[idx].get("phonemized_text", "")
        if info["phonemes"] is None:
            info["phonemes"] = ""

        # first check if we have lyrics aligned
        if lyrics_aligned != "":
            # 10% of the time we swap to the full lyrics
            if random.random() < 0.1:
                info["lyrics"] = lyrics
            # 90% of the time we use the aligned lyrics
            else:
                info["lyrics"] = lyrics_aligned
        # if we don't have aligned lyrics, we use the full lyrics
        else:
            info["lyrics"] = lyrics

        info["tags_and_lyrics"] = (info["tags"], info["lyrics"])

        # padding mask is passed for compatibility, but shouldn't be used
        padding_mask = torch.ones(input_vae_embeds.shape[-1])
        info["padding_mask"] = padding_mask.bool()

        return (input_vae_embeds, info)


class GeneralMemmapMapDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        dataset_dir: str,
        vae_memmap_filename: str = "data_vae_val.bin",
        semantic_memmap_filename: str = "data_semantic_val.bin",
        metas_filename: str = "metas_val.jsonl",
        vae_dim: int = 128,
        vae_rate_hz: int = 100,
        vae_n_tokens: int = 3000,
        vae_use_n_tokens: int = 3000,
        vae_use_float16: bool = False,
        semantic_rate_hz: int = 25,
        semantic_n_tokens: int = 750,
        semantic_use_n_tokens: int = 750,
    ):
        """For use in training conditional diffusion model.

        When a metas file is provided, the metadata is loaded and returned with the data.
        This can be used for lyric and tags conditioning.

        """
        super().__init__()
        self.dataset_dir = dataset_dir
        self.vae_dim = vae_dim
        self.vae_memmap_filename = vae_memmap_filename
        self.semantic_memmap_filename = semantic_memmap_filename
        self.vae_n_tokens = vae_n_tokens
        self.semantic_n_tokens = semantic_n_tokens
        self.metas_filename = metas_filename
        self.vae_use_n_tokens = vae_use_n_tokens
        self.semantic_use_n_tokens = semantic_use_n_tokens
        self.vae_use_float16 = vae_use_float16
        self.vae_rate_hz = vae_rate_hz
        self.semantic_rate_hz = semantic_rate_hz

        # open vae memmap
        if vae_use_float16:
            vae_data = np.memmap(
                os.path.join(dataset_dir, vae_memmap_filename),
                dtype=np.float16,
                mode="r",
            )
        else:
            vae_data = np.memmap(
                os.path.join(dataset_dir, vae_memmap_filename),
                dtype=np.float32,
                mode="r",
            )

        vae_data = vae_data.reshape(-1, vae_n_tokens, vae_dim)
        self.vae_data = vae_data
        print(
            "vae_data",
            os.path.join(dataset_dir, vae_memmap_filename),
            self.vae_data.shape,
        )

        # open semantic memmap
        semantic_data = np.memmap(
            os.path.join(dataset_dir, semantic_memmap_filename),
            dtype=np.uint16,
            mode="r",
        )
        semantic_data = semantic_data.reshape(-1, semantic_n_tokens, 1)
        self.semantic_data = semantic_data[:, :, 0]
        print(
            "semantic_data",
            os.path.join(dataset_dir, semantic_memmap_filename),
            self.semantic_data.shape,
        )

        assert (
            vae_data.shape[0] == semantic_data.shape[0]
        )  # must have same number of rows

        print(f"Found {vae_data.shape[0]} examples.")

        self.metas = read_jsonl(
            os.path.join(dataset_dir, metas_filename), progress=False
        )
        print(
            "Loaded metadata for",
            len(self.metas),
            "examples from",
            os.path.join(dataset_dir, metas_filename),
        )
        assert len(self.metas) == self.vae_data.shape[0]

    def __len__(self):
        return self.vae_data.shape[0]

    def __getitem__(self, idx):
        vae_embeds = torch.from_numpy(
            self.vae_data[idx, : self.vae_use_n_tokens, ...].copy()
        ).float()

        info = {}
        info["idx"] = idx
        info["seconds_start"] = 0
        info["seconds_total"] = 30.0

        # construct a mask based on the number of tokens in the metadata
        # we only have to pad the vae tokens as we use the semantic pad token as input?
        if "n_vae_tokens" in self.metas[idx]:
            n_vae_tokens = self.metas[idx]["n_vae_tokens"]
        else:
            n_vae_tokens = self.vae_use_n_tokens

        padding_mask = torch.ones(self.vae_use_n_tokens)
        if n_vae_tokens < self.vae_use_n_tokens:
            padding_mask[n_vae_tokens:] = 0

        info["padding_mask"] = padding_mask.bool()

        # semantic codes
        semantic_codes = torch.from_numpy(
            self.semantic_data[idx, : self.semantic_use_n_tokens, ...].copy()
        ).long()
        info["semantic_codes"] = semantic_codes

        info["lyrics"] = self.metas[idx].get("text", "")
        info["tags"] = self.metas[idx].get("tags", [])
        info["tags_and_lyrics"] = (info["tags"], info["lyrics"])
        info["phonemes"] = self.metas[idx].get("phonemized_text", "")
        if info["phonemes"] is None:
            info["phonemes"] = ""

        vae_embeds = vae_embeds.permute(1, 0)  # channels, seq_len

        return (vae_embeds, info)


from suno_utils.models.dac.nn.quantize_2 import ResidualVectorQuantize
from suno_utils.utils.s3 import read_from_s3


def load_rvq(codec_path: str):
    sd = read_from_s3(codec_path, read_f=torch.load)
    model = ResidualVectorQuantize(
        input_dim=128,
        n_codebooks=12,
        codebook_size=2048,
        codebook_dim=8,
        quantizer_dropout=0.0,
    )

    model.load_state_dict(
        {k[10:]: v for k, v in sd["state_dict"].items() if k.startswith("quantize")}
    )
    model.eval()

    for param_name, param in model.named_parameters():
        param.requires_grad = False

    return model


def decode_vq(rvq, codes, n_quantizers):
    z_q = 0
    for i, quantizer in enumerate(rvq.quantizers[:n_quantizers]):
        _z_q = quantizer.embed_code(codes[:, :, i]).transpose(1, 2)
        _z_q = quantizer.out_proj(_z_q)
        z_q += _z_q.transpose(1, 2)
    return z_q


class CodecMemmapMapDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        dataset_dir: str,
        codec_memmap_filename: str = "data_codec_val.bin",
        semantic_memmap_filename: str = "data_semantic_val.bin",
        metas_filename: str = "metas_val.jsonl",
        codec_n_codebooks: int = 12,
        codec_n_tokens: int = 9000,
        codec_use_n_tokens: int = 9000,
        semantic_n_tokens: int = 9000,
        semantic_use_n_tokens: int = 9000,
    ):
        """For use in training conditional diffusion model.

        When a metas file is provided, the metadata is loaded and returned with the data.
        This can be used for lyric and tags conditioning.

        """
        super().__init__()
        self.dataset_dir = dataset_dir
        self.codec_memmap_filename = codec_memmap_filename
        self.semantic_memmap_filename = semantic_memmap_filename
        self.codec_n_tokens = codec_n_tokens
        self.semantic_n_tokens = semantic_n_tokens
        self.metas_filename = metas_filename
        self.codec_use_n_tokens = codec_use_n_tokens
        self.semantic_use_n_tokens = semantic_use_n_tokens

        # load rvq from codec
        self.rvq = load_rvq("s3://suno-data/georg/models/codec/dac_2c_25x12.pt")

        # open vae memmap
        codec_data = np.memmap(
            os.path.join(dataset_dir, codec_memmap_filename),
            dtype=np.uint16,
            mode="r",
        )
        codec_data = codec_data.reshape(-1, codec_n_tokens, codec_n_codebooks)
        self.codec_data = codec_data

        # open semantic memmap
        semantic_data = np.memmap(
            os.path.join(dataset_dir, semantic_memmap_filename),
            dtype=np.uint16,
            mode="r",
        )
        semantic_data = semantic_data.reshape(-1, semantic_n_tokens, 1)
        self.semantic_data = semantic_data[:, :, 0]

        print(
            "codec_data",
            os.path.join(dataset_dir, codec_memmap_filename),
            self.codec_data.shape,
        )
        print(
            "semantic_data",
            os.path.join(dataset_dir, semantic_memmap_filename),
            self.semantic_data.shape,
        )

        assert (
            codec_data.shape[0] == semantic_data.shape[0]
        )  # must have same number of rows

        self.memmap_tokens = codec_data.shape[1]

        print(f"Found {codec_data.shape[0]} examples.")

        self.metas = read_jsonl(
            os.path.join(dataset_dir, metas_filename), progress=False
        )
        print(
            "Loaded metadata for",
            len(self.metas),
            "examples from",
            os.path.join(dataset_dir, metas_filename),
        )
        assert len(self.metas) == self.codec_data.shape[0]

    def __len__(self):
        return self.codec_data.shape[0]

    def __getitem__(self, idx):
        codec_codes = torch.from_numpy(
            self.codec_data[idx, : self.memmap_tokens, ...].copy()
        ).long()

        # decode the codec codes back to continuous space
        with torch.no_grad():
            vae_embeds = decode_vq(self.rvq, codec_codes, self.codec_use_n_tokens)

        info = {}
        info["idx"] = idx
        info["seconds_start"] = 0
        info["seconds_total"] = 10.0

        # construct a mask based on the number of tokens in the metadata
        # we only have to pad the vae tokens as we use the semantic pad token as input?
        n_tokens = self.metas[idx]["n_tokens"]

        padding_mask = torch.ones(self.codec_use_n_tokens)
        padding_mask[n_tokens:] = 0

        info["padding_mask"] = padding_mask

        # semantic codes
        semantic_codes = torch.from_numpy(
            self.semantic_data[idx, : self.semantic_use_n_tokens, ...].copy()
        ).long()
        info["semantic_codes"] = semantic_codes

        info["lyrics"] = self.metas[idx].get("text", "")
        info["tags"] = self.metas[idx].get("tags", [])
        info["tags_and_lyrics"] = (info["tags"], info["lyrics"])

        vae_embeds = vae_embeds.permute(1, 0)  # channels, seq_len

        return (vae_embeds, info)


class MemmapDataset(torch.utils.data.IterableDataset):
    def __init__(
        self,
        vae_memmap_path: str,
        semantic_memmap_path: str = None,
        codec_memmap_path: str = None,
        custom_metadata_fn: Optional[Callable[[str], str]] = None,
    ):
        super().__init__()
        self.filenames = []
        self.custom_metadata_fn = custom_metadata_fn

        vae_data = np.memmap(vae_memmap_path, dtype=np.float32, mode="r")
        vae_data = vae_data.reshape(-1, N_TOKENS_MEMMAP, VAE_DIM)
        self.vae_data = vae_data

        # debug
        # self.vae_data = np.zeros((100, N_TOKENS_MEMMAP, VAE_DIM)).astype(np.float32)

        self.semantic_data = None
        self.codec_data = None
        print("vae_data", self.vae_data.shape)

        if semantic_memmap_path is not None:
            semantic_data = np.memmap(semantic_memmap_path, dtype=np.uint16, mode="r")
            semantic_data = semantic_data.reshape(-1, N_TOKENS_MEMMAP)

            assert (
                vae_data.shape[0] == semantic_data.shape[0]
            )  # must have same number of rows
            self.semantic_data = semantic_data

            # debug
            # self.semantic_data = np.zeros((100, N_TOKENS_MEMMAP)).astype(np.int16)

            print("semantic_data", self.semantic_data.shape)

        if codec_memmap_path is not None:
            codec_data = np.memmap(codec_memmap_path, dtype=np.uint16, mode="r")
            codec_data = codec_data.reshape(-1, N_TOKENS_MEMMAP, CODEC_N_CODEBOOKS)
            self.codec_data = codec_data

            # debug
            # self.codec_data = np.zeros(
            #    (100, N_TOKENS_MEMMAP, CODEC_N_CODEBOOKS)
            # ).astype(np.int16)
            print("codec_data", self.codec_data.shape)

        print(f"Found {vae_data.shape[0]} examples.")

    def __len__(self):
        return self.vae_data.shape[0]

    def __iter__(self):
        while True:
            rand_idx = np.random.randint(0, self.vae_data.shape[0] - 1)
            start_time = time.time()

            info = {}

            t_start = 0
            t_end = 10.0
            seconds_start = 0
            seconds_total = t_end

            vae_embed = torch.from_numpy(self.vae_data[rand_idx, ...].copy())

            if self.semantic_data is not None:
                semantic_codes = torch.from_numpy(
                    self.semantic_data[rand_idx, ...].copy()
                ).long()
                info["semantic_codes"] = semantic_codes

            if self.codec_data is not None:
                codec_codes = torch.from_numpy(
                    self.codec_data[rand_idx, ...].copy()
                ).long()
                info["codec_codes"] = codec_codes

            info["idx"] = rand_idx
            info["timestamps"] = (t_start, t_end)
            info["seconds_start"] = seconds_start
            info["seconds_total"] = seconds_total

            end_time = time.time()

            info["load_time"] = end_time - start_time

            # disabling this inferace for now and hardcoding conditioning
            # this means we do not do any conditioning dropout rn
            # if self.custom_metadata_fn is not None:
            #    custom_metadata = self.custom_metadata_fn(info, semantic_codes)
            #    info.update(custom_metadata)
            #
            #    if "__reject__" in info and info["__reject__"]:
            #        return self[random.randrange(len(self))]

            vae_embed = vae_embed.permute(1, 0)  # channels, seq_len

            yield (vae_embed, info)


class SampleDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        paths,
        sample_size=65536,
        sample_rate=48000,
        keywords=None,
        relpath=None,
        random_crop=True,
        force_channels="stereo",
        custom_metadata_fn: Optional[Callable[[str], str]] = None,
    ):
        super().__init__()
        self.filenames = []
        self.relpath = relpath

        self.augs = torch.nn.Sequential(
            PhaseFlipper(),
        )

        self.pad_crop = PadCrop_Normalized_T(
            sample_size, sample_rate, randomize=random_crop
        )

        self.force_channels = force_channels

        self.encoding = torch.nn.Sequential(
            Stereo() if self.force_channels == "stereo" else torch.nn.Identity(),
            Mono() if self.force_channels == "mono" else torch.nn.Identity(),
        )

        self.filenames = get_audio_filenames(paths, keywords)

        print(f"Found {len(self.filenames)} files")

        self.sr = sample_rate

        self.custom_metadata_fn = custom_metadata_fn

    def load_file(self, filename):
        ext = filename.split(".")[-1]

        if ext == "mp3":
            with AudioFile(filename) as f:
                audio = f.read(f.frames)
                audio = torch.from_numpy(audio)
                in_sr = f.samplerate
        else:
            audio, in_sr = torchaudio.load(filename, format=ext)

        if in_sr != self.sr:
            resample_tf = T.Resample(in_sr, self.sr)
            audio = resample_tf(audio)

        return audio

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, idx):
        audio_filename = self.filenames[idx]
        try:
            start_time = time.time()
            audio = self.load_file(audio_filename)

            audio, t_start, t_end, seconds_start, seconds_total, padding_mask = (
                self.pad_crop(audio)
            )

            # Run augmentations on this sample (including random crop)
            if self.augs is not None:
                audio = self.augs(audio)

            audio = audio.clamp(-1, 1)

            # Encode the file to assist in prediction
            if self.encoding is not None:
                audio = self.encoding(audio)

            info = {}

            info["path"] = audio_filename

            if self.relpath is not None:
                info["relpath"] = path.relpath(audio_filename, self.relpath)

            info["timestamps"] = (t_start, t_end)
            info["seconds_start"] = seconds_start
            info["seconds_total"] = seconds_total
            info["padding_mask"] = padding_mask

            end_time = time.time()

            info["load_time"] = end_time - start_time

            if self.custom_metadata_fn is not None:
                custom_metadata = self.custom_metadata_fn(info, audio)
                info.update(custom_metadata)

                if "__reject__" in info and info["__reject__"]:
                    return self[random.randrange(len(self))]

            return (audio, info)
        except Exception as e:
            print(f"Couldn't load file {audio_filename}: {e}")
            return self[random.randrange(len(self))]


def group_by_keys(
    data, keys=wds.tariterators.base_plus_ext, lcase=True, suffixes=None, handler=None
):
    """Return function over iterator that groups key, value pairs into samples.
    :param keys: function that splits the key into key and extension (base_plus_ext)
    :param lcase: convert suffixes to lower case (Default value = True)
    """
    current_sample = None
    for filesample in data:
        assert isinstance(filesample, dict)
        fname, value = filesample["fname"], filesample["data"]
        prefix, suffix = keys(fname)
        if wds.tariterators.trace:
            print(
                prefix,
                suffix,
                current_sample.keys() if isinstance(current_sample, dict) else None,
            )
        if prefix is None:
            continue
        if lcase:
            suffix = suffix.lower()
        if current_sample is None or prefix != current_sample["__key__"]:
            if wds.tariterators.valid_sample(current_sample):
                yield current_sample
            current_sample = dict(__key__=prefix, __url__=filesample["__url__"])
        if suffix in current_sample:
            print(
                f"{fname}: duplicate file name in tar file {suffix} {current_sample.keys()}"
            )
        if suffixes is None or suffix in suffixes:
            current_sample[suffix] = value
    if wds.tariterators.valid_sample(current_sample):
        yield current_sample


wds.tariterators.group_by_keys = group_by_keys

# S3 code and WDS preprocessing code based on implementation by Scott Hawley originally in https://github.com/zqevans/audio-diffusion/blob/main/dataset/dataset.py


def get_s3_contents(
    dataset_path,
    s3_url_prefix=None,
    filter="",
    recursive=True,
    debug=False,
    profile=None,
):
    """
    Returns a list of full S3 paths to files in a given S3 bucket and directory path.
    """
    # Ensure dataset_path ends with a trailing slash
    if dataset_path != "" and not dataset_path.endswith("/"):
        dataset_path += "/"
    # Use posixpath to construct the S3 URL path
    bucket_path = posixpath.join(s3_url_prefix or "", dataset_path)
    # Construct the `aws s3 ls` command
    cmd = ["aws", "s3", "ls", bucket_path]

    if profile is not None:
        cmd.extend(["--profile", profile])

    if recursive:
        # Add the --recursive flag if requested
        cmd.append("--recursive")

    # Run the `aws s3 ls` command and capture the output
    run_ls = subprocess.run(cmd, capture_output=True, check=True)
    # Split the output into lines and strip whitespace from each line
    contents = run_ls.stdout.decode("utf-8").split("\n")
    contents = [x.strip() for x in contents if x]
    # Remove the timestamp from lines that begin with a timestamp
    contents = [
        (
            re.sub(r"^\S+\s+\S+\s+\d+\s+", "", x)
            if re.match(r"^\S+\s+\S+\s+\d+\s+", x)
            else x
        )
        for x in contents
    ]
    # Construct a full S3 path for each file in the contents list
    contents = [
        posixpath.join(s3_url_prefix or "", x) for x in contents if not x.endswith("/")
    ]
    # Apply the filter, if specified
    if filter:
        contents = [x for x in contents if filter in x]
    # Remove redundant directory names in the S3 URL
    if recursive:
        # Get the main directory name from the S3 URL
        main_dir = "/".join(bucket_path.split("/")[3:])
        # Remove the redundant directory names from each file path
        contents = [x.replace(f"{main_dir}", "").replace("//", "/") for x in contents]
    # Print debugging information, if requested
    if debug:
        print("contents = \n", contents)
    # Return the list of S3 paths to files
    return contents


def get_all_s3_urls(
    names=[],  # list of all valid [LAION AudioDataset] dataset names
    # list of subsets you want from those datasets, e.g. ['train','valid']
    subsets=[""],
    s3_url_prefix=None,  # prefix for those dataset names
    recursive=True,  # recursively list all tar files in all subdirs
    filter_str="tar",  # only grab files with this substring
    # print debugging info -- note: info displayed likely to change at dev's whims
    debug=False,
    profiles={},  # dictionary of profiles for each item in names, e.g. {'dataset1': 'profile1', 'dataset2': 'profile2'}
):
    "get urls of shards (tar files) for multiple datasets in one s3 bucket"
    urls = []
    for name in names:
        # If s3_url_prefix is not specified, assume the full S3 path is included in each element of the names list
        if s3_url_prefix is None:
            contents_str = name
        else:
            # Construct the S3 path using the s3_url_prefix and the current name value
            contents_str = posixpath.join(s3_url_prefix, name)
        if debug:
            print(f"get_all_s3_urls: {contents_str}:")
        for subset in subsets:
            subset_str = posixpath.join(contents_str, subset)
            if debug:
                print(f"subset_str = {subset_str}")
            # Get the list of tar files in the current subset directory
            profile = profiles.get(name, None)
            tar_list = get_s3_contents(
                subset_str,
                s3_url_prefix=None,
                recursive=recursive,
                filter=filter_str,
                debug=debug,
                profile=profile,
            )
            for tar in tar_list:
                # Escape spaces and parentheses in the tar filename for use in the shell command
                tar = tar.replace(" ", "\ ").replace("(", "\(").replace(")", "\)")
                # Construct the S3 path to the current tar file
                s3_path = posixpath.join(name, subset, tar) + " -"
                # Construct the AWS CLI command to download the current tar file
                if s3_url_prefix is None:
                    request_str = f"pipe:aws s3 --cli-connect-timeout 0 cp {s3_path}"
                else:
                    request_str = f"pipe:aws s3 --cli-connect-timeout 0 cp {posixpath.join(s3_url_prefix, s3_path)}"
                if profiles.get(name):
                    request_str += f" --profile {profiles.get(name)}"
                if debug:
                    print("request_str = ", request_str)
                # Add the constructed URL to the list of URLs
                urls.append(request_str)
    return urls


def log_and_continue(exn):
    """Call in an exception handler to ignore any exception, isssue a warning, and continue."""
    print(f"Handling webdataset error ({repr(exn)}). Ignoring.")
    return True


def is_valid_sample(sample):
    has_json = "json" in sample
    has_audio = "audio" in sample
    is_silent = is_silence(sample["audio"])
    is_rejected = "__reject__" in sample["json"] and sample["json"]["__reject__"]

    return has_json and has_audio and not is_silent and not is_rejected


class S3DatasetConfig:
    def __init__(
        self,
        id: str,
        s3_path: str,
        custom_metadata_fn: Optional[Callable[[str], str]] = None,
        profile: Optional[str] = None,
    ):
        self.id = id
        self.s3_path = s3_path
        self.custom_metadata_fn = custom_metadata_fn
        self.profile = profile
        self.urls = []

    def load_data_urls(self):
        self.urls = get_all_s3_urls(
            names=[self.s3_path],
            s3_url_prefix=None,
            recursive=True,
            profiles={self.s3_path: self.profile} if self.profile else {},
        )

        return self.urls


def audio_decoder(key, value):
    # Get file extension from key
    ext = key.split(".")[-1]

    if ext in AUDIO_KEYS:
        return torchaudio.load(io.BytesIO(value))
    else:
        return None


def collation_fn(samples):
    batched = list(zip(*samples))
    result = []
    for b in batched:
        if isinstance(b[0], (int, float)):
            b = np.array(b)
        elif isinstance(b[0], torch.Tensor):
            b = torch.stack(b)
        elif isinstance(b[0], np.ndarray):
            b = np.array(b)
        else:
            b = b
        result.append(b)
    return result


class S3WebDataLoader:
    def __init__(
        self,
        datasets: List[S3DatasetConfig],
        batch_size,
        sample_size,
        sample_rate=48000,
        num_workers=8,
        epoch_steps=1000,
        random_crop=True,
        force_channels="stereo",
        augment_phase=True,
        **data_loader_kwargs,
    ):
        self.datasets = datasets

        self.sample_size = sample_size
        self.sample_rate = sample_rate
        self.random_crop = random_crop
        self.force_channels = force_channels
        self.augment_phase = augment_phase

        urls = [dataset.load_data_urls() for dataset in datasets]

        # Flatten the list of lists of URLs
        urls = [url for dataset_urls in urls for url in dataset_urls]

        self.dataset = wds.DataPipeline(
            wds.ResampledShards(urls),
            wds.tarfile_to_samples(handler=log_and_continue),
            wds.decode(audio_decoder, handler=log_and_continue),
            wds.map(self.wds_preprocess, handler=log_and_continue),
            wds.select(is_valid_sample),
            wds.to_tuple("audio", "json", handler=log_and_continue),
            wds.batched(batch_size, partial=False, collation_fn=collation_fn),
        ).with_epoch(epoch_steps // num_workers if num_workers > 0 else epoch_steps)

        self.data_loader = wds.WebLoader(
            self.dataset, num_workers=num_workers, **data_loader_kwargs
        )

    def wds_preprocess(self, sample):
        found_key, rewrite_key = "", ""
        for k, v in sample.items():  # print the all entries in dict
            for akey in AUDIO_KEYS:
                if k.endswith(akey):
                    # to rename long/weird key with its simpler counterpart
                    found_key, rewrite_key = k, akey
                    break
            if "" != found_key:
                break
        if "" == found_key:  # got no audio!
            return None  # try returning None to tell WebDataset to skip this one

        audio, in_sr = sample[found_key]
        if in_sr != self.sample_rate:
            resample_tf = T.Resample(in_sr, self.sample_rate)
            audio = resample_tf(audio)

        if self.sample_size is not None:
            # Pad/crop and get the relative timestamp
            pad_crop = PadCrop_Normalized_T(
                self.sample_size,
                randomize=self.random_crop,
                sample_rate=self.sample_rate,
            )
            audio, t_start, t_end, seconds_start, seconds_total, padding_mask = (
                pad_crop(audio)
            )
            sample["json"]["seconds_start"] = seconds_start
            sample["json"]["seconds_total"] = seconds_total
            sample["json"]["padding_mask"] = padding_mask
        else:
            t_start, t_end = 0, 1

        # Check if audio is length zero, initialize to a single zero if so
        if audio.shape[-1] == 0:
            audio = torch.zeros(1, 1)

        # Make the audio stereo and augment by randomly inverting phase
        augs = torch.nn.Sequential(
            Stereo() if self.force_channels == "stereo" else torch.nn.Identity(),
            Mono() if self.force_channels == "mono" else torch.nn.Identity(),
            PhaseFlipper() if self.augment_phase else torch.nn.Identity(),
        )

        audio = augs(audio)

        sample["json"]["timestamps"] = (t_start, t_end)

        if "text" in sample["json"]:
            sample["json"]["prompt"] = sample["json"]["text"]

        # Check for custom metadata functions
        for dataset in self.datasets:
            if dataset.custom_metadata_fn is None:
                continue

            if dataset.s3_path in sample["__url__"]:
                custom_metadata = dataset.custom_metadata_fn(sample["json"], audio)
                sample["json"].update(custom_metadata)

        if (
            found_key != rewrite_key
        ):  # rename long/weird key with its simpler counterpart
            del sample[found_key]

        sample["audio"] = audio

        # Add audio to the metadata as well for conditioning
        sample["json"]["audio"] = audio

        return sample


# Custom worker initialization function
def worker_init_fn(worker_id):
    # Set the seed for each worker
    slurm_procid = os.environ.get("SLURM_PROCID")
    if slurm_procid is None:
        seed = worker_id
    else:
        # combine SLURM_PROCID and worker_id to get a unique seed for each worker
        # this assumes that the number of workers is less than 1024
        seed = (int(slurm_procid) * 1024) + worker_id
    print(slurm_procid, worker_id, seed)
    pl.seed_everything(seed)


def create_dataloader_from_config(
    dataset_config,
    batch_size,
    sample_size,
    sample_rate,
    audio_channels=2,
    num_workers=4,
):
    dataset_type = dataset_config.get("dataset_type", None)

    assert dataset_type is not None, "Dataset type must be specified in dataset config"

    if audio_channels == 1:
        force_channels = "mono"
    else:
        force_channels = "stereo"

    if dataset_type == "audio_dir":
        audio_dir_configs = dataset_config.get("datasets", None)

        assert (
            audio_dir_configs is not None
        ), 'Directory configuration must be specified in datasets["dataset"]'

        training_dirs = []

        custom_metadata_fn = None
        custom_metadata_module_path = dataset_config.get("custom_metadata_module", None)

        if custom_metadata_module_path is not None:
            spec = importlib.util.spec_from_file_location(
                "metadata_module", custom_metadata_module_path
            )
            metadata_module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(metadata_module)

            custom_metadata_fn = metadata_module.get_custom_metadata

        for audio_dir_config in audio_dir_configs:
            audio_dir_path = audio_dir_config.get("path", None)
            assert (
                audio_dir_path is not None
            ), "Path must be set for local audio directory configuration"
            training_dirs.append(audio_dir_path)

        train_set = SampleDataset(
            training_dirs,
            sample_rate=sample_rate,
            sample_size=sample_size,
            random_crop=dataset_config.get("random_crop", True),
            force_channels=force_channels,
            custom_metadata_fn=custom_metadata_fn,
            relpath=training_dirs[
                0
            ],  # TODO: Make relpath relative to each training dir
        )

        return torch.utils.data.DataLoader(
            train_set,
            batch_size,
            shuffle=True,
            num_workers=num_workers,
            persistent_workers=True,
            pin_memory=True,
            drop_last=True,
            collate_fn=collation_fn,
        )

    elif dataset_type == "memmap":
        custom_metadata_fn = None
        custom_metadata_module_path = dataset_config.get("custom_metadata_module", None)

        if custom_metadata_module_path is not None:
            spec = importlib.util.spec_from_file_location(
                "metadata_module", custom_metadata_module_path
            )
            metadata_module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(metadata_module)

            custom_metadata_fn = metadata_module.get_custom_metadata

        train_set = MemmapDataset(
            dataset_config.get("train_vae_memmap_path", None),
            semantic_memmap_path=dataset_config.get("train_semantic_memmap_path", None),
            codec_memmap_path=dataset_config.get("train_codec_memmap_path", None),
            custom_metadata_fn=custom_metadata_fn,
        )

        train_dataloader = torch.utils.data.DataLoader(
            train_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            collate_fn=collation_fn,
        )

        val_dataloader = None
        val_vae_memmap_path = dataset_config.get("train_vae_memmap_path", None)
        if val_vae_memmap_path is not None:
            val_set = MemmapDataset(
                dataset_config.get("val_vae_memmap_path", None),
                semantic_memmap_path=dataset_config.get(
                    "val_semantic_memmap_path", None
                ),
                codec_memmap_path=dataset_config.get("val_codec_memmap_path", None),
                custom_metadata_fn=custom_metadata_fn,
            )

            val_dataloader = torch.utils.data.DataLoader(
                val_set,
                batch_size,
                num_workers=num_workers,
                # persistent_workers=True,
                worker_init_fn=worker_init_fn,
                pin_memory=True,
                drop_last=True,
                collate_fn=collation_fn,
            )

    elif dataset_type == "dummy_memmap":
        train_set = DummyMemmapMapDataset(
            vae_dim=dataset_config.get("vae_dim", 128),
            vae_n_tokens=dataset_config.get("vae_n_tokens_memmap", 3000),
            semantic_n_tokens=dataset_config.get("semantic_n_tokens_memmap", 750),
        )

        train_dataloader = torch.utils.data.DataLoader(
            train_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            collate_fn=collation_fn,
        )

        val_set = DummyMemmapMapDataset(
            vae_dim=dataset_config.get("vae_dim", 128),
            vae_n_tokens=dataset_config.get("vae_n_tokens_memmap", 3000),
            semantic_n_tokens=dataset_config.get("semantic_n_tokens_memmap", 750),
        )

        val_dataloader = torch.utils.data.DataLoader(
            val_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            collate_fn=collation_fn,
        )

        return train_dataloader, val_dataloader

    elif dataset_type == "general_context_memmap":
        train_set = ContextMemmapMapDataset(
            dataset_dir=dataset_config.get("dataset_dir", None),
            vae_memmap_filename=dataset_config.get("train_vae_memmap_filename", None),
            semantic_memmap_filename=dataset_config.get(
                "train_semantic_memmap_filename", None
            ),
            vae_dim=dataset_config.get("vae_dim", 128),
            vae_n_tokens=dataset_config.get("vae_n_tokens_memmap", 3000),
            vae_use_float16=dataset_config.get("vae_use_float16", False),
            semantic_n_tokens=dataset_config.get("semantic_n_tokens_memmap", 750),
            metas_filename=dataset_config.get("train_metas_filename", None),
            vae_rate_hz=dataset_config.get("vae_rate_hz", 100),
            semantic_rate_hz=dataset_config.get("semantic_rate_hz", 25),
        )

        train_dataloader = torch.utils.data.DataLoader(
            train_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            shuffle=True,
            collate_fn=collation_fn,
        )

        val_set = ContextMemmapMapDataset(
            dataset_dir=dataset_config.get("dataset_dir", None),
            vae_memmap_filename=dataset_config.get("val_vae_memmap_filename", None),
            semantic_memmap_filename=dataset_config.get(
                "val_semantic_memmap_filename", None
            ),
            vae_dim=dataset_config.get("vae_dim", 128),
            vae_n_tokens=dataset_config.get("vae_n_tokens_memmap", 3000),
            vae_use_float16=dataset_config.get("vae_use_float16", False),
            semantic_n_tokens=dataset_config.get("semantic_n_tokens_memmap", 750),
            metas_filename=dataset_config.get("val_metas_filename", None),
            vae_rate_hz=dataset_config.get("vae_rate_hz", 100),
            semantic_rate_hz=dataset_config.get("semantic_rate_hz", 25),
        )

        val_dataloader = torch.utils.data.DataLoader(
            val_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            shuffle=False,
            collate_fn=collation_fn,
        )

        return train_dataloader, val_dataloader

    elif dataset_type == "general_memmap":
        train_set = GeneralMemmapMapDataset(
            dataset_dir=dataset_config.get("dataset_dir", None),
            vae_memmap_filename=dataset_config.get("train_vae_memmap_filename", None),
            semantic_memmap_filename=dataset_config.get(
                "train_semantic_memmap_filename", None
            ),
            vae_dim=dataset_config.get("vae_dim", 128),
            vae_n_tokens=dataset_config.get("vae_n_tokens_memmap", 3000),
            vae_use_float16=dataset_config.get("vae_use_float16", False),
            semantic_n_tokens=dataset_config.get("semantic_n_tokens_memmap", 750),
            metas_filename=dataset_config.get("train_metas_filename", None),
            semantic_use_n_tokens=dataset_config.get("semantic_use_n_tokens", 750),
            vae_use_n_tokens=dataset_config.get("vae_use_n_tokens", 3000),
            vae_rate_hz=dataset_config.get("vae_rate_hz", 100),
            semantic_rate_hz=dataset_config.get("semantic_rate_hz", 25),
            right_padding_mask_rate=dataset_config.get("right_padding_mask_rate", 0.0),
        )

        train_dataloader = torch.utils.data.DataLoader(
            train_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            shuffle=True,
            collate_fn=collation_fn,
        )

        val_set = GeneralMemmapMapDataset(
            dataset_dir=dataset_config.get("dataset_dir", None),
            vae_memmap_filename=dataset_config.get("val_vae_memmap_filename", None),
            semantic_memmap_filename=dataset_config.get(
                "val_semantic_memmap_filename", None
            ),
            vae_dim=dataset_config.get("vae_dim", 128),
            vae_n_tokens=dataset_config.get("vae_n_tokens_memmap", 3000),
            vae_use_float16=dataset_config.get("vae_use_float16", False),
            semantic_n_tokens=dataset_config.get("semantic_n_tokens_memmap", 750),
            metas_filename=dataset_config.get("val_metas_filename", None),
            semantic_use_n_tokens=dataset_config.get("semantic_use_n_tokens", 750),
            vae_use_n_tokens=dataset_config.get("vae_use_n_tokens", 3000),
            vae_rate_hz=dataset_config.get("vae_rate_hz", 100),
            semantic_rate_hz=dataset_config.get("semantic_rate_hz", 25),
            right_padding_mask_rate=dataset_config.get("right_padding_mask_rate", 0.0),
        )

        val_dataloader = torch.utils.data.DataLoader(
            val_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            shuffle=False,
            collate_fn=collation_fn,
        )

        return train_dataloader, val_dataloader

    elif dataset_type == "codec_memmap":
        train_set = CodecMemmapMapDataset(
            dataset_dir=dataset_config.get("dataset_dir", None),
            codec_memmap_filename=dataset_config.get(
                "train_codec_memmap_filename", None
            ),
            semantic_memmap_filename=dataset_config.get(
                "train_semantic_memmap_filename", None
            ),
            codec_n_codebooks=dataset_config.get("codec_n_codebooks", 12),
            codec_n_tokens=dataset_config.get("codec_n_tokens_memmap", 750),
            semantic_n_tokens=dataset_config.get("semantic_n_tokens_memmap", 750),
            metas_filename=dataset_config.get("train_metas_filename", None),
            semantic_use_n_tokens=dataset_config.get("semantic_use_n_tokens", 750),
            codec_use_n_tokens=dataset_config.get("decode_use_n_tokens", 750),
        )

        train_dataloader = torch.utils.data.DataLoader(
            train_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            shuffle=True,
            collate_fn=collation_fn,
        )

        val_set = CodecMemmapMapDataset(
            dataset_dir=dataset_config.get("dataset_dir", None),
            codec_memmap_filename=dataset_config.get("val_codec_memmap_filename", None),
            semantic_memmap_filename=dataset_config.get(
                "val_semantic_memmap_filename", None
            ),
            codec_n_codebooks=dataset_config.get("codec_n_codebooks", 12),
            codec_n_tokens=dataset_config.get("codec_n_tokens_memmap", 750),
            semantic_n_tokens=dataset_config.get("semantic_n_tokens_memmap", 750),
            metas_filename=dataset_config.get("val_metas_filename", None),
            semantic_use_n_tokens=dataset_config.get("semantic_use_n_tokens", 750),
            codec_use_n_tokens=dataset_config.get("codec_use_n_tokens", 750),
        )

        val_dataloader = torch.utils.data.DataLoader(
            val_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            shuffle=False,
            collate_fn=collation_fn,
        )

        return train_dataloader, val_dataloader

    elif dataset_type == "memmap_vae":
        train_set = VAEMemmapMapDataset(
            dataset_config.get("train_vae_memmap_path", None),
            dataset_config.get("train_metas_path", None),
            dataset_config.get("vae_dim", 128),
            dataset_config.get("n_tokens_memmap", 1000),
        )

        train_dataloader = torch.utils.data.DataLoader(
            train_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            shuffle=True,
            collate_fn=collation_fn,
        )

        val_set = VAEMemmapMapDataset(
            dataset_config.get("val_vae_memmap_path", None),
            dataset_config.get("val_metas_path", None),
            dataset_config.get("vae_dim", 128),
            dataset_config.get("n_tokens_memmap", 1000),
        )

        val_dataloader = torch.utils.data.DataLoader(
            val_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            shuffle=True,
            collate_fn=collation_fn,
        )

        return train_dataloader, val_dataloader

    elif dataset_type == "memmap_semantic_autoencoder":
        custom_metadata_fn = None
        custom_metadata_module_path = dataset_config.get("custom_metadata_module", None)

        if custom_metadata_module_path is not None:
            spec = importlib.util.spec_from_file_location(
                "metadata_module", custom_metadata_module_path
            )
            metadata_module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(metadata_module)

            custom_metadata_fn = metadata_module.get_custom_metadata

        train_set = SemanticToAutoencoderMemmapDataset(
            dataset_config.get("train_semantic_memmap_path", None),
            dataset_config.get("train_vae_memmap_path", None),
            custom_metadata_fn=custom_metadata_fn,
        )

        train_dataloader = torch.utils.data.DataLoader(
            train_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            collate_fn=collation_fn,
        )

        val_set = SemanticToAutoencoderMemmapDataset(
            dataset_config.get("val_semantic_memmap_path", None),
            dataset_config.get("val_vae_memmap_path", None),
            custom_metadata_fn=custom_metadata_fn,
        )

        val_dataloader = torch.utils.data.DataLoader(
            val_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            collate_fn=collation_fn,
        )

        return train_dataloader, val_dataloader

    elif dataset_type == "memmap_discrete_vae":
        custom_metadata_fn = None
        custom_metadata_module_path = dataset_config.get("custom_metadata_module", None)

        if custom_metadata_module_path is not None:
            spec = importlib.util.spec_from_file_location(
                "metadata_module", custom_metadata_module_path
            )
            metadata_module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(metadata_module)

            custom_metadata_fn = metadata_module.get_custom_metadata

        train_set = DiscreteVAEtoVAEMemmapDataset(
            dataset_config.get("train_discrete_memmap_path", None),
            dataset_config.get("train_vae_memmap_path", None),
            dataset_config.get("vae_dim", 128),
            dataset_config.get("n_tokens_memmap", 1000),
            custom_metadata_fn=custom_metadata_fn,
        )

        train_dataloader = torch.utils.data.DataLoader(
            train_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            collate_fn=collation_fn,
        )

        val_set = DiscreteVAEtoVAEMemmapDataset(
            dataset_config.get("val_discrete_memmap_path", None),
            dataset_config.get("val_vae_memmap_path", None),
            dataset_config.get("vae_dim", 128),
            dataset_config.get("n_tokens_memmap", 1000),
            custom_metadata_fn=custom_metadata_fn,
        )

        val_dataloader = torch.utils.data.DataLoader(
            val_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            collate_fn=collation_fn,
        )

        return train_dataloader, val_dataloader

    elif dataset_type == "memmap_semantic":
        custom_metadata_fn = None
        custom_metadata_module_path = dataset_config.get("custom_metadata_module", None)

        if custom_metadata_module_path is not None:
            spec = importlib.util.spec_from_file_location(
                "metadata_module", custom_metadata_module_path
            )
            metadata_module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(metadata_module)

            custom_metadata_fn = metadata_module.get_custom_metadata

        train_set = SemanticToCodecMemmapDataset(
            dataset_config.get("train_semantic_memmap_path", None),
            dataset_config.get("train_codec_memmap_path", None),
            custom_metadata_fn=custom_metadata_fn,
        )

        train_dataloader = torch.utils.data.DataLoader(
            train_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            collate_fn=collation_fn,
        )

        val_set = SemanticToCodecMemmapDataset(
            dataset_config.get("val_semantic_memmap_path", None),
            dataset_config.get("val_codec_memmap_path", None),
            custom_metadata_fn=custom_metadata_fn,
        )

        val_dataloader = torch.utils.data.DataLoader(
            val_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            collate_fn=collation_fn,
        )

        return train_dataloader, val_dataloader

    elif dataset_type == "memmap_denoising":
        custom_metadata_fn = None
        custom_metadata_module_path = dataset_config.get("custom_metadata_module", None)

        if custom_metadata_module_path is not None:
            spec = importlib.util.spec_from_file_location(
                "metadata_module", custom_metadata_module_path
            )
            metadata_module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(metadata_module)

            custom_metadata_fn = metadata_module.get_custom_metadata

        train_set = DenoisingMemmapDataset(
            dataset_config.get("train_input_memmap_path", None),
            dataset_config.get("train_corrupt_memmap_path", None),
            custom_metadata_fn=custom_metadata_fn,
        )

        train_dataloader = torch.utils.data.DataLoader(
            train_set,
            batch_size,
            num_workers=num_workers,
            persistent_workers=True,
            worker_init_fn=worker_init_fn,
            pin_memory=True,
            drop_last=True,
            collate_fn=collation_fn,
        )

        val_dataloader = None
        val_input_memmap_path = dataset_config.get("val_input_memmap_path", None)
        if val_input_memmap_path is not None:
            val_set = DenoisingMemmapDataset(
                dataset_config.get("val_input_memmap_path", None),
                dataset_config.get("val_corrupt_memmap_path", None),
                custom_metadata_fn=custom_metadata_fn,
            )

            val_dataloader = torch.utils.data.DataLoader(
                val_set,
                batch_size,
                num_workers=num_workers,
                persistent_workers=True,
                worker_init_fn=worker_init_fn,
                pin_memory=True,
                drop_last=True,
                collate_fn=collation_fn,
            )

        return train_dataloader, val_dataloader

    elif dataset_type == "s3":
        dataset_configs = []

        for s3_config in dataset_config["datasets"]:
            custom_metadata_fn = None
            custom_metadata_module_path = s3_config.get("custom_metadata_module", None)

            if custom_metadata_module_path is not None:
                spec = importlib.util.spec_from_file_location(
                    "metadata_module", custom_metadata_module_path
                )
                metadata_module = importlib.util.module_from_spec(spec)
                spec.loader.exec_module(metadata_module)

                custom_metadata_fn = metadata_module.get_custom_metadata

            dataset_configs.append(
                S3DatasetConfig(
                    id=s3_config["id"],
                    s3_path=s3_config["s3_path"],
                    custom_metadata_fn=custom_metadata_fn,
                    profile=s3_config.get("profile", None),
                )
            )

        return S3WebDataLoader(
            dataset_configs,
            sample_rate=sample_rate,
            sample_size=sample_size,
            batch_size=batch_size,
            random_crop=dataset_config.get("random_crop", True),
            num_workers=num_workers,
            persistent_workers=True,
            force_channels=force_channels,
            epoch_steps=dataset_config.get("epoch_steps", 2000),
        ).data_loader