import pandas as pd
import numpy as np

from sklearn.metrics import average_precision_score

#Expects ann-fsinst-distmat_nlbopti.txt and ann-fsinst-distmat_exactpitch40.txt
#12042 * 360 distmats

parser = argparse.ArgumentParser(description='Compute Average Precision Scores')
parser.add_argument('-distmat', metavar='mtcroot', type=str, help='path to distance matrix as computed by nlbalign')
parser.add_argument('-mtcroot', metavar='mtcroot', type=str, help='path to MTC to find metadata', default='/Users/pvk/data/MTC/')
parser.add_argument('-mtcjsonroot', metavar='mtcjsonroot', type=str, help='mtcroot for json files as generated by krn2json', default='/Users/pvk/data/MTCjson/')
args = parser.parse_args()

mtcannroot = mtcroot + 'MTC-ANN-2.0.1/'
mtcfsroot = mtcroot + 'MTC-FS-INST-2.0/'

mtcannjsondir = mtcjsonroot + 'MTC-ANN-2.0.1/json/'
mtcfsjsondir = mtcjsonroot + 'MTC-FS-INST-2.0/json/'

mtcann_metadata = pd.read_csv(mtcannroot+'metadata/MTC-ANN-tune-family-labels.csv', na_filter=False, index_col=0, header=None, encoding='utf8', names=['ID','TuneFamily'])
mtcfsinst_metadata = pd.read_csv(mtcfsroot+'metadata/MTC-FS-INST-2.0.csv', na_filter=False, index_col=0, header=None, encoding='utf8', names=["filename","songid","source_id","serial_number","page","singer_id_s","date_of_recording","place_of_recording","latitude","longitude","textfamily_id","title","firstline","tunefamily_id","tunefamily","type","voice_stanza_number","voice_stanza","image_filename_s","audio_filename","variation","confidence","comment","MTC_title","author"])

#group by tune family
mtcann_metadata_bytf = mtcann_metadata.groupby('TuneFamily')

#read the distmats (generated by nlbdistmat from libmusical)
distmat_nlbopti = pd.read_csv('ann-fsinst-distmat_nlbopti.txt', index_col=0, sep='\t', header=0)
distmat_exactpitch = pd.read_csv('ann-fsinst-distmat_exactpitch40.txt', index_col=0, sep='\t', header=0)

#get ids of all tune family members for a given song id
#nlbid should be in MTC-ANN
def getAllTuneFamilyMembers(nlbid):
    return list(mtcann_metadata_bytf.groups[mtcann_metadata.loc[nlbid]['TuneFamily']])

#tune family labels of ann are in mtcann_metadata

#for each query, construct score and true label
#score is normalized rank [0,1] in sorted result list

#nlbid must be an nlbid in MTC-ANN
def getLabelAndScore(nlbid, distmat):
    tfmembers = getAllTuneFamilyMembers(nlbid) #these are the positive
    y_true = 1*distmat.index.isin(tfmembers) #vector with 1's at the position of the relevant documents
    y_score = distmat[nlbid].values
    y = list(zip(y_true, y_score))
    y = sorted(y, key=lambda x: x[1])
    #remove query
    y = y[1:]
    y_true, y_score = zip(*y)
    y_score = np.linspace(1.0, 0.0, num=len(y_score), endpoint=True)
    return y_true, y_score

#store aveage precision scores per query (=mtcann song)
ap_scores = pd.DataFrame(index=mtcann_metadata.index)

def addScores(ap_scores, distmat, name="DistMeasure"):
    for nlbid in mtcann_metadata.index:
        y_true, y_score = getLabelAndScore(nlbid,distmat)
        ap = average_precision_score(y_true, y_score)
        ap_scores.loc[nlbid,name] = ap

addScores(ap_scores, distmat_nlbopti, "opti_full")
addScores(ap_scores, distmat_nlbopti[-360:], "opti_ann") #last 360 rows are mtc-ann
addScores(ap_scores, distmat_exactpitch, "exactpitch_full")
addScores(ap_scores, distmat_exactpitch[-360:], "exactpitch_ann") ##last 360 rows are mtc-ann

#write to disc
with open('ap_scores.csv','w') as f:
    ap_scores.to_csv(f)
