import os
import glob
import random
import torch
import numpy as np
from torch.utils import data
from suno_utils.audio import Audio
from torchaudio_augmentations import (
    RandomApply,
    Noise,
    Gain,
    PitchShift,
    Compose,
)


class ArtistSimDataset(data.Dataset):
    def __init__(
        self,
        metadata,
        task_indices,
        split="train",
        input_length_s=30.0,
        sample_rate=24000,
        num_samples=-1,
    ):
        assert split in ["train", "valid"]
        self.metadata = metadata
        self.artist_to_indices = task_indices["artist_sim"][split]
        self.artists = list(self.artist_to_indices.keys())
        self.input_length_s = input_length_s
        self.sample_rate = sample_rate
        self.num_samples = num_samples
        self.split = split

        # get augmentation
        if split == "train":
            self._get_augmentations()

        print(f"{len(self.artists)} artists are available for artist_sim {split} set")

    def _get_augmentations(self):
        # Stochastic data augmentation
        transforms = [
            RandomApply([Noise(min_snr=0.1, max_snr=0.5)], p=0.3),
            RandomApply([Gain()], p=0.2),
            RandomApply(
                [
                    PitchShift(
                        n_samples=int(self.input_length_s * self.sample_rate),
                        sample_rate=self.sample_rate,
                        pitch_shift_min=-3.0,
                        pitch_shift_max=3.0,
                    )
                ],
                p=0.4,
            ),
        ]
        self.augmentation = Compose(transforms=transforms)

    def __getitem__(self, index):
        # read data
        if self.split == "train":
            # Select a random artist
            artist = random.choice(self.artists)

            # Select two unique tracks from the artist
            indices = random.sample(self.artist_to_indices[artist], 2)
            audio_path_1 = self.metadata[indices[0]]["filepath"]
            audio_path_2 = self.metadata[indices[1]]["filepath"]

        elif self.split == "valid":
            artist = self.artists[index]
            audio_path_1 = self.metadata[self.artist_to_indices[artist][0]]["filepath"]
            audio_path_2 = self.metadata[self.artist_to_indices[artist][1]]["filepath"]

        audio_1 = Audio.from_file(audio_path_1, sample_rate=self.sample_rate)
        audio_2 = Audio.from_file(audio_path_2, sample_rate=self.sample_rate)

        if self.split == "train":
            # random crop
            try:
                start_ms_1 = random.randint(
                    0, (audio_1.duration_ms - int(1000 * self.input_length_s) - 1)
                )
                start_ms_2 = random.randint(
                    0, (audio_2.duration_ms - int(1000 * self.input_length_s) - 1)
                )
                start_s_1 = start_ms_1 / 1000
                start_s_2 = start_ms_2 / 1000
            except ValueError:
                start_s_1 = 0.0
                start_s_2 = 0.0
            wav_1 = audio_1.get_segment(
                from_s=start_s_1, to_s=start_s_1 + self.input_length_s
            ).array_float
            wav_2 = audio_2.get_segment(
                from_s=start_s_2, to_s=start_s_2 + self.input_length_s
            ).array_float

            # augmentation
            wav_1 = (
                self.augmentation(torch.from_numpy(wav_1).unsqueeze(0))
                .squeeze(0)
                .numpy()
            )
            wav_2 = (
                self.augmentation(torch.from_numpy(wav_2).unsqueeze(0))
                .squeeze(0)
                .numpy()
            )
        elif self.split == "valid":
            # crop first and last 30s
            start_s_1 = 0.0
            start_s_2 = 0.0
            wav_1 = audio_1.get_segment(
                from_s=start_s_1, to_s=start_s_1 + self.input_length_s
            ).array_float
            wav_2 = audio_2.get_segment(
                from_s=start_s_2, to_s=start_s_2 + self.input_length_s
            ).array_float

        # zero padding
        if len(wav_1) < int(self.sample_rate * self.input_length_s):
            pad = int(self.sample_rate * self.input_length_s) - len(wav_1)
            wav_1 = np.pad(wav_1, (0, pad), mode="constant", constant_values=0)

        if len(wav_2) < int(self.sample_rate * self.input_length_s):
            pad = int(self.sample_rate * self.input_length_s) - len(wav_2)
            wav_2 = np.pad(wav_2, (0, pad), mode="constant", constant_values=0)

        return wav_1, wav_2

    def __len__(self):
        if self.num_samples > 0:
            return self.num_samples
        else:
            return len(self.artists)