#! /usr/bin/env python3

# this executes the chord extractor from polyffusion on a list of midi files
import sys 
import os
import sys
import pretty_midi
from multiprocessing import Pool, cpu_count, set_start_method
from tqdm import tqdm
from control_toys.data import fast_scandir
from functools import partial
import argparse

chord_extractor_path = '/Users/shawley/github/polyffusion/polyffusion/chord_extractor'
# add chord extract path to python sys
sys.path.append(chord_extractor_path)

from chord_class import ChordClass

#try:
from chord_extractor.main import transcribe_cb1000_midi
chord_extractor_available = True
#except:
#    chord_extractor_available = False


def extract_chords(args, midi_file):
    "this will read in one midi file and estimate the chords"
    chords_path = midi_file.replace('.mid', '_chords.txt')
    try:
        transcribe_cb1000_midi(midi_file, chords_path)
    except Exception as e:
        print(f"\nError processing {midi_file}: {e}. Skipping")



if __name__ == '__main__':
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    p.add_argument('--start-method', type=str, default='fork',
                   choices=['fork', 'forkserver', 'spawn'],
                   help='the multiprocessing start method')
    p.add_argument('--skip-versions', default=True, help='skip extra versions of the same song')
    p.add_argument("midi_dirs", nargs='+', help="directories containing MIDI files")
    args = p.parse_args()
    print("args = ",args)

    if not chord_extractor_available:
        print("Error: chord_extractor not available. Please install it, e.g. via 'ln -s ~/diffusion/polyffusion/polyffusion/chord_extractor .'")
        sys.exit(1)

    set_start_method(args.start_method)
    midi_dirs = args.midi_dirs

    if os.path.isdir(midi_dirs[0]):
        midi_files = []
        for mdir in midi_dirs:
            m_subdirs, mf = fast_scandir(mdir, ['mid', 'midi'])
            if mf != []: midi_files = midi_files + mf
    elif os.path.isfile(midi_dirs[0]):
        midi_files = midi_dirs

    if args.skip_versions: 
        midi_files = [f for f in midi_files if '/versions/' not in f]
    #print("midi_files = ",midi_files) # just a check for debugging
        
    extract_one = partial(extract_chords, args)
    with Pool(cpu_count()) as p:
        # TODO: how to get functions called from p.imap to return values?
        list(tqdm(p.imap(extract_one, midi_files), total=len(midi_files), desc='Extracting chords from MIDI files'))

    print("Finished")