import os
from glob import glob
from torch import Tensor
from typing import Tuple
import subprocess
import torchaudio
from abc import abstractmethod
from clmr.datasets import Dataset
import random


def preprocess_audio(source, target, sample_rate):
    p = subprocess.Popen(
        ["ffmpeg", "-i", source, "-ac", "1", "-ar", str(sample_rate), target, "-loglevel", "quiet"]
    )
    p.wait()


class HUM2SONG(Dataset):
    """Create a Dataset for any folder of audio files.
    Args:
        root (str): Path to the directory where the dataset is found or downloaded.
        src_ext_audio (str): The extension of the audio files to analyze.
    """

    def __init__(
        self,
        root: str,
        src_ext_audio: str = ".mp3",
        n_classes: int = 1,
    ) -> None:
        super(HUM2SONG, self).__init__(root)

        self._path = root
        self._src_ext_audio = src_ext_audio
        self.n_classes = n_classes

        self.hum = sorted(glob(
            os.path.join(self._path, "hum", "**", "*{}".format(self._src_ext_audio)),
            recursive=True,
        ))
        self.song = sorted(glob(
            os.path.join(self._path, "song", "**", "*{}".format(self._src_ext_audio)),
            recursive=True,
        ))
        n_hum = len(self.hum)
        self.label = [1.0] * n_hum
        for i in range(n_hum):
            h = self.hum[i]
            self.hum.append(h)
            ran_songs = random.sample(self.song, k=10)
            for s in ran_songs:
                if s[s.rfind('/')+1:] != h[h.rfind('/')+1:]:
                    self.song.append(s)
                    self.label.append(0)

        if len(self.hum) == 0:
            raise RuntimeError(
                "Dataset not found. Please place the audio files in the {} folder.".format(
                    self._path
                )
            )

    def file_path(self, n: int, it: int) -> str:
        if it == 0:
          fp = self.hum[n]
        elif it == 1:
          fp = self.song[n]
        return fp

    def __getitem__(self, n: int) -> Tuple[Tensor, Tensor]:
        """Load the n-th sample from the dataset.
        Args:
            n (int): The index of the sample to be loaded
        Returns:
            Tuple [Tensor, Tensor]: ``(waveform, label)``
        """
        hum, _ = self.load(n, 0)
        song, _ = self.load(n, 1)
        label = self.label[n]
        return hum, song, label

    def __len__(self) -> int:
        return len(self.hum)
      
    def target_file_path(self, n: int, it: int) -> str:
        fp = self.file_path(n, it)
        file_basename, _ = os.path.splitext(fp)
        return file_basename + self._ext_audio

    def preprocess(self, n: int, sample_rate: int):
        for it in [0, 1]:
          fp = self.file_path(n, it)
          target_fp = self.target_file_path(n, it)
          assert fp[fp.rfind('/'):] != target_fp[target_fp.rfind('/'):]

          if not os.path.exists(target_fp):
              preprocess_audio(fp, target_fp, sample_rate)

    def load(self, n, it: int):
        target_fp = self.target_file_path(n, it)
        try:
            audio, sample_rate = torchaudio.load(target_fp)
        except OSError as e:
            print("File not found, try running `python preprocess.py` first.\n\n", e)
            return
        return audio, sample_rate

import os
from glob import glob
from torch import Tensor
from typing import Tuple
import subprocess
import torchaudio
from abc import abstractmethod
from clmr.datasets import Dataset
import random


def preprocess_audio(source, target, sample_rate):
    p = subprocess.Popen(
        ["ffmpeg", "-i", source, "-ac", "1", "-ar", str(sample_rate), target, "-loglevel", "quiet"]
    )
    p.wait()


class HUM2SONG(Dataset):
    """Create a Dataset for any folder of audio files.
    Args:
        root (str): Path to the directory where the dataset is found or downloaded.
        src_ext_audio (str): The extension of the audio files to analyze.
    """

    def __init__(
        self,
        root: str,
        src_ext_audio: str = ".mp3",
        n_classes: int = 1,
    ) -> None:
        super(HUM2SONG, self).__init__(root)

        self._path = root
        self._src_ext_audio = src_ext_audio
        self.n_classes = n_classes

        self.hum = sorted(glob(
            os.path.join(self._path, "hum", "**", "*{}".format(self._src_ext_audio)),
            recursive=True,
        ))
        self.song = sorted(glob(
            os.path.join(self._path, "song", "**", "*{}".format(self._src_ext_audio)),
            recursive=True,
        ))
        n_hum = len(self.hum)
        self.label = [1.0] * n_hum
        for i in range(n_hum):
            h = self.hum[i]
            self.hum.append(h)
            ran_songs = random.sample(self.song, k=10)
            for s in ran_songs:
                if s[s.rfind('/')+1:] != h[h.rfind('/')+1:]:
                    self.song.append(s)
                    self.label.append(0)

        if len(self.hum) == 0:
            raise RuntimeError(
                "Dataset not found. Please place the audio files in the {} folder.".format(
                    self._path
                )
            )

    def file_path(self, n: int, it: int) -> str:
        if it == 0:
          fp = self.hum[n]
        elif it == 1:
          fp = self.song[n]
        return fp

    def __getitem__(self, n: int) -> Tuple[Tensor, Tensor]:
        """Load the n-th sample from the dataset.
        Args:
            n (int): The index of the sample to be loaded
        Returns:
            Tuple [Tensor, Tensor]: ``(waveform, label)``
        """
        hum, _ = self.load(n, 0)
        song, _ = self.load(n, 1)
        label = self.label[n]
        return hum, song, label

    def __len__(self) -> int:
        return len(self.hum)
      
    def target_file_path(self, n: int, it: int) -> str:
        fp = self.file_path(n, it)
        file_basename, _ = os.path.splitext(fp)
        return file_basename + self._ext_audio

    def preprocess(self, n: int, sample_rate: int):
        for it in [0, 1]:
          fp = self.file_path(n, it)
          target_fp = self.target_file_path(n, it)
          assert fp[fp.rfind('/'):] != target_fp[target_fp.rfind('/'):]

          if not os.path.exists(target_fp):
              preprocess_audio(fp, target_fp, sample_rate)

    def load(self, n, it: int):
        target_fp = self.target_file_path(n, it)
        try:
            audio, sample_rate = torchaudio.load(target_fp)
        except OSError as e:
            print("File not found, try running `python preprocess.py` first.\n\n", e)
            return
        return audio, sample_rate

Create a Dataset for any folder of audio files.     Args:         root (str): Path to the directory where the dataset is found or downloaded.         src_ext_audio (str): The extension of the audio files to analyze. Load the n-th sample from the dataset.         Args:             n (int): The index of the sample to be loaded         Returns:             Tuple [Tensor, Tensor]: ``(waveform, label)``