import numpy as np
import os
import trojzvuky as trz
import mir_eval
import sys
import json
from librosa import note_to_hz
import csv
sys.path.insert(0, ".\\piano_transcription-master")
from ground_truth import anotator

#Test metody spektrálních špiček se známým počtem tónů na prvním datasetu
est_freq = []
ref_freq = []
est_freq_cast =[]
ref_freq_cast = []
initialdir = r".\dataset\dataset1"
soubory = os.listdir(initialdir)

for i in soubory:
    nazev = initialdir + "\\" + i
    temp = anotator(i)
    print(i)
    akord = trz.Chord(nazev)
    akord.spektrum, akord.rate, akord.delka, akord.audio, akord.pocet_tonu, akord.pravda = akord.audioloader(pocet_tonu=len(temp))
    ref_freq.append(akord.pravda)
    osa = akord.freq_osa(akord.delka, akord.rate)
    f,*_ = akord.f0_finder(akord.spektrum, akord.rate, akord.delka, osa)
    est_freq.append(f)
    if np.min(akord.pravda) > 138.59:
        ref_freq_cast.append(akord.pravda)
        est_freq_cast.append(f)

timestamps = np.arange(len(ref_freq))
timestamps_cast = np.arange(len(ref_freq_cast))

metrics_moje = mir_eval.multipitch.metrics(ref_time=timestamps, ref_freqs=ref_freq, est_time=timestamps, est_freqs=est_freq)  
metrics_moje_cast = mir_eval.multipitch.metrics(ref_time=timestamps_cast, ref_freqs=ref_freq_cast, est_time=timestamps_cast, est_freqs=est_freq_cast)  

hodnoty_moje_znamy = {
    "precission" : metrics_moje[0],
    "recall" : metrics_moje[1],
    "accuracy" : metrics_moje[2],
    "e_sub" : metrics_moje[3],
    "e_miss" : metrics_moje[4],
    "e_fa": metrics_moje[5],
    "e_tot": metrics_moje[6]
}

hodnoty_moje_znamy_cast = {
    "precission" : metrics_moje_cast[0],
    "recall" : metrics_moje_cast[1],
    "accuracy" : metrics_moje_cast[2],
    "e_sub" : metrics_moje_cast[3],
    "e_miss" : metrics_moje_cast[4],
    "e_fa": metrics_moje_cast[5],
    "e_tot": metrics_moje_cast[6]
}

with open("metriky_moje_metoda_znamy.json", "w") as f:
    json.dump(hodnoty_moje_znamy,f)

with open("metriky_moje_metoda_znamy_D3+.json", "w") as f:
     json.dump(hodnoty_moje_znamy_cast,f)

#Test metody spektrálních špiček s neznámým počtem tónů a PTR na prvním datasetu
est_freq_moje = []
est_freq_bytedance = []
ref_freq = []

est_freq_moje_cast=[]
est_freq_btd_cast=[]
ref_freq_cast= []

for i in soubory:
    nazev = initialdir + "\\" + i
    temp = anotator(i)
    print(i)
    akord = trz.Chord(nazev)
    akord.spektrum, akord.rate, akord.delka, akord.audio, akord.pocet_tonu, akord.pravda, transcript = akord.audioloader()
    ref_freq.append(akord.pravda)
    est_freq_bytedance.append(transcript)
    osa = akord.freq_osa(akord.delka, akord.rate)
    f,*_ = akord.f0_finder(akord.spektrum, akord.rate, akord.delka, osa)
    est_freq_moje.append(f)
    if np.min(akord.pravda) > 138.59:
        ref_freq_cast.append(akord.pravda)
        est_freq_moje_cast.append(f)
        est_freq_btd_cast.append(transcript)


timestamps = np.arange(len(ref_freq))
timestamps_cast = np.arange(len(ref_freq_cast))

metrics_moje_btd = mir_eval.multipitch.metrics(ref_time=timestamps, ref_freqs=ref_freq, est_time=timestamps, est_freqs=est_freq_moje)  
metrics_bytedance = mir_eval.multipitch.metrics(ref_time=timestamps, ref_freqs=ref_freq, est_time=timestamps, est_freqs=est_freq_bytedance)
metrics_moje_btd_cast = mir_eval.multipitch.metrics(ref_time=timestamps_cast, ref_freqs=ref_freq_cast, est_time=timestamps_cast, est_freqs=est_freq_moje_cast)  
metrics_bytedance_cast = mir_eval.multipitch.metrics(ref_time=timestamps_cast, ref_freqs=ref_freq_cast, est_time=timestamps_cast, est_freqs=est_freq_btd_cast)

hodnoty_moje_neznamy = {
    "precission" : metrics_moje_btd[0],
    "recall" : metrics_moje_btd[1],
    "accuracy" : metrics_moje_btd[2],
    "e_sub" : metrics_moje_btd[3],
    "e_miss" : metrics_moje_btd[4],
    "e_fa": metrics_moje_btd[5],
    "e_tot": metrics_moje_btd[6]
}

hodnoty_bytedance = {
    "precission" : metrics_bytedance[0],
    "recall" : metrics_bytedance[1],
    "accuracy" : metrics_bytedance[2],
    "e_sub" : metrics_bytedance[3],
    "e_miss" : metrics_bytedance[4],
    "e_fa" : metrics_bytedance[5],
    "e_tot": metrics_bytedance[6]
}

hodnoty_moje_neznamy_cast = {
    "precission" : metrics_moje_btd_cast[0],
    "recall" : metrics_moje_btd_cast[1],
    "accuracy" : metrics_moje_btd_cast[2],
    "e_sub" : metrics_moje_btd_cast[3],
    "e_miss" : metrics_moje_btd_cast[4],
    "e_fa": metrics_moje_btd_cast[5],
    "e_tot": metrics_moje_btd_cast[6]
}

hodnoty_btd_cast = {
    "precission" : metrics_bytedance_cast[0],
    "recall" : metrics_bytedance_cast[1],
    "accuracy" : metrics_bytedance_cast[2],
    "e_sub" : metrics_bytedance_cast[3],
    "e_miss" : metrics_bytedance_cast[4],
    "e_fa": metrics_bytedance_cast[5],
    "e_tot": metrics_bytedance_cast[6]
}

with open("metriky_moje_metoda_neznamy.json", "w") as f:
    json.dump(hodnoty_moje_neznamy,f)

with open("metriky_PTR.json", "w") as f:
     json.dump(hodnoty_bytedance,f)

with open("metriky_moje_metoda_neznamy_D3+.json", "w") as f:
    json.dump(hodnoty_moje_neznamy_cast,f)

with open("metriky_PTR_D3+.json", "w") as f:
    json.dump(hodnoty_btd_cast,f)



#Test metody spektrálních špiček se známým počtem na druhém datasetu (pouze trojzvuky)
est_freq = []
est_freq_cast = []
ref_freq_cast = []
ref_freq = []
initialdir = r".\dataset\piano_triads"
soubory = os.listdir(initialdir)
skipped_row = ""

def anotator2(nazev):
    '''Vrací základní pravdu odvozenou z přiloženého .csv souboru'''
    nazev_kratky = nazev.replace(".wav", "")
    pravdiva_matice = np.array([])
    with open('triads.csv') as csv_file:
        csv_reader = csv.reader(csv_file)
        header_row = next(csv_reader)  
        for i, row in enumerate(csv_reader):
            if row[0] != nazev_kratky:  
                continue
           
            for K in row[1:4]:
                temp = K
                cislo = int(K[-1]) - 1
                temp = temp.replace(K[-1], str(cislo))
                if temp[1] == "s":
                    temp=temp.replace("s","#")
                pravdiva_matice=np.append(pravdiva_matice, temp.replace("_",""))  
            pravdiva_matice = note_to_hz(pravdiva_matice)
            for m in pravdiva_matice:
                if m > 5000:
                    return None
            return pravdiva_matice

for i in soubory:
    nazev = initialdir + "\\" + i
    spravna = anotator2(i)
    if type(spravna) != np.ndarray:
        continue
    print(i)
    akord = trz.Chord(nazev)
    akord.spektrum, akord.rate, akord.delka, akord.audio, akord.pocet_tonu, akord.pravda = akord.audioloader(pocet_tonu=3)
    ref_freq.append(spravna)
    osa = akord.freq_osa(akord.delka, akord.rate)
    f,*_ = akord.f0_finder(akord.spektrum, akord.rate, akord.delka, osa)
    est_freq.append(f)
    if np.min(spravna) > 138.59:   
        ref_freq_cast.append(spravna)
        est_freq_cast.append(f)

timestamps = np.arange(len(ref_freq))
timestamps_cast = np.arange(len(ref_freq_cast))
metrics_moje = mir_eval.multipitch.metrics(ref_time=timestamps, ref_freqs=ref_freq, est_time=timestamps, est_freqs=est_freq)
metrics_moje_cast = mir_eval.multipitch.metrics(ref_time=timestamps_cast, ref_freqs=ref_freq_cast, est_time=timestamps_cast, est_freqs=est_freq_cast)  


hodnoty_moje_dataset2 = {
    "precission" : metrics_moje[0],
    "recall" : metrics_moje[1],
    "accuracy" : metrics_moje[2],
    "e_sub" : metrics_moje[3],
    "e_miss" : metrics_moje[4],
    "e_fa": metrics_moje[5],
    "e_tot": metrics_moje[6]
}


hodnoty_moje_cast_dataset2 = {
    "precission" : metrics_moje_cast[0],
    "recall" : metrics_moje_cast[1],
    "accuracy" : metrics_moje_cast[2],
    "e_sub" : metrics_moje_cast[3],
    "e_miss" : metrics_moje_cast[4],
    "e_fa": metrics_moje_cast[5],
    "e_tot": metrics_moje_cast[6]
}


with open("metriky_moje_metoda_znamy_dataset2.json", "w") as f:
    json.dump(hodnoty_moje_dataset2,f)

with open("metriky_moje_metoda_znamy_dataset2_D3+.json", "w") as f:
     json.dump(hodnoty_moje_cast_dataset2,f)



#Test metody spektrálních špiček s neznámým počtem tónů na druhém datasetu
est_freq_moje = []
est_freq_moje_cast = []
est_freq_bytedance = []
est_freq_bytedance_cast = []
ref_freq_cast = []
ref_freq = []
skipped_row = ""

for i in soubory:
    nazev = initialdir + "\\" + i
    spravna = anotator2(i)
    if type(spravna) != np.ndarray:
        continue
    print(i)
    akord = trz.Chord(nazev)
    akord.spektrum, akord.rate, akord.delka, akord.audio, akord.pocet_tonu, akord.pravda, transcript= akord.audioloader()
    ref_freq.append(spravna)
    est_freq_bytedance.append(transcript)
    osa = akord.freq_osa(akord.delka, akord.rate)
    f,*_ = akord.f0_finder(akord.spektrum, akord.rate, akord.delka, osa)
    est_freq_moje.append(f)
    if np.min(spravna) > 138.59:   
        ref_freq_cast.append(spravna)
        est_freq_moje_cast.append(f)
        est_freq_bytedance_cast.append(transcript)


timestamps = np.arange(len(ref_freq))
timestamps_cast = np.arange(len(ref_freq_cast))
metrics_moje_btd = mir_eval.multipitch.metrics(ref_time=timestamps, ref_freqs=ref_freq, est_time=timestamps, est_freqs=est_freq_moje)
metrics_btd = mir_eval.multipitch.metrics(ref_time=timestamps, ref_freqs=ref_freq, est_time=timestamps, est_freqs=est_freq_bytedance)
metrics_moje_cast = mir_eval.multipitch.metrics(ref_time=timestamps_cast, ref_freqs=ref_freq_cast, est_time=timestamps_cast, est_freqs=est_freq_moje_cast)  
metrics_bytedance_cast = mir_eval.multipitch.metrics(ref_time=timestamps_cast, ref_freqs=ref_freq_cast, est_time=timestamps_cast, est_freqs=est_freq_bytedance_cast)  

hodnoty_moje_neznamy_dataset2 = {
    "precission" : metrics_moje_btd[0],
    "recall" : metrics_moje_btd[1],
    "accuracy" : metrics_moje_btd[2],
    "e_sub" : metrics_moje_btd[3],
    "e_miss" : metrics_moje_btd[4],
    "e_fa": metrics_moje_btd[5],
    "e_tot": metrics_moje_btd[6]
}

hodnoty_ptr_dataset2 = {
    "precission" : metrics_btd[0],
    "recall" : metrics_btd[1],
    "accuracy" : metrics_btd[2],
    "e_sub" : metrics_btd[3],
    "e_miss" : metrics_btd[4],
    "e_fa": metrics_btd[5],
    "e_tot": metrics_btd[6]
}

hodnoty_moje_neznamy_cast_dataset2 = {
    "precission" : metrics_moje_cast[0],
    "recall" : metrics_moje_cast[1],
    "accuracy" : metrics_moje_cast[2],
    "e_sub" : metrics_moje_cast[3],
    "e_miss" : metrics_moje_cast[4],
    "e_fa": metrics_moje_cast[5],
    "e_tot": metrics_moje_cast[6]
}

hodnoty_ptr_cast_dataset2 = {
    "precission" : metrics_bytedance_cast[0],
    "recall" : metrics_bytedance_cast[1],
    "accuracy" : metrics_bytedance_cast[2],
    "e_sub" : metrics_bytedance_cast[3],
    "e_miss" : metrics_bytedance_cast[4],
    "e_fa": metrics_bytedance_cast[5],
    "e_tot": metrics_bytedance_cast[6]
}

with open("metriky_moje_metoda_neznamy_dataset2.json", "w") as f:
    json.dump(hodnoty_moje_neznamy_dataset2,f)

with open("metriky_PTR_dataset2+.json", "w") as f:
     json.dump(hodnoty_ptr_dataset2,f)

with open("metriky_moje_metoda_neznamy_dataset2_D3+.json", "w") as f:
     json.dump(hodnoty_moje_neznamy_cast_dataset2,f)

with open("metriky_PTR_dataset2_D3+.json", "w") as f:
     json.dump(hodnoty_ptr_cast_dataset2,f)