import librosa
import matplotlib.pyplot as plt
import numpy as np
import librosa.display
from scipy import signal
from scipy.fftpack import fft
import IPython.display as ipd
import sys
import os
from gui import sample_select, pocet_demo
sys.path.insert(0, ".\\piano_transcription-master")
from ground_truth import anotator, bytedance


class Chord:
    def __init__(self, name, pocet_tonu=None):
        self.name = name
        self.spektrum = None
        self.rate = None
        self.delka = None
        self.audio = None
        if pocet_tonu:
            self.pocet_tonu = pocet_tonu
        else:
            self.pocet_tonu = None
        self.pravda = None

    def audioloader(self, pocet_tonu=None):
        '''Načítá zvukový soubor a nastavuje atributy: název, pole vzorků, f_vz, Fourierovu transformaci, délku souboru ve vzorcích,
          základní pravdu, počet tónů'''
        filename = self.name
        data, rate = librosa.load(filename, sr=librosa.get_samplerate(filename))
        fft_out = fft(data)
        N = len(data)
        pravda = anotator(os.path.basename(filename))
        if pocet_tonu:
            pocet = pocet_tonu
            return fft_out, rate, N, data, pocet, pravda
        else:
            _ = bytedance(filename)
            pocet = _[0]
            bytedance_vystup = _[1] 
            return fft_out, rate, N, data, pocet, pravda, bytedance_vystup
    
    def freq_osa(self, N, rate):
        '''Nastavuje frekvenční osu pro vizualizaci FT'''
        xf = []
        for i in range(N):
            xf.append((i*rate)/N)
        return xf
    
    def plotter(self, fce, osa, nazev, proved = False):
        '''Pokud proved==True, dovoluje vykreslit a eventuálně uložit vizualizaci modulového spektra'''
        if proved:
            plt.xscale("log")
            plt.plot(osa, np.abs(fce))
            plt.xlabel("frekvence [Hz]")
            plt.ylabel("Modul [-]")
            plt.xlim(left=27, right=20000)
            plt.title(nazev)
            #plt.savefig('spektrum_multipitch.svg')
            plt.show()
            
         

    def f0_finder(self, fce, sr, N, osa):
        '''Vyhledává kmitočty nejsilnějších složek spektra nahrávky akordu'''
        semitone = 2**(1/12)
        f = []
        p = []  
        fce[int(len(osa)/2):] = 0 # Uvážení pouze kladné části spektra pod vzorkovacím kmitočtem
        fce[:int((27*N/sr))] = 0 # Vynulování složek menších než 27 Hz
        self.spektrum=fce.copy()
        i = 1
        fake_i = 0
        pocet_iteraci = self.pocet_tonu
        self.plotter(fce, osa, "Spektrum zadaného akordu")
        while i <= pocet_iteraci:
            candidate = (np.argmax(np.abs(fce)))
            f_center = (candidate*sr)/N
            a = np.abs(self.spektrum[int(candidate/2)])
            b = np.abs(self.spektrum[int(candidate/0.89)])
            if a - b > 20: # jsi druhá harmonická? Ano 
                fce[int((f_center/(semitone**0.2))*(N/sr)):int((f_center*(semitone**0.2))*(N/sr))] = 0 #Vynulovat
                #print(f"{f_center} byla 2. harmonickou, vynulováno")
                if np.abs(fce[int((f_center/2)*(N/sr))]) == 0: #Máme zapsaný tvůj fundament? Ano
                    #print(f"{f_center/2} je zapsano")
                    pocet_iteraci += 1 
                    i += 1
                    fake_i += 1
                    continue
                else:
                    #print("zapsan jeji fundament")
                    f_center = f_center/2
            if i == 1:
                fce[:int((f_center/2)*N/sr)] = 0
            if f_center >= librosa.note_to_hz("E7"): # Posílení fundamentů v poslední oktávě klavíru
                fce[candidate:int(4600*(N/sr))] *= 8
            f.append(round(f_center, 2))
            p.append(int(round(librosa.hz_to_midi(f[i-1-fake_i]), 0)))
            fce[int((f_center/(semitone))*(N/sr)):int((f_center*(semitone))*(N/sr))] = 0 # Odfiltrování složky s maximální amplitudou
            if i != self.pocet_tonu:
                self.plotter(fce, osa, f"Spektrum zadaného akordu po odebrání {i-fake_i}. tónu.")
            i += 1
        f.sort()
        f = np.array(f)
        p.sort()
        if self.pocet_tonu == 4 and p[-1] - p[0] > 12:
            p[-1] = p[0] + 12
        if self.pocet_tonu != 0:
            tony = librosa.midi_to_note(p, unicode=False)
        else:
            tony = librosa.midi_to_note(69, unicode=False)
        
        tony2=[]
        for i in range(0, len(tony)):
            tony2.append(tony[i][:-1])
        return f, p, tony, tony2 #Přednání hodnot do funkcí určující tóninu a obrat
    
def obrat_checker3(p):  
    '''Pro trojzvuky kontroluje obrat podle rozdílů MIDI p hodnot. Vrací název a kód pro obrat. '''
    if (p[1]-p[0]) == 4:  
        if (p[2]-p[0] == 7):
            return "dur kvintakord", 0
        else:
            return "moll sextakord", 2
    elif (p[1]-p[0]) == 3:  
        if (p[2]-p[0] == 8):
            return "dur sextakord", 2
        else:
            return "moll kvintakord", 0
    elif (p[1]-p[0]) == 5:  
        if (p[2]-p[0] == 9):
            return "dur kvartsextakord", 1
        else:
            return "moll kvartsextakord", 1
    else:# (p[1]-p[0]) != (3 or 4 or 5):
        return "x", None
    
def typ_obrat_matcher(p):
        '''Pro čtyřzvuky a septakordy, vrací název, kód obratu, kód typu (specifikovaný v textu práce) a kytarovou značku
        (0 - septakord, 3 - kvintsextakord, 2 - terckvartakord, 1 - sekundakord)'''
        r = (p[1]-p[0], p[2]-p[1], p[3]-p[2])
        vypis_akordu = {
            (4, 4, 3) : ["zvětšeně velký septakord", 7, 0, "maj/5#"],
            (4, 4, 1) : ["měkce velký kvintsextakord", 6, 3, "mmaj"],
            (4, 3, 4) : ["tvrdě velký septakord", 5, 0, "maj"],
            (4, 3, 3) : ["tvrdě malý septakord", 3, 0, "7"],
            (4, 3, 2) : ["měkce malý kvintsextakord", 4, 3, "m7"],
            (4, 3, 1) : ["zvětšeně velký kvintsextakord", 7, 3, "maj/5#"],
            (4, 2, 3) : ["zmenšeně malý terckvartakord", 8, 2, "m7/5b"],
            (4, 1, 4) : ["tvrdě velký terckvartakord", 5, 2, "maj"],
            (4, 1, 3) : ["měkce velký terckvartakord", 6, 2, "mmaj"],
            (3, 4, 4) : ["měkce velký septakord", 6, 0, "mmaj"],
            (3, 4, 3) : ["měkce malý septakord", 4, 0, "m7"],
            (3, 4, 2) : ["zmenšeně malý kvintsextakord", 8, 3, "m7/5b"],
            (3, 4, 1) : ["tvrdě velký kvintsextakord", 5, 3, "maj"],
            (3, 3, 4) : ["zmenšeně malý septakord", 8, 0, "m7/5b"],
            (3, 3, 3) : ["zmenšeně zmenšený septakord (více možností)", 9, 0, "dim"],
            (3, 3, 2) : ["tvrdě malý kvintsextakord", 3, 3, "7"],
            (3, 2, 4) : ["tvrdě malý terckvartakord", 3, 2, "7"],
            (3, 2, 3) : ["měkce malý terckvartakord", 4, 2, "m7"],
            (3, 1, 4) : ["zvětšeně velký terckvartakord", 7, 2, "maj/5#"],
            (2, 4, 3) : ["tvrdě malý sekundakord", 3, 1, "7"],
            (2, 3, 4) : ["měkce malý sekundakord", 4, 1, "m7"],       
            (2, 3, 3) : ["zmenšeně malý sekundakord", 8, 1, "m7/5b"],
            (1, 4, 4) : ["zvětšeně velký sekundakord", 7, 1, "maj/5#"],
            (1, 4, 3) : ["tvrdě velký sekundakord", 5, 1, "maj"],
            (1, 3, 4) : ["měkce velký sekundakord", 6, 1, "mmaj"],
            #Neseptakordy
            (4, 3, 5) : ["dur kvintakord", 10, 0, "dur čtyřzvuk"],
            (3, 5, 4) : ["dur sextakord ", 10, 2, "dur čtyřzvuk"],
            (5, 4, 3) : ["dur kvartsextakord ", 10, 1, "dur čtyřzvuk"],
            (3, 4, 5) : ["moll kvintakord", 10, 0, "moll čtyřzvuk"],
            (4, 5, 3) : ["moll sextakord ", 10, 2, "moll čtyřzvuk"],
            (5, 3, 4) : ["moll kvartsextakord ", 10, 1, "moll čtyřzvuk"]
        }
        if r in vypis_akordu:
            return vypis_akordu[r]
        else:
            return ['0', 0, 0, '0']     
           
def f0_finder_mic(fce, sr, N, osa, pocet_tonu):
    '''Vyhledává kmitočty nejsilnějších složek spektra při použití mikrofonu'''
    semitone = 2**(1/12)
    f = []
    p = []
    
    fce[int(len(osa)/2):] = 0 # Uvážení pouze kladné části spektra pod vzorkovacím kmitočtem
    fce[:int((27*N/sr))] = 0 # Vynulování složek menších než 27 Hz
    spic=max(fce)
    fce=fce/spic
    spektrum=fce.copy()
    # plt.plot(fce)
    # plt.show()
    i = 1
    fake_i = 0
    pocet_iteraci = pocet_tonu
    # plt.plot(fce)
    # plt.xscale("log")
    # plt.show()
    while i <= pocet_iteraci:
        candidate = (np.argmax(np.abs(fce)))
        f_center = (candidate*sr)/N
        
        a = np.abs(spektrum[int(candidate/2)])
        b = np.abs(spektrum[int(candidate/0.89)])
        if a - b > 0.01 and a>0.1: # jsi druhá harmonická? Ano
            fce[int((f_center/(semitone**0.2))*(N/sr)):int((f_center*(semitone**0.2))*(N/sr))] = 0 #Vynulovat
            #print(f"{f_center} byla 2. harmonickou, vynulováno")
            if np.abs(fce[int((f_center/2)*(N/sr))]) == 0: #Máme zapsaný tvůj fundament? Ano
                #print(f"{f_center/2} je zapsano")
                pocet_iteraci += 1 
                i += 1
                fake_i += 1
                continue
            else:
                #print("zapsan jeji fundament")
                f_center = f_center/2
        if i == 1:
            fce[:int((f_center/2)*N/sr)] = 0
        if f_center > librosa.note_to_hz("E7"): # Posílení fundamentů v poslední oktávě klavíru
            fce[candidate:] *= 4
        f.append(round(f_center, 3))
        p.append(int(round(librosa.hz_to_midi(f[i-1-fake_i]), 0)))
        fce[int((f_center/(semitone))*(N/sr)):int((f_center*(semitone))*(N/sr))] = 0 # Odfiltrování složky s maximální amplitudou
        i += 1
    f.sort()

    p.sort()
    
    
    if pocet_tonu == 4 and p[-1] - p[0] > 12:
        p[-1] = p[0] + 12

    tony = librosa.midi_to_note(p, unicode=False)
    tony2=[]
    for i in range(0, len(tony)):
        tony2.append(tony[i][:-1])
    return  p, tony, tony2, f

def analyza_dataset_demo():
    '''Provede demonstraci Analyzátoru akordů klavíru na uživatelem vybraných položkách datasetu'''
    lets_continue = True
    with open("AAK_vysledky.txt", "w") as f:
        pass
    while lets_continue:
        num = pocet_demo()
        nazev = sample_select()
        kod = os.path.basename(os.path.normpath(nazev)) 
        textak = open('AAK_vysledky.txt', 'a')
        if nazev != '':
            print('Analýza probíhá')
            textak.write(f"Testovaný vzorek: {kod}\n")
            textak.flush()
        else:
            textak.write("Žádný vzorek nevybrán.\n")
            textak.flush()
            break
        akord1 = Chord(nazev, pocet_tonu=num)
        akord1.spektrum, akord1.rate, akord1.delka, akord1.audio, akord1.pocet_tonu, akord1.pravda, *_ = akord1.audioloader(pocet_tonu=akord1.pocet_tonu)
        osa = akord1.freq_osa(akord1.delka, akord1.rate)
        akord1.plotter(akord1.spektrum, osa, "Spektrum zadaného akordu")
        f, p, tony, tony2 = akord1.f0_finder(akord1.spektrum, akord1.rate, akord1.delka, osa)
        if akord1.pocet_tonu == 3: 
            popis = obrat_checker3(p)
            if popis == ("x", None):
                vypis = f"Výstup určování fundamentů je {tony}\n"
            else:
                vypis = f"Zadaný akord je trojzvuk: {tony2[popis[1]]}{popis[0]} s nejnižším tónem v {tony[0][-1]}. oktávě klavíru.\n\n"
        elif akord1.pocet_tonu == 4:
            popis = typ_obrat_matcher(p)
            if popis == ("x", None):
                vypis = f"Výstup určování fundamentů je {tony}\n"
            else:
                vypis = f"Zadaný akord je čtyřzvuk: {tony2[popis[2]]} {popis[0]} s nenjnižším tónem v {tony[0][-1]}. oktávě klavíru.\n\n" #tedy {tony2[popis[2]]}{popis[3]}
        else:
            vypis = f"Váš vstup sestává z {akord1.pocet_tonu} tónu/tónů: {tony}.\n\n"
        textak.write(vypis)
        textak.flush()
        f, p, tony, tony2 = None, None, None, None
    textak.write("Konec!\n")
    textak.close


