# coding=utf-8
# Copyright 2023 The TensorFlow Datasets Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""NSynth Dataset."""

from __future__ import annotations

import csv
import os

import numpy as np
from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
import tensorflow_datasets.public_api as tfds

_FULL_DESCRIPTION = """\
Full NSynth Dataset is split into train, valid, and test sets, with no
instruments overlapping between the train set and the valid/test sets.
"""

_GANSYNTH_DESCRIPTION = """\
NSynth Dataset limited to acoustic instruments in the MIDI pitch interval
[24, 84]. Uses alternate splits that have overlap in instruments (but not exact
notes) between the train set and valid/test sets. This variant was originally
introduced in the ICLR 2019 GANSynth paper (https://arxiv.org/abs/1902.08710).
"""

_F0_AND_LOUDNESS_ADDENDUM = """\
This version additionally contains estimates for F0 using CREPE
(Kim et al., 2018) and A-weighted perceptual loudness in decibels. Both signals
are provided at a frame rate of 250Hz.
"""

# From http://proceedings.mlr.press/v70/engel17a.html

_NUM_SECS = 4
_AUDIO_RATE = 16000  # 16 kHz
_F0_AND_LOUDNESS_RATE = 250  # 250 Hz
_CREPE_FRAME_SIZE = 1024
_LD_N_FFT = 2048
_LD_RANGE = 120.0
_REF_DB = 20.7  # White noise, amplitude=1.0, n_fft=2048

_INSTRUMENT_FAMILIES = [
    "bass",
    "brass",
    "flute",
    "guitar",
    "keyboard",
    "mallet",
    "organ",
    "reed",
    "string",
    "synth_lead",
    "vocal",
]
_INSTRUMENT_SOURCES = ["acoustic", "electronic", "synthetic"]
_QUALITIES = [
    "bright",
    "dark",
    "distortion",
    "fast_decay",
    "long_release",
    "multiphonic",
    "nonlinear_env",
    "percussive",
    "reverb",
    "tempo-synced",
]

_BASE_DOWNLOAD_PATH = (
    "http://download.magenta.tensorflow.org/datasets/nsynth/nsynth-"
)

_SPLITS = ["train", "valid", "test"]


class NsynthConfig(tfds.core.BuilderConfig):
  """BuilderConfig for NSynth Dataset."""

  def __init__(
      self, gansynth_subset=False, estimate_f0_and_loudness=False, **kwargs
  ):
    """Constructs a NsynthConfig.

    Args:
      gansynth_subset: bool, whether to use the subset of the dataset introduced
        in the ICLR 2019 GANSynth paper (Engel, et al. 2018). This subset uses
        acoustic-only instrument sources and limits the pitches to the interval
        [24, 84]. The train and test splits are also modified so that
        instruments (but not specific notes) overlap between them. See
        https://arxiv.org/abs/1902.08710 for more details.
      estimate_f0_and_loudness: bool, whether to estimate fundamental frequency
        (F0) and loudness for the audio (at 250 Hz) and add them to the set of
        features.
      **kwargs: keyword arguments forwarded to super.
    """
    name_parts = []
    if gansynth_subset:
      name_parts.append("gansynth_subset")
    else:
      name_parts.append("full")
    if estimate_f0_and_loudness:
      name_parts.append("f0_and_loudness")
    v230 = tfds.core.Version("2.3.0")
    v231 = tfds.core.Version("2.3.1")
    v232 = tfds.core.Version("2.3.2")
    v233 = tfds.core.Version("2.3.3")
    super(NsynthConfig, self).__init__(
        name=".".join(name_parts),
        version=v233,
        supported_versions=[v232, v231, v230],
        release_notes={
            "2.3.3": (
                "F0 computed with fix in CREPE wave normalization "
                "(https://github.com/marl/crepe/issues/49)."
            ),
            "2.3.2": "Use Audio feature.",
            "2.3.1": "F0 computed with normalization fix in CREPE.",
            "2.3.0": "New `loudness_db` feature in decibels (unormalized).",
        },
        **kwargs,
    )  # pytype: disable=wrong-arg-types  # gen-stub-imports
    self.gansynth_subset = gansynth_subset
    self.estimate_f0_and_loudness = estimate_f0_and_loudness


class Builder(tfds.core.BeamBasedBuilder):
  """A large-scale and high-quality dataset of annotated musical notes."""

  BUILDER_CONFIGS = [
      NsynthConfig(description=_FULL_DESCRIPTION),
      NsynthConfig(gansynth_subset=True, description=_GANSYNTH_DESCRIPTION),
      NsynthConfig(
          gansynth_subset=True,
          estimate_f0_and_loudness=True,
          description=_GANSYNTH_DESCRIPTION + _F0_AND_LOUDNESS_ADDENDUM,
      ),
  ]

  def _info(self):
    features = {
        "id": tf.string,
        "audio": tfds.features.Audio(
            shape=(_AUDIO_RATE * _NUM_SECS,),
            dtype=np.float32,
            sample_rate=_AUDIO_RATE,
        ),
        "pitch": tfds.features.ClassLabel(num_classes=128),
        "velocity": tfds.features.ClassLabel(num_classes=128),
        "instrument": {
            # We read the list of labels in _split_generators.
            "label": tfds.features.ClassLabel(num_classes=1006),
            "family": tfds.features.ClassLabel(names=_INSTRUMENT_FAMILIES),
            "source": tfds.features.ClassLabel(names=_INSTRUMENT_SOURCES),
        },
        "qualities": {quality: tf.bool for quality in _QUALITIES},
    }
    if self.builder_config.estimate_f0_and_loudness:
      f0_and_ld_shape = (_F0_AND_LOUDNESS_RATE * _NUM_SECS,)
      features["f0"] = {
          "hz": tfds.features.Tensor(shape=f0_and_ld_shape, dtype=np.float32),
          "midi": tfds.features.Tensor(shape=f0_and_ld_shape, dtype=np.float32),
          "confidence": tfds.features.Tensor(
              shape=f0_and_ld_shape, dtype=np.float32
          ),
      }
      features["loudness"] = {
          "db": tfds.features.Tensor(shape=f0_and_ld_shape, dtype=np.float32)
      }
    return self.dataset_info_from_configs(
        features=tfds.features.FeaturesDict(features),
        homepage="https://g.co/magenta/nsynth-dataset",
        metadata=tfds.core.BeamMetadataDict(
            sample_rate=_AUDIO_RATE,
        ),
    )

  def _split_generators(self, dl_manager):
    """Returns splits."""

    dl_urls = {}
    dl_urls["examples"] = {
        split: _BASE_DOWNLOAD_PATH + "%s.tfrecord.tar" % split
        for split in _SPLITS
    }
    dl_urls["instrument_labels"] = _BASE_DOWNLOAD_PATH + "instrument_labels.txt"
    if self.builder_config.gansynth_subset:
      dl_urls["gansynth_splits"] = _BASE_DOWNLOAD_PATH + "gansynth_splits.csv"
    dl_paths = dl_manager.download_and_extract(dl_urls)

    with tf.io.gfile.GFile(dl_paths["instrument_labels"]) as f:
      instrument_labels = f.read().strip().splitlines()
    self.info.features["instrument"]["label"].names = instrument_labels

    split_ids = {s: set() for s in _SPLITS}
    split_dirs = {s: [dl_paths["examples"][s]] for s in _SPLITS}
    if self.builder_config.gansynth_subset:
      # Generator needs to see all original splits for each new split.
      split_dirs = {s: dl_paths["examples"].values() for s in _SPLITS}
      with tf.io.gfile.GFile(dl_paths["gansynth_splits"]) as f:
        reader = csv.DictReader(f)
        for row in reader:
          split_ids[row["split"]].add(row["id"])

    return [
        tfds.core.SplitGenerator(  # pylint: disable=g-complex-comprehension
            name=split,
            gen_kwargs={
                "tfrecord_dirs": split_dirs[split],
                "ids": split_ids[split],
                "split": split,
            },
        )
        for split in _SPLITS
    ]

  def _build_pcollection(self, pipeline, tfrecord_dirs, ids, split):
    """Build PCollection of examples for split."""
    beam = tfds.core.lazy_imports.apache_beam

    def _emit_base_example(ex):
      """Maps an input example to a TFDS example."""
      beam.metrics.Metrics.counter(split, "base-examples").inc()
      features = ex.features.feature
      id_ = features["note_str"].bytes_list.value[0]
      return id_, {
          "id": id_,
          "audio": np.array(
              features["audio"].float_list.value, dtype=np.float32
          ),
          "pitch": features["pitch"].int64_list.value[0],
          "velocity": features["velocity"].int64_list.value[0],
          "instrument": {
              "label": tf.compat.as_text(
                  features["instrument_str"].bytes_list.value[0]
              ),
              "family": tf.compat.as_text(
                  features["instrument_family_str"].bytes_list.value[0]
              ),
              "source": tf.compat.as_text(
                  features["instrument_source_str"].bytes_list.value[0]
              ),
          },
          "qualities": {
              q: features["qualities"].int64_list.value[i]
              for (i, q) in enumerate(_QUALITIES)
          },
      }

    def _in_split(id_ex, split_ids):
      unused_id, ex = id_ex
      if not split_ids or tf.compat.as_text(ex["id"]) in split_ids:
        beam.metrics.Metrics.counter(split, "in-split").inc()
        return True
      return False

    def _estimate_f0(id_ex):
      """Estimate the fundamental frequency using CREPE and add to example."""
      id_, ex = id_ex
      beam.metrics.Metrics.counter(split, "estimate-f0").inc()

      audio = ex["audio"]

      # Copied from magenta/ddsp/spectral_ops.py
      # Pad end so that `num_frames = _NUM_SECS * _F0_AND_LOUDNESS_RATE`.
      hop_size = _AUDIO_RATE / _F0_AND_LOUDNESS_RATE
      n_samples = len(audio)
      n_frames = _NUM_SECS * _F0_AND_LOUDNESS_RATE
      n_samples_padded = (n_frames - 1) * hop_size + _CREPE_FRAME_SIZE
      n_padding = n_samples_padded - n_samples
      assert n_padding % 1 == 0
      audio = np.pad(audio, (0, int(n_padding)), mode="constant")
      crepe_step_size = 1000 / _F0_AND_LOUDNESS_RATE  # milliseconds

      _, f0_hz, f0_confidence, _ = tfds.core.lazy_imports.crepe.predict(
          audio,
          sr=_AUDIO_RATE,
          viterbi=True,
          step_size=crepe_step_size,
          center=False,
          verbose=0,
      )
      f0_midi = tfds.core.lazy_imports.librosa.core.hz_to_midi(f0_hz)
      # Set -infs introduced by hz_to_midi to 0.
      f0_midi[f0_midi == -np.inf] = 0
      # Set nans to 0 in confidence.
      f0_confidence = np.nan_to_num(f0_confidence)
      ex = dict(ex)
      ex["f0"] = {
          "hz": f0_hz.astype(np.float32),
          "midi": f0_midi.astype(np.float32),
          "confidence": f0_confidence.astype(np.float32),
      }
      return id_, ex

    def _calc_loudness(id_ex):
      """Compute loudness, add to example (ref is white noise, amplitude=1)."""
      id_, ex = id_ex
      beam.metrics.Metrics.counter(split, "compute-loudness").inc()

      audio = ex["audio"]

      # Copied from magenta/ddsp/spectral_ops.py
      # Get magnitudes.
      hop_size = int(_AUDIO_RATE // _F0_AND_LOUDNESS_RATE)

      # Add padding to the end
      n_samples_initial = int(audio.shape[-1])
      n_frames = int(np.ceil(n_samples_initial / hop_size))
      n_samples_final = (n_frames - 1) * hop_size + _LD_N_FFT
      pad = n_samples_final - n_samples_initial
      audio = np.pad(audio, ((0, pad),), "constant")

      librosa = tfds.core.lazy_imports.librosa
      spectra = librosa.stft(
          audio, n_fft=_LD_N_FFT, hop_length=hop_size, center=False
      ).T

      # Compute power
      amplitude = np.abs(spectra)
      amin = 1e-20  # Avoid log(0) instabilities.
      power_db = np.log10(np.maximum(amin, amplitude))
      power_db *= 20.0

      # Perceptual weighting.
      frequencies = librosa.fft_frequencies(sr=_AUDIO_RATE, n_fft=_LD_N_FFT)
      a_weighting = librosa.A_weighting(frequencies)[np.newaxis, :]
      loudness = power_db + a_weighting

      # Set dynamic range.
      loudness -= _REF_DB
      loudness = np.maximum(loudness, -_LD_RANGE)

      # Average over frequency bins.
      mean_loudness_db = np.mean(loudness, axis=-1)

      ex = dict(ex)
      ex["loudness"] = {"db": mean_loudness_db.astype(np.float32)}
      return id_, ex

    examples = (
        pipeline
        | beam.Create([os.path.join(dir_, "*") for dir_ in tfrecord_dirs])
        | beam.io.tfrecordio.ReadAllFromTFRecord(
            coder=beam.coders.ProtoCoder(tf.train.Example)
        )
        | beam.Map(_emit_base_example)
        | beam.Filter(_in_split, split_ids=ids)
    )
    if self.builder_config.estimate_f0_and_loudness:
      examples = (
          examples
          | beam.Reshuffle()
          | beam.Map(_estimate_f0)
          | beam.Map(_calc_loudness)
      )

    return examples
