import numpy as np
import librosa
import sys
sys.path.append('../../../data/')
from src.music.utilities.processing_models import piano_detection_model
from src.music.config import CHKPT_PATH_PIANO_EVAL

PIANO_SOLO_DETECTOR = piano_detection_model.PianoSoloDetector(CHKPT_PATH_PIANO_EVAL)
exclude_playlist_folders = ['synth_audio_recorded', 'from_url']

def clean_start_and_end_blanks(probs):
    if len(probs) > 20:
        # clean up to 10s in each direction
        n_zeros_start = 0
        for i in range(10):
            if probs[i] <= 0.001:
                n_zeros_start += 1
            else:
                break
        n_zeros_end = 0
        for i in range(10):
            if probs[-(i + 1)] <= 0.001:
                n_zeros_end += 1
            else:
                break
        if n_zeros_end == 0:
            return probs[n_zeros_start:]
        else:
            return probs[n_zeros_start:-n_zeros_end]
    else:
        return probs

def calculate_piano_solo_prob(audio_path, verbose=False):
    """Calculate the piano solo probability of all downloaded mp3s, and append
    the probability to the meta csv file. Code from https://github.com/bytedance/GiantMIDI-Piano
    """
    try:
        error_msg = 'Error in audio loading?'
        (audio, _) = librosa.core.load(audio_path, sr=piano_detection_model.SR, mono=True)
        error_msg += ' Nope. Error in solo prediction?'
        probs = PIANO_SOLO_DETECTOR.predict(audio)
        # probs = clean_start_and_end_blanks(probs)  # remove blanks at start and end (<=10s each way). If not piano, the rest of the song will be enough to tell.
        piano_solo_prob = np.mean(probs)
        error_msg += ' Nope. '
        return piano_solo_prob, ''
    except:
        return None, error_msg + 'Yes.'