import os
import sys
import argparse
import pickle

from midi_processing import extract_data, analyze_data

parser = argparse.ArgumentParser()
parser.add_argument('midi_pack', help='folder name for midi pack in midi_packs folder', type=str)
parser.add_argument('name', help='name for experiment', type=str)
parser.add_argument('--b', help='lengh of sequence in bars', type=int)
parser.add_argument('-a', help='analize data', action='store_true')
args = parser.parse_args()

'''SETTINGS'''
MIDI_PACK_NAME = args.midi_pack
EXPERIMENT_NAME = args.name
BARS_IN_SEQ = args.b

midi_folder_path = os.path.join('midi_packs', MIDI_PACK_NAME)

# analyze data set for intresting intruments
if args.a:
    analyze_data(midi_folder_path)
    sys.exit()


'''MODEL WORKFLOW DIALOG'''
number_of_instruments = int(input('Please specify number of instruments\n'))
model_workflow = dict()
input_list = []
for i in range(number_of_instruments):
    input_string = input('Please specify a workflow step\n')
    tokens = input_string.split()
    if tokens[-1] == 'melody':
        model_workflow[i] = (tokens[0], tokens[1])
    else:
        model_workflow[i] = ((tokens[1], tokens[0]), tokens[2])

# make folder for new experiment if no exist
try:
  os.mkdir(os.path.join('training_sets', EXPERIMENT_NAME))
except:
   pass

# extract process
for key, (instrument, how) in model_workflow.items():
    if how == 'melody':
        instrument_name = instrument
    else:
        instrument_name = instrument[1]
  
    save_path = os.path.join('training_sets', EXPERIMENT_NAME, instrument_name.lower() + '_data.pkl')
    x_train, y_train, program = extract_data(midi_folder_path=midi_folder_path, how=how, 
                                              instrument=instrument, bar_in_seq=BARS_IN_SEQ)
    pickle.dump((x_train, y_train, program), open(save_path,'wb'))

pickle.dump(model_workflow, open(os.path.join('training_sets', EXPERIMENT_NAME, 'workflow.pkl'),'wb'))