text: #!/usr/bin/env python
  """Script to generate Altered and Corrupted Midi Excerpt (ACME) datasets"""
  import argparse
  import json
  import logging
  import os
  import shutil
  import sys
  from glob import glob
  from pathlib import Path
  from zipfile import BadZipfile
  
  import numpy as np
  from tqdm import tqdm
  
  from mdtk import degradations, downloaders, fileio
  from mdtk.df_utils import get_random_excerpt
  from mdtk.formatters import FORMATTERS, create_corpus_csvs
  
  logo_path = Path(__file__, "..", "img", "logo.txt").resolve()
  with open(logo_path, "r") as ff:
      LOGO = ff.read()
  
  
  DESCRIPTION = "Make datasets of altered and corrupted midi excerpts."
  
  
  def parse_degradation_kwargs(kwarg_dict):
      """Convenience function to parse a dictionary of keyword arguments for
      the functions within the degradations module. All keys in the kwarg_dict
      supplied should start with the function name and be followed by a double
      underscore. For example: func_name__kwarg_name. An example supplied
      dictionary:
          {
              'func1__kwarg1': 7,
              'func2__kwarg1': "hello",
              'func2__kwarg2': "world"
          }
      This results in the returned dictionary:
          {
              'func1': {'kwarg1': 7},
              'func2': {'kwarg1': "hello", 'kwarg2': "world}
          }
  
      Parameters
      ----------
      kwarg_dict : dict
          Dict containing keyword arguments to parse
  
      Returns
      -------
      func_kwargs : dict
          Dict with keys matching the names of the functions. The corresponding
          value is a dictionary of the keyword arguments for the function.
      """
      func_names = degradations.DEGRADATIONS.keys()
      func_kwargs = {name: {} for name in func_names}
      if kwarg_dict is None:
          return func_kwargs
      for kk, kwarg_value in kwarg_dict.items():
          try:
              func_name, kwarg_name = kk.split("__", 1)
          except ValueError:
              raise ValueError(f"Supplied keyword [{kk}] must have a double underscore")
          assert func_name in func_names, f"{func_name} not in {func_names}"
          func_kwargs[func_name][kwarg_name] = kwarg_value
      return func_kwargs
  
  
  def clean_download_cache(dir_path=downloaders.DEFAULT_CACHE_PATH, prompt=True):
      """
      Delete the download cache directory and all files in it.
  
      Parameters
      ----------
      dir_path : str
          The path to the download cache base directory.
  
      prompt : bool
          Prompt the user before deleting.
  
      Returns
      -------
      clean_ok : bool
          True if the cache has been cleaned and deleted successfully, the cache
          dir doesn't exist, or the user chooses not to delete it at the prompt.
          False otherwise.
      """
      if not os.path.exists(dir_path):
          print(f"Download cache ({dir_path}) is clear and deleted already.")
          return True
  
      if prompt:
          response = input(f"Delete download cache ({dir_path})? [y/N]: ")
  
      if (not prompt) or response in ["y", "ye", "yes"]:
          try:
              shutil.rmtree(dir_path)
          except Exception:
              print("Could not delete download cache. Please do so manually.")
              return False
      return True
  
  
  def parse_args(args_input=None):
      """Convenience function for parsing user supplied command line args"""
      parser = argparse.ArgumentParser(
          description=DESCRIPTION, formatter_class=argparse.ArgumentDefaultsHelpFormatter
      )
      default_outdir = os.path.join(os.getcwd(), "acme")
      parser.add_argument(
          "-o",
          "--output-dir",
          type=str,
          default=default_outdir,
          help="the directory to write the dataset to.",
      )
      parser.add_argument(
          "--config",
          default=None,
          help="Load a json config "
          "file, in the format created by measure_errors.py. "
          "This will override --degradations, --degradation-"
          "dist, and --clean-prop.",
      )
      parser.add_argument(
          "--formats",
          metavar="format",
          help="Create "
          "custom versions of the acme data for easier loading "
          "with our provided pytorch Dataset classes. Choices are"
          f" {list(FORMATTERS.keys())}. Specify none to avoid "
          "creation",
          nargs="*",
          default=list(FORMATTERS.keys()),
      )
      parser.add_argument(
          "--local-midi-dirs",
          metavar="midi_dir",
          type=str,
          nargs="*",
          help="directories containing midi files to include in the dataset",
          default=[],
      )
      parser.add_argument(
          "--local-csv-dirs",
          metavar="csv_dir",
          type=str,
          nargs="*",
          help="directories containing csv files to include in the dataset",
          default=[],
      )
      parser.add_argument(
          "--recursive",
          action="store_true",
          help="Search local"
          " dataset directories recursively for all midi or csv "
          "files.",
      )
      parser.add_argument(
          "--datasets",
          metavar="dataset_name",
          nargs="*",
          default=downloaders.DATASETS,
          help="datasets to download and use. Must match names "
          "of classes in the downloaders module. By default, "
          "will use cached downloaded data if available. To clear "
          "the cache, run the script with the --clean flag. To "
          'download no data, provide an input of "None"',
      )
      parser.add_argument(
          "--degradations",
          metavar="deg_name",
          nargs="*",
          choices=list(degradations.DEGRADATIONS.keys()),
          default=list(degradations.DEGRADATIONS.keys()),
          help="degradations to use on the data. Must match "
          "names of functions in the degradations module. By "
          "default, will use them all.",
      )
      parser.add_argument(
          "--excerpt-length",
          metavar="ms",
          type=int,
          help="The length of the excerpt (in ms) to take from "
          "each piece. The excerpt will start on a note onset "
          "and include all notes whose onset lies within this "
          "number of ms after the first note.",
          default=5000,
      )
      parser.add_argument(
          "--min-notes",
          metavar="N",
          type=int,
          default=10,
          help="The minimum number of notes required for an excerpt to be valid.",
      )
      parser.add_argument(
          "--degradation-kwargs",
          metavar="json_file_or_string",
          help="json file or json-formatted string with keyword "
          "arguments for the degradation functions. First "
          "provide the degradation name, then a double "
          "underscore, then the keyword argument name, followed "
          "by the value to use for the kwarg. e.g. "
          '`{"time_shift__align_onset": true, '
          '"pitch_shift__min_pitch": 5}`',
          default=None,
      )
      parser.add_argument(
          "--degradation-dist",
          metavar="relative_probability",
          nargs="*",
          default=None,
          help="A list of relative "
          "probabilities that each degradation will used. Must "
          "be the same length as --degradations. Defaults to a "
          "uniform distribution.",
          type=float,
      )
      parser.add_argument(
          "--clean-prop",
          type=float,
          help="The proportion of excerpts in the final dataset that should be clean.",
          default=1 / (1 + len(degradations.DEGRADATIONS)),
      )
      parser.add_argument(
          "--splits",
          metavar=("train", "valid", "test"),
          nargs=3,
          type=float,
          help="The relative sizes of the "
          "train, validation, and test sets respectively.",
          default=[0.8, 0.1, 0.1],
      )
      parser.add_argument(
          "--seed",
          type=int,
          default=None,
          help="The numpy seed to use when creating the dataset.",
      )
      parser.add_argument(
          "--clean",
          action="store_true",
          help="Clear and delete"
          f" the download cache {downloaders.DEFAULT_CACHE_PATH}"
          " (and do nothing else).",
      )
      parser.add_argument(
          "-v", "--verbose", action="store_true", help="Verbose printing."
      )
      parser.add_argument(
          "--no-prompt", action="store_true", help="Dont prompt user for response."
      )
      args = parser.parse_args(args=args_input)
      return args
  
  
  if __name__ == "__main__":
      ARGS = parse_args()
  
      if ARGS.clean:
          clean_ok = clean_download_cache(
              downloaders.DEFAULT_CACHE_PATH, prompt=not ARGS.no_prompt
          )
          sys.exit(0 if clean_ok else 1)
  
      if ARGS.seed is None:
          seed = np.random.randint(0, 2 ** 32)
          print(f"No random seed supplied. Setting to {seed}.")
      else:
          seed = ARGS.seed
          print(f"Setting random seed to {seed}.")
      np.random.seed(seed)
  
      # Load given degradation_kwargs
      degradation_kwargs = {}
      if ARGS.degradation_kwargs is not None:
          if os.path.exists(ARGS.degradation_kwargs):
              # If file exists, assume that is what was passed
              with open(ARGS.degradation_kwargs, "r") as json_file:
                  degradation_kwargs = json.load(json_file)
          else:
              # File doesn't exist, assume json string was passed
              degradation_kwargs = json.loads(ARGS.degradation_kwargs)
      degradation_kwargs = parse_degradation_kwargs(degradation_kwargs)
      if ARGS.verbose:
          print(f"Using degradation kwargs: {degradation_kwargs}")
  
      # Load config
      if ARGS.config is not None:
          with open(ARGS.config, "r") as file:
              config = json.load(file)
          if ARGS.verbose:
              print(f"Loading from config file {ARGS.config}.")
          if "degradation_dist" in config:
              ARGS.degradation_dist = np.array(config["degradation_dist"])
              ARGS.degradations = list(degradations.DEGRADATIONS.keys())
          if "clean_prop" in config:
              ARGS.clean_prop = config["clean_prop"]
      # Warn user they specified kwargs for degradation not being used
      for deg, args in degradation_kwargs.items():
          if deg not in ARGS.degradations and len(args) > 0:
              logging.warning(
                  f'--degradation_kwargs contains args for unused degradation "{deg}".'
              )
  
      # Exit if degradation-dist is a diff length to degradations
      if ARGS.degradation_dist is None:
          ARGS.degradation_dist = np.ones(len(ARGS.degradations))
      assert len(ARGS.degradation_dist) == len(ARGS.degradations), (
          "Given degradation_dist is not the same length as degradations:\n"
          f"len({ARGS.degradation_dist}) != len({ARGS.degradations})"
      )
  
      # Check that no probabilities are invalid
      assert (
          min(ARGS.degradation_dist) >= 0
      ), "--degradation-dist values must not be negative."
      assert (
          sum(ARGS.degradation_dist) > 0
      ), "Some --degradation-dist value must be positive."
      assert (
          0 <= ARGS.clean_prop <= 1
      ), "--clean-prop must be between 0 and 1 (inclusive)."
      assert min(ARGS.splits) >= 0, "--splits values must not be negative."
      assert sum(ARGS.splits) > 0, "Some --splits value must be positive."
  
      # Parse formats
      if len(ARGS.formats) == 1 and ARGS.formats[0].lower() == "none":
          formats = []
      else:
          assert all([name in FORMATTERS.keys() for name in ARGS.formats]), (
              f"all provided formats {ARGS.formats} must be in the "
              "list of available formats {list(FORMATTERS.keys())}"
          )
          formats = ARGS.formats
  
      # These will be tuples of (dataset, relative_path, full_path, note_df).
      # dataset: The name of the dataset the note_df is drawn from. This will be
      #          the excerpt's base directory within output/clean or
      #          output/altered in the generated ACME dataset.
      # relative_path: The relative path of the corresponding file, including
      #                basename, representing the excerpts path within its
      #                dataset base directory.
      # full_path: The full path to the input file. Used for printing errors.
      # note_df: The cleaned note_df read from the input file with the given
      #          input_kwargs.
      # This list will be sorted before shuffling, and the tuples are structured
      # in such a way that the sorting is identical to previous versions of this
      # script to ensure backwards compatability.
      input_data = []
      input_kwargs = {"single_track": True, "non_overlapping": True}
  
      # Instantiate downloaders =================================================
      OVERWRITE = None
      ds_names = ARGS.datasets
      if len(ds_names) == 1 and ds_names[0].lower() == "none":
          ds_names = []
      else:
          assert all([name in downloaders.DATASETS for name in ds_names]), (
              f"all provided dataset names {ds_names} must be in the "
              "list of available datasets for download "
              f"{downloaders.DATASETS}"
          )
  
      if not ds_names and not ARGS.local_csv_dirs and not ARGS.local_midi_dirs:
          raise ValueError(
              "You must provide one of --datasets, --local-csv-dirs, or --local-midi-dirs"
          )
      # Instantiated downloader classes
      downloader_dict = {
          ds_name: getattr(downloaders, ds_name)(
              cache_path=downloaders.DEFAULT_CACHE_PATH
          )
          for ds_name in ds_names
      }
  
      # Clear and set up output dir =============================================
      if os.path.exists(ARGS.output_dir):
          if ARGS.verbose:
              print(f"Clearing stale data from {ARGS.output_dir}.")
  
          response = None
          if not ARGS.no_prompt:
              response = input(f"Delete output_dir ({ARGS.output_dir})? [y/N]: ")
  
          if (ARGS.no_prompt) or response in ["y", "ye", "yes"]:
              try:
                  shutil.rmtree(ARGS.output_dir)
              except Exception:
                  print("Could not delete output dir. Please do so manually.")
                  sys.exit(1)
          else:
              print(
                  "You must specify an empty directory as --output-dir, or specify a "
                  "path which doesn't yet exist"
              )
              sys.exit(1)
  
      os.makedirs(ARGS.output_dir, exist_ok=True)
      for out_subdir in ["clean", "altered"]:
          output_dirs = [
              os.path.join(ARGS.output_dir, out_subdir, name) for name in ds_names
          ]
          for path in output_dirs:
              os.makedirs(path, exist_ok=True)
  
      # Load data from downloaders ==============================================
      print("Loading data from downloaders, this could take a while...")
      for dataset in downloader_dict:
          downloader = downloader_dict[dataset]
          dataset_base = os.path.join(downloaders.DEFAULT_CACHE_PATH, dataset)
          dataset_base_len = len(dataset_base) + len(os.path.sep)
          output_path = os.path.join(dataset_base, "data")
  
          try:
              try:
                  downloader.download_csv(
                      output_path=output_path, overwrite=OVERWRITE, verbose=ARGS.verbose
                  )
                  ext = "csv"
                  input_func = fileio.csv_to_df
              except NotImplementedError:
                  downloader.download_midi(
                      output_path=output_path, overwrite=OVERWRITE, verbose=ARGS.verbose
                  )
                  ext = "mid"
                  input_func = fileio.midi_to_df
          except BadZipfile:
              print(
                  "The download cache contains invalid data. Run "
                  "`make_dataset.py --clean` to clean the cache, then try to "
                  "re-create the ACME dataset.",
                  file=sys.stderr,
              )
              sys.exit(1)
  
          for filename in tqdm(
              glob(os.path.join(output_path, "**", f"*.{ext}"), recursive=True),
              desc=f"Loading data from {dataset}",
          ):
  
              note_df = input_func(filename, **input_kwargs)
              if note_df is not None:
                  rel_path = filename[(dataset_base_len + 5) :]
                  input_data.append((dataset, rel_path, filename, note_df))
  
      # Load user data ==========================================================
      for data_type in ["midi", "csv"]:
          if data_type == "midi":
              local_dirs = ARGS.local_midi_dirs
              ext = "mid"
              df_load_func = fileio.midi_to_df
          else:
              local_dirs = ARGS.local_csv_dirs
              ext = "csv"
              df_load_func = fileio.csv_to_df
  
          for path in local_dirs:
              # Bugfix for paths ending in /
              if len(path) > 1 and path[-1] == os.path.sep:
                  path = path[:-1]
              dataset = f"local_{os.path.basename(path)}"
              dataset_base_len = len(path) + len(os.path.sep)
  
              if ARGS.recursive:
                  path = os.path.join(path, "**")
              for filepath in tqdm(
                  glob(os.path.join(path, f"*.{ext}"), recursive=ARGS.recursive),
                  desc=f"Loading user {data_type} from {path}",
              ):
  
                  note_df = df_load_func(filepath, **input_kwargs)
                  if note_df is not None:
                      rel_path = filepath[dataset_base_len:]
                      input_data.append((dataset, rel_path, filepath, note_df))
  
      # All data is loaded. Sort and shuffle. ===================================
      # output to output_dir/clean/dataset_name/filename.csv
      # The reason for this is we know there will be no filename duplicates
      input_data.sort()
      np.random.shuffle(input_data)  # This is important for join_notes
  
      meta_file = open(os.path.join(ARGS.output_dir, "metadata.csv"), "w")
  
      # Perform degradations and write degraded data to output ==================
      # output to output_dir/degraded/dataset_name/filename.csv
      # The reason for this is that there could be filename duplicates (as above,
      # it's assumed there shouldn't be duplicates within each dataset!), and
      # this allows for easy matching of source and target data
      deg_choices = ARGS.degradations
      goal_deg_dist = ARGS.degradation_dist
  
      # Remove any degs with goal_deg_dist == 0 and normalize
      non_zero = [i for i, p in enumerate(goal_deg_dist) if p != 0]
      deg_choices = np.array(deg_choices)[non_zero]
      goal_deg_dist = np.array(goal_deg_dist)[non_zero]
      goal_deg_dist /= np.sum(goal_deg_dist)
  
      # Add none for no degradation
      if 0 < ARGS.clean_prop < 1:
          deg_choices = np.insert(deg_choices, 0, "none")
          goal_deg_dist *= 1 - ARGS.clean_prop
          goal_deg_dist = np.insert(goal_deg_dist, 0, ARGS.clean_prop)
      elif ARGS.clean_prop == 1:
          deg_choices = np.array(["none"])
          goal_deg_dist = np.array([1])
      nr_degs = len(goal_deg_dist)
  
      # Normalize split proportions and remove 0s
      split_names = ["train", "valid", "test"]
      split_props = np.array(ARGS.splits)
      split_props /= np.sum(split_props)
      non_zero = [i for i, p in enumerate(split_props) if p != 0]
      split_names = np.array(split_names)[non_zero]
      split_props = np.array(split_props)[non_zero]
      nr_splits = len(split_names)
  
      # Write out deg_choices to degradation_ids.csv
      with open(os.path.join(ARGS.output_dir, "degradation_ids.csv"), "w") as file:
          file.write("id,degradation_name\n")
          for i, deg_name in enumerate(deg_choices):
              file.write(f"{i},{deg_name}\n")
  
      # The idea is to keep track of the current distribution of degradations
      # and then sample in reverse order of the difference between this and
      # the goal distribution. We do the same for splits.
      deg_counts = np.zeros(nr_degs)
      split_counts = np.zeros(nr_splits)
  
      meta_file.write("altered_csv_path,degraded,degradation_id,clean_csv_path,split\n")
      for i, data in enumerate(tqdm(input_data, desc="Degrading data")):
          dataset, rel_path, file_path, note_df = data
          rel_path = f"{rel_path[:-3]}csv"
          # First, get the degradation order for this iteration.
          # Get the current distribution of degradations
          if np.sum(deg_counts) == 0:  # First iteration, set to uniform
              current_deg_dist = np.ones(nr_degs) / nr_degs
              current_split_dist = np.ones(nr_splits) / nr_splits
          else:
              current_deg_dist = deg_counts / np.sum(deg_counts)
              current_split_dist = split_counts / np.sum(split_counts)
  
          # Grab an excerpt from this df
          excerpt = get_random_excerpt(
              note_df,
              min_notes=ARGS.min_notes,
              excerpt_length=ARGS.excerpt_length,
              first_onset_range=(0, 200),
              iterations=10,
          )
  
          # If no valid excerpt was found, skip this piece
          if excerpt is None:
              logging.warning(
                  "Unable to find valid excerpt from file "
                  f"{file_path}. Lengthen --excerpt-length or "
                  "lower --min-notes. Skipping.",
              )
              continue
  
          # Try degradations in reverse order of the difference between
          # their current distribution and their desired distribution.
          diffs = goal_deg_dist - current_deg_dist
          degs_sorted = sorted(zip(diffs, deg_choices, list(range(len(deg_choices)))))[
              ::-1
          ]
  
          # Calculate split in the same way (but only save the first)
          split_diffs = split_props - current_split_dist
          _, split_name, split_num = sorted(
              zip(split_diffs, split_names, list(range(nr_splits)))
          )[-1]
  
          # Make default labels for no degradation
          clean_path = os.path.join("clean", dataset, rel_path)
          altered_path = clean_path
          deg_binary = 0
  
          # Try to perform a degradation
          degraded = None
          for diff, deg_name, deg_num in degs_sorted:
              # Break for no degradation
              if deg_name == "none":
                  break
  
              # Try the degradation
              deg_fun = degradations.DEGRADATIONS[deg_name]
              deg_fun_kwargs = degradation_kwargs[deg_name]  # degradation_kwargs
              # at top of main call
              logging.disable(logging.WARNING)
              degraded = deg_fun(excerpt, **deg_fun_kwargs)
              logging.disable(logging.NOTSET)
  
              if degraded is not None:
                  # Update labels
                  deg_binary = 1
                  altered_path = os.path.join("altered", dataset, rel_path)
  
                  # Write degraded csv
                  altered_outpath = os.path.join(ARGS.output_dir, altered_path)
                  fileio.df_to_csv(degraded, altered_outpath)
                  break
  
          # Write data
          if not (degraded is None and ARGS.clean_prop == 0):
              # Update counts
              deg_counts[deg_num] += 1
              split_counts[split_num] += 1
  
              # Write clean csv
              clean_outpath = os.path.join(ARGS.output_dir, clean_path)
              fileio.df_to_csv(excerpt, clean_outpath)
  
              # Write metadata
              meta_file.write(
                  f"{altered_path},{deg_binary},{deg_num}," f"{clean_path},{split_name}\n"
              )
          else:
              logging.warning(
                  "Unable to degrade chosen excerpt from "
                  f"{file_path} and no clean excerpts requested."
                  " Skipping.",
              )
  
      meta_file.close()
  
      for f in formats:
          create_corpus_csvs(ARGS.output_dir, FORMATTERS[f])
  
      print(f'\n{10*"="} Finished! {10*"="}\n')
      print("Count of degradations:")
      for deg_name, count in zip(deg_choices, deg_counts):
          print(f"\t* {deg_name}: {int(count)}")
  
      print(
          f"\nYou will find the generated data at {ARGS.output_dir} "
          "with subdirectories"
      )
      print("\t* clean - contains the extracted clean excerpts")
      print(
          "\t* altered - contains the excerpts altered by the degradations "
          "described in metadata.csv"
      )
      print("\nmetadata.csv describes:")
      print("\t* (the id number for) the type of degradation used for the alteration")
      print("\t* the path for the altered and clean files")
      print("\t* which split (train, valid, test) the file should be used in")
      print("\t* in which corpus and on what line the file is located")
  
      print(
          "\ndegradation_ids.csv is a mapping of degradation name to the id "
          "number used in metadata.csv"
      )
  
      for f in formats:
          print(f"\n{FORMATTERS[f]['message']}")
  
      print(
          "\nTo reproduce this dataset again, run the script with argument "
          f"--seed {seed}"
      )
      print(LOGO)
  