from piano_transcription_inference import PianoTranscription, sample_rate, load_audio
import librosa 
import os
import numpy as np

def bytedance(nazev):
    # Load audio
    (audio, _) = load_audio(nazev, sr=sample_rate, mono=True)

    # Transcriptor
    transcriptor = PianoTranscription(device='cpu')    # 'cuda' | 'cpu'

    transcribed_dict = transcriptor.transcribe(audio, midi_path=False)

    bytedance_vystup = []
    pocet = 0
    #print(transcribed_dict["est_note_events"])
    for i in transcribed_dict["est_note_events"]:
        delka_noty = i["offset_time"] - i["onset_time"]
        bytedance_vystup.append(i["midi_note"])
        if delka_noty > 0.3: #
            pocet +=1
    return pocet, librosa.midi_to_hz(bytedance_vystup)

def anotator(chord_name):
    '''Vrací matici frekvencí podle názvu souboru z datasetu vytvořeného pro tuto práci'''
    chord_name = chord_name.rstrip(".wav")
    septakordy = {
        3 : [4, 3, 3],
        4 : [3, 4, 3],
        5 : [4, 3, 4],
        6 : [3, 4, 4],
        7 : [4, 4, 3],
        8 : [3, 3, 4],
        9 : [3, 3, 3]
    }
    chord_notes = [0,0,0]
    # rozdělí název akordu podle podtržítek
    if len(chord_name) == 2:
        my_array = np.ones(1) * librosa.note_to_hz(chord_name)
        return my_array
    parts = chord_name.split('_')
    inversion = int(parts[-2])
    octave = parts[-1]
    # získáme název tónu a kód obratu
    koren = parts[2]
    if len(koren) == 2:
        _ = list(koren)
        _[-1] = "#"
        koren = ''.join(_)
    pismeno = koren[0]
    tonality = pismeno.isupper()
    if tonality:
        a, b = 4, 3
    else: 
        a, b = 3, 4
    chord_notes[0] = librosa.note_to_midi(koren + octave)    
    if len(parts) == 5:#151_h_G_1_2
        chord_notes[1] = chord_notes[0] + a 
        chord_notes[2] =  chord_notes[1] + b 
        if inversion == 1:
            chord_notes[0] += 12
        elif inversion == 2:
            chord_notes[0] += 12
            chord_notes[1] += 12
    elif len(parts) == 6: 
        typ = int(parts[3])
        chord_notes.append("")
        if typ == 10:
            chord_notes[1] = chord_notes[0] + a 
            chord_notes[2] =  chord_notes[1] + b 
            chord_notes[3] = chord_notes[0] + 12
            if inversion != 0:
                chord_notes[0] += 12
                chord_notes[3] += a
            if inversion == 2:
                chord_notes[0] == chord_notes[0] + a + 12
                chord_notes[1] == chord_notes[1] + b + 12
        else: 
            chord_notes[1] = chord_notes[0] + septakordy[typ][0]
            chord_notes[2] = chord_notes[1] + septakordy[typ][1]
            chord_notes[3] = chord_notes[2] + septakordy[typ][2]
            if inversion == 1:
                chord_notes[0] += 12
            elif inversion == 2:
                chord_notes[0] += 12
                chord_notes[1] += 12
            elif inversion == 3:
                chord_notes[0] += 12
                chord_notes[1] += 12
                chord_notes[2] += 12
    chord_notes.sort()
    chord_letters = librosa.midi_to_note(chord_notes)
    octave_nums = []
    for i in chord_letters:
        octave_nums.append(i[-1])
    if octave not in octave_nums:
        chord_notes2 = []
        for i in chord_notes:
            chord_notes2.append(i-12)
        chord_letters = librosa.midi_to_note(chord_notes2)  
        return np.round(librosa.midi_to_hz(chord_notes2), decimals=2)

    else:
        return np.round(librosa.midi_to_hz(chord_notes), decimals=2)

