import os
import sys
import numpy as np
import pickle
sys.path.append('../../')

from src.music.utils import  get_out_path
from src.music.config import ALL_NOISE, ALL_AUGMENTATIONS, NB_AUG, NOISE_INJECTED
from src.music.utilities.midi_processor import encode_midi_structured, encode_midi_chunks_structured

nb_noise = ALL_NOISE.shape[0]
nb_aug = ALL_AUGMENTATIONS.shape[0]

def sample_augmentations(n):
    return ALL_AUGMENTATIONS[np.random.choice(np.arange(nb_aug), size=n, replace=False)]

def sample_noise():
    return ALL_NOISE[np.random.choice(np.arange(nb_noise))]

def processed2encoded(processed_path, encoded_path=None, augmentation=False, nb_aug=None, noise_injection=False, verbose=False, level=0):
    assert processed_path.split('.')[-1] in ['mid', 'midi']
    if not encoded_path:
        encoded_path, _, _ = get_out_path(in_path=processed_path, in_word='processed', out_word='encoded', out_extension='.pickle')

    if verbose: print(' ' * level + f'Encoding {processed_path}')
    if os.path.exists(encoded_path):
        if verbose: print(' ' * (level + 2) + 'Midi file is already encoded.')
        return encoded_path, ''

    if augmentation:
        assert isinstance(nb_aug, int)
    error_msg = 'Error in encoding. '
    try:
        error_msg = 'Error in encoding midi?'
        nb_noise = 1 if noise_injection else 0
        encoded_main, encoded_aug, encoded_noisy = encode_midi_structured(processed_path, nb_aug, nb_noise)

        # make sure augmentations are not out of bounds
        error_msg = ' Nope. Error in saving encoding?'
        with open(encoded_path, 'wb') as f:
            pickle.dump(dict(main=encoded_main, aug=encoded_aug, noisy=encoded_noisy), f)
        error_msg = ' Nope.'
        if verbose:
            extra = f' Saved to {encoded_path}' if encoded_path else ''
            print(' ' * (level + 2) + f'Success! {extra}')
        return encoded_path, ''
    except:
        if verbose: print(' ' * (level + 2) + 'Transcription failed.')
        if os.path.exists(encoded_path):
            os.remove(encoded_path)
        return None, error_msg + ' Yes.'