import time
import os
import sys
sys.path.append('../../')
import pretty_midi as pm
import numpy as np

from src.music.utils import  get_out_path
from src.music.config import MIN_LEN, MIN_NB_NOTES, MAX_GAP_IN_SONG, REMOVE_FIRST_AND_LAST


def sort_notes(notes):
    starts = np.array([n.start for n in notes])
    index_sorted = np.argsort(starts)
    return [notes[i] for i in index_sorted].copy()


def delete_notes_end_after_start(notes):
    indexes_to_keep = [i for i, n in enumerate(notes) if n.start < n.end]
    return [notes[i] for i in indexes_to_keep].copy()

def compute_largest_gap(notes):
    gaps = []
    latest_note_end_so_far = notes[0].end
    for i in range(len(notes) - 1):
        note_start = notes[i + 1].start
        if latest_note_end_so_far < note_start:
            gaps.append(note_start - latest_note_end_so_far)
        latest_note_end_so_far = max(latest_note_end_so_far, notes[i+1].end)
    if len(gaps) > 0:
        largest_gap = np.max(gaps)
    else:
        largest_gap = 0
    return largest_gap

def analyze_instrument(inst):
    # test that piano plays throughout
    init = time.time()
    notes = inst.notes.copy()
    nb_notes = len(notes)
    start = notes[0].start
    end = inst.get_end_time()
    duration = end - start
    largest_gap = compute_largest_gap(notes)
    return nb_notes, start, end, duration, largest_gap

def remove_beginning_and_end(midi, end_time):
    notes = midi.instruments[0].notes.copy()
    new_notes = [n for n in notes if n.start > REMOVE_FIRST_AND_LAST and n.end < end_time - REMOVE_FIRST_AND_LAST]
    midi.instruments[0].notes = new_notes
    return midi

def remove_blanks_beginning_and_end(midi):
    # remove blanks and the beginning and the end
    shift = midi.instruments[0].notes[0].start
    for n in midi.instruments[0].notes:
        n.start = max(0, n.start - shift)
        n.end = max(0, n.end - shift)
    for ksc in midi.key_signature_changes:
        ksc.time = max(0, ksc.time - shift)
    for tsc in midi.time_signature_changes:
        tsc.time = max(0, tsc.time - shift)
    for pb in midi.instruments[0].pitch_bends:
        pb.time = max(0, pb.time - shift)
    for cc in midi.instruments[0].control_changes:
        cc.time = max(0, cc.time - shift)
    return midi

def is_valid_inst(largest_gap, duration, nb_notes, gap_counts=True):
    error_msg = ''
    valid = True
    if largest_gap > MAX_GAP_IN_SONG and gap_counts:
        valid = False
        error_msg += f'wide gap ({largest_gap:.2f} secs), '
    if duration < (MIN_LEN + 2 * REMOVE_FIRST_AND_LAST):
        valid = False
        error_msg += f'too short ({duration:.2f} secs), '
    if nb_notes < MIN_NB_NOTES * duration / 60:  # nb of notes needs to be superior to the minimum number / min * the duration in minute
        valid = False
        error_msg += f'too few notes ({nb_notes}), '
    return valid, error_msg

def midi2processed(midi_path, processed_path=None, apply_filtering=True, verbose=False, level=0):
    assert midi_path.split('.')[-1] in ['mid', 'midi']
    if not processed_path:
        processed_path, _, _ = get_out_path(in_path=midi_path, in_word='midi', out_word='processed', out_extension='.mid')

    if verbose: print(' ' * level + f'Processing {midi_path}.')

    if os.path.exists(processed_path):
        if verbose: print(' ' * (level + 2) + 'Processed midi file already exists.')
        return processed_path, ''
    error_msg = 'Error in scrubbing. '
    try:
        inst_error_msg = ''
        # load mid file
        error_msg += 'Error in midi loading?'
        midi = pm.PrettyMIDI(midi_path)
        error_msg += ' Nope. Removing invalid notes?'
        midi.remove_invalid_notes()  # filter invalid notes
        error_msg += ' Nope. Filtering instruments?'
        # filter instruments
        instruments = midi.instruments.copy()
        new_instru = []
        instruments_data = []
        for i_inst, inst in enumerate(instruments):
            if inst.program <= 7 and not inst.is_drum and len(inst.notes) > 5:
                # inst is a piano
                # check data
                inst.notes = sort_notes(inst.notes)  # sort notes
                inst.notes = delete_notes_end_after_start(inst.notes)  # delete invalid notes
                nb_notes, start, end, duration, largest_gap = analyze_instrument(inst)
                is_valid, err_msg = is_valid_inst(largest_gap=largest_gap, duration=duration, nb_notes=nb_notes, gap_counts='maestro' not in midi_path)
                if is_valid or not apply_filtering:
                    new_instru.append(inst)
                    instruments_data.append([nb_notes, start, end, duration, largest_gap])
                else:
                    inst_error_msg += 'inst1: ' + err_msg + '\n'
        instruments_data = np.array(instruments_data)
        error_msg += ' Nope. Taking one instrument?'

        if len(new_instru) == 0:
            error_msg = f'No piano instrument. {inst_error_msg}'
            assert False
        elif len(new_instru) > 1:
            # take instrument playing the most notes
            instrument = new_instru[np.argmax(instruments_data[:, 0])]
        else:
            instrument = new_instru[0]
        instrument.program = 0  # set the instrument to Grand Piano.
        midi.instruments = [instrument]  # put instrument in midi file
        error_msg += ' Nope. Removing blanks?'
        # remove first and last REMOVE_FIRST_AND_LAST seconds (avoid clapping and jingles)
        end_time = midi.get_end_time()
        if apply_filtering: midi = remove_beginning_and_end(midi, end_time)

        # remove beginning and end
        midi = remove_blanks_beginning_and_end(midi)
        error_msg += ' Nope. Saving?'

        # save midi file
        midi.write(processed_path)
        error_msg += ' Nope.'
        if verbose:
            extra = f' Saved to {processed_path}' if midi_path else ''
            print(' ' * (level + 2) + f'Success! {extra}')
        return processed_path, ''
    except:
        if verbose: print(' ' * (level + 2) + 'Scrubbing failed.')
        if os.path.exists(processed_path):
            os.remove(processed_path)
        return None, error_msg + ' Yes.'