import numpy as np
import mir_eval as me
import madmom

def get_score(predictions, anotace, indices, data):
    # IN: predicted beat times, annotations, file indices of given predictions
    # OUT: Score dictionary, also printed

    score = {
        "be_fscore": [],
        "be_recall": [],
        "be_precision": [],
        "be_cmlc": [],
        "be_cmlt": [],
        "be_amlc": [],
        "be_amlt": [],
        "be_cemgil": [],
        "be_goto": [],
        "be_pscore": [],
        "be_infogain": [],
        "be_cmlcm": [],
        "br_fscore": [],
        "br_recall": [],
        "br_precision": [],
        "br_cmlc": [],
        "br_cmlt": [],
        "br_amlc": [],
        "br_amlt": [],
        "br_cemgil": [],
        "br_goto": [],
        "br_pscore": [],
        "br_infogain": [],
        "br_cmlcm": [],
        "gt_fscore": [],
        "gt_recall": [],
        "gt_precision": [],
        "gt_cmlc": [],
        "gt_cmlt": [],
        "gt_amlc": [],
        "gt_amlt": [],
        "gt_cemgil": [],
        "gt_goto": [],
        "gt_pscore": [],
        "gt_infogain": [],
        "gt_cmlcm": [],
        "hw_fscore": [],
        "hw_recall": [],
        "hw_precision": [],
        "hw_cmlc": [],
        "hw_cmlt": [],
        "hw_amlc": [],
        "hw_amlt": [],
        "hw_cemgil": [],
        "hw_goto": [],
        "hw_pscore": [],
        "hw_infogain": [],
        "hw_cmlcm": [],
        "sm_fscore": [],
        "sm_recall": [],
        "sm_precision": [],
        "sm_cmlc": [],
        "sm_cmlt": [],
        "sm_amlc": [],
        "sm_amlt": [],
        "sm_cemgil": [],
        "sm_goto": [],
        "sm_pscore": [],
        "sm_infogain": [],
        "sm_cmlcm": [],
    }

    for i in range(len(predictions)):
        dataset_name = data.at[indices[i], 'dataset']
        predikce_corr = np.array([x for x in predictions[i] if x >= 0])
        anotace_corr = np.array([x for x in anotace[i][0] if x >= 0])
        _, tempp, tempr = me.onset.f_measure(anotace_corr, predikce_corr, window=0.07)
        tempf = me.beat.f_measure(anotace_corr, predikce_corr) # This was just to ensure that beat f-score is no different than onset f-score. (it's not)
        cmlc, cmlt, amlc, amlt = me.beat.continuity(anotace_corr, predikce_corr)
        cemgil = me.beat.cemgil(anotace_corr, predikce_corr)
        goto = me.beat.goto(anotace_corr, predikce_corr)
        pscore = me.beat.p_score(anotace_corr, predikce_corr)
        infogain = me.beat.information_gain(anotace_corr, predikce_corr)
        cmlcm, cmltm = madmom.evaluation.beats.cml(predikce_corr,anotace_corr) #Just to check that both algorithms have same result
        if dataset_name == 'BR':
            score["br_fscore"].append(tempf)
            score["br_recall"].append(tempr)
            score["br_precision"].append(tempp)
            score["br_cmlc"].append(cmlc)
            score["br_cmlt"].append(cmlt)
            score["br_amlc"].append(amlc)
            score["br_amlt"].append(amlt)
            score["br_cemgil"].append(cemgil)
            score["br_goto"].append(goto)
            score["br_pscore"].append(pscore)
            score["br_infogain"].append(infogain)
            score["br_cmlcm"].append(cmlcm)
        elif dataset_name == 'BE':
            tempf, tempp, tempr = me.onset.f_measure(anotace_corr, predikce_corr, window=0.07)
            score["be_fscore"].append(tempf)
            score["be_recall"].append(tempr)
            score["be_precision"].append(tempp)
            score["be_cmlc"].append(cmlc)
            score["be_cmlt"].append(cmlt)
            score["be_amlc"].append(amlc)
            score["be_amlt"].append(amlt)
            score["be_cemgil"].append(cemgil)
            score["be_goto"].append(goto)
            score["be_pscore"].append(pscore)
            score["be_infogain"].append(infogain)
            score["be_cmlcm"].append(cmlcm)
        elif dataset_name == 'GT':
            tempf, tempp, tempr = me.onset.f_measure(anotace_corr, predikce_corr, window=0.07)
            score["gt_fscore"].append(tempf)
            score["gt_recall"].append(tempr)
            score["gt_precision"].append(tempp)
            score["gt_cmlc"].append(cmlc)
            score["gt_cmlt"].append(cmlt)
            score["gt_amlc"].append(amlc)
            score["gt_amlt"].append(amlt)
            score["gt_cemgil"].append(cemgil)
            score["gt_goto"].append(goto)
            score["gt_pscore"].append(pscore)
            score["gt_infogain"].append(infogain)
            score["gt_cmlcm"].append(cmlcm)
        elif dataset_name == 'HW':
            tempf, tempp, tempr = me.onset.f_measure(anotace_corr, predikce_corr, window=0.07)
            score["hw_fscore"].append(tempf)
            score["hw_recall"].append(tempr)
            score["hw_precision"].append(tempp)
            score["hw_cmlc"].append(cmlc)
            score["hw_cmlt"].append(cmlt)
            score["hw_amlc"].append(amlc)
            score["hw_amlt"].append(amlt)
            score["hw_cemgil"].append(cemgil)
            score["hw_goto"].append(goto)
            score["hw_pscore"].append(pscore)
            score["hw_infogain"].append(infogain)
            score["hw_cmlcm"].append(cmlcm)
        elif dataset_name == 'SM':
            score["sm_fscore"].append(tempf)
            score["sm_recall"].append(tempr)
            score["sm_precision"].append(tempp)
            score["sm_cmlc"].append(cmlc)
            score["sm_cmlt"].append(cmlt)
            score["sm_amlc"].append(amlc)
            score["sm_amlt"].append(amlt)
            score["sm_cemgil"].append(cemgil)
            score["sm_goto"].append(goto)
            score["sm_pscore"].append(pscore)
            score["sm_infogain"].append(infogain)
            score["sm_cmlcm"].append(cmlcm)

    print('BRD fscore = {:0.3f}'.format(np.mean(score["br_fscore"])))
    print('BRD recall = {:0.3f}'.format(np.mean(score["br_recall"])))
    print('BRD precision = {:0.3f}'.format(np.mean(score["br_precision"])))
    print('BRD cmlc = {:0.3f}'.format(np.mean(score["br_cmlc"])))
    print('BRD cmlt = {:0.3f}'.format(np.mean(score["br_cmlt"])))
    print('BRD amlc = {:0.3f}'.format(np.mean(score["br_amlc"])))
    print('BRD amlt = {:0.3f}'.format(np.mean(score["br_amlt"])))
    print('BRD cemgil = {:0.3f}'.format(np.mean(score["br_cemgil"])))
    print('BRD goto = {:0.3f}'.format(np.mean(score["br_goto"])))
    print('BRD pscore = {:0.3f}'.format(np.mean(score["br_pscore"])))
    print('BRD infogain = {:0.3f}'.format(np.mean(score["br_infogain"])))
    print('BRD cmlcm = {:0.3f}'.format(np.mean(score["br_cmlcm"])))

    print('BEATLES fscore = {:0.3f}'.format(np.mean(score["be_fscore"])))
    print('BEATLES recall = {:0.3f}'.format(np.mean(score["be_recall"])))
    print('BEATLES precision = {:0.3f}'.format(np.mean(score["be_precision"])))
    print('BEATLES cmlc = {:0.3f}'.format(np.mean(score["be_cmlc"])))
    print('BEATLES cmlt = {:0.3f}'.format(np.mean(score["be_cmlt"])))
    print('BEATLES amlc = {:0.3f}'.format(np.mean(score["be_amlc"])))
    print('BEATLES amlt = {:0.3f}'.format(np.mean(score["be_amlt"])))
    print('BEATLES cemgil = {:0.3f}'.format(np.mean(score["be_cemgil"])))
    print('BEATLES goto = {:0.3f}'.format(np.mean(score["be_goto"])))
    print('BEATLES pscore = {:0.3f}'.format(np.mean(score["be_pscore"])))
    print('BEATLES infogain = {:0.3f}'.format(np.mean(score["be_infogain"])))
    print('BEATLES cmlcm = {:0.3f}'.format(np.mean(score["be_cmlcm"])))

    print('GTZAN fscore = {:0.3f}'.format(np.mean(score["gt_fscore"])))
    print('GTZAN recall = {:0.3f}'.format(np.mean(score["gt_recall"])))
    print('GTZAN precision = {:0.3f}'.format(np.mean(score["gt_precision"])))
    print('GTZAN cmlc = {:0.3f}'.format(np.mean(score["gt_cmlc"])))
    print('GTZAN cmlt = {:0.3f}'.format(np.mean(score["gt_cmlt"])))
    print('GTZAN amlc = {:0.3f}'.format(np.mean(score["gt_amlc"])))
    print('GTZAN amlt = {:0.3f}'.format(np.mean(score["gt_amlt"])))
    print('GTZAN cemgil = {:0.3f}'.format(np.mean(score["gt_cemgil"])))
    print('GTZAN goto = {:0.3f}'.format(np.mean(score["gt_goto"])))
    print('GTZAN pscore = {:0.3f}'.format(np.mean(score["gt_pscore"])))
    print('GTZAN infogain = {:0.3f}'.format(np.mean(score["gt_infogain"])))
    print('GTZAN cmlcm = {:0.3f}'.format(np.mean(score["gt_cmlcm"])))

    print('HAINSWORTH fscore = {:0.3f}'.format(np.mean(score["hw_fscore"])))
    print('HAINSWORTH recall = {:0.3f}'.format(np.mean(score["hw_recall"])))
    print('HAINSWORTH precision = {:0.3f}'.format(np.mean(score["hw_precision"])))
    print('HAINSWORTH cmlc = {:0.3f}'.format(np.mean(score["hw_cmlc"])))
    print('HAINSWORTH cmlt = {:0.3f}'.format(np.mean(score["hw_cmlt"])))
    print('HAINSWORTH amlc = {:0.3f}'.format(np.mean(score["hw_amlc"])))
    print('HAINSWORTH amlt = {:0.3f}'.format(np.mean(score["hw_amlt"])))
    print('HAINSWORTH cemgil = {:0.3f}'.format(np.mean(score["hw_cemgil"])))
    print('HAINSWORTH goto = {:0.3f}'.format(np.mean(score["hw_goto"])))
    print('HAINSWORTH pscore = {:0.3f}'.format(np.mean(score["hw_pscore"])))
    print('HAINSWORTH infogain = {:0.3f}'.format(np.mean(score["hw_infogain"])))
    print('HAINSWORTH cmlcm = {:0.3f}'.format(np.mean(score["hw_cmlcm"])))

    print('SMC fscore = {:0.3f}'.format(np.mean(score["sm_fscore"])))
    print('SMC recall = {:0.3f}'.format(np.mean(score["sm_recall"])))
    print('SMC precision = {:0.3f}'.format(np.mean(score["sm_precision"])))
    print('SMC cmlc = {:0.3f}'.format(np.mean(score["sm_cmlc"])))
    print('SMC cmlt = {:0.3f}'.format(np.mean(score["sm_cmlt"])))
    print('SMC amlc = {:0.3f}'.format(np.mean(score["sm_amlc"])))
    print('SMC amlt = {:0.3f}'.format(np.mean(score["sm_amlt"])))
    print('SMC cemgil = {:0.3f}'.format(np.mean(score["sm_cemgil"])))
    print('SMC goto = {:0.3f}'.format(np.mean(score["sm_goto"])))
    print('SMC pscore = {:0.3f}'.format(np.mean(score["sm_pscore"])))
    print('SMC infogain = {:0.3f}'.format(np.mean(score["sm_infogain"])))
    print('SMC cmlcm = {:0.3f}'.format(np.mean(score["sm_cmlcm"])))
    return score


def get_madmom_score_sorted(scores, data, indices):
    # IN: precomputed madmom metrics for whole dataset, data.csv, indices of given files
    # OUT: Score dictionary, also printed
    score = {
        "br_fscore": [],
        "br_recall": [],
        "br_precision": [],
        "be_fscore": [],
        "be_recall": [],
        "be_precision": [],
        "gt_fscore": [],
        "gt_recall": [],
        "gt_precision": [],
        "hw_fscore": [],
        "hw_recall": [],
        "hw_precision": [],
        "sm_fscore": [],
        "sm_recall": [],
        "sm_precision": [],
    }

    for i in range(len(indices)):
        dataset_name = data.at[indices[i], "dataset"]
        fscore = scores.at[indices[i], "fscore"]
        recall = scores.at[indices[i], "recall"]
        precision = scores.at[indices[i], "precision"]


        if dataset_name == 'BR':
            score["br_fscore"].append(fscore)
            score["br_recall"].append(recall)
            score["br_precision"].append(precision)
        if dataset_name == 'BE':
            score["be_fscore"].append(fscore)
            score["be_recall"].append(recall)
            score["be_precision"].append(precision)
        elif dataset_name == 'GT':
            score["gt_fscore"].append(fscore)
            score["gt_recall"].append(recall)
            score["gt_precision"].append(precision)
        elif dataset_name == 'HW':
            score["hw_fscore"].append(fscore)
            score["hw_recall"].append(recall)
            score["hw_precision"].append(precision)
        elif dataset_name == 'SM':
            score["sm_fscore"].append(fscore)
            score["sm_recall"].append(recall)
            score["sm_precision"].append(precision)

    print('BRD fscore = {:0.3f}'.format(np.mean(score["br_fscore"])))
    print('BRD recall = {:0.3f}'.format(np.mean(score["br_recall"])))
    print('BRD precision = {:0.3f}'.format(np.mean(score["br_precision"])))

    print('BEATLES fscore = {:0.3f}'.format(np.mean(score["be_fscore"])))
    print('BEATLES recall = {:0.3f}'.format(np.mean(score["be_recall"])))
    print('BEATLES precision = {:0.3f}'.format(np.mean(score["be_precision"])))

    print('GTZAN fscore = {:0.3f}'.format(np.mean(score["gt_fscore"])))
    print('GTZAN recall = {:0.3f}'.format(np.mean(score["gt_recall"])))
    print('GTZAN precision = {:0.3f}'.format(np.mean(score["gt_precision"])))

    print('HAINSWORTH fscore = {:0.3f}'.format(np.mean(score["hw_fscore"])))
    print('HAINSWORTH recall = {:0.3f}'.format(np.mean(score["hw_recall"])))
    print('HAINSWORTH precision = {:0.3f}'.format(np.mean(score["hw_precision"])))

    print('SMC fscore = {:0.3f}'.format(np.mean(score["sm_fscore"])))
    print('SMC recall = {:0.3f}'.format(np.mean(score["sm_recall"])))
    print('SMC precision = {:0.3f}'.format(np.mean(score["sm_precision"])))
    return score