from defs import wavs
from defs import beats
from defs import paths as pt
import os
import librosa
import librosa.display
import numpy as np
import gc
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')


# funkce na vytvoreni grafu synchronizace nahravek
def plot_sync(path, wav_ref, wav_new, beats_frame_ref, beats_frame_new, stngs):
    fig = plt.figure(figsize=(16, 8))

    plt.subplot(2, 1, 1)
    librosa.display.waveplot(wav_ref, sr=stngs.sr)
    # plt.title('Synchronizace nahrávek', fontsize=28)
    # plt.xlabel('', fontsize=20)
    # plt.xticks(fontsize=18)
    # plt.ylabel('amplituda [-]', fontsize=20)
    # plt.yticks(fontsize=18)
    ax1 = plt.gca()

    plt.subplot(2, 1, 2)
    librosa.display.waveplot(wav_new, sr=stngs.sr)
    # plt.title('Analyzovaná nahrávka', fontsize=28)
    # plt.xlabel('čas [s]', fontsize=20)
    # plt.xticks(fontsize=18)
    # plt.ylabel('amplituda [-]', fontsize=20)
    # plt.yticks(fontsize=18)
    ax2 = plt.gca()

    plt.tight_layout()
    trans_figure = fig.transFigure.inverted()
    lines = []

    beats_time_ref = librosa.frames_to_time(beats_frame_ref, sr=stngs.sr, hop_length=stngs.hop_length)
    beats_time_new = librosa.frames_to_time(beats_frame_new, sr=stngs.sr, hop_length=stngs.hop_length)
    beats_index_list = np.vstack((beats_time_ref, beats_time_new)).T
    # zmenou indexace muzeme zobrazit kazdou dobu nebo pouze prvni
    # for tp1, tp2 in beats_index_list[::]:
    # zobrazujeme kazdou sestou dobu (nahravky maji 6/8 metrum)
    for tp1, tp2 in beats_index_list[::6]:
        coord1 = trans_figure.transform(ax1.transData.transform([tp1, 0]))
        coord2 = trans_figure.transform(ax2.transData.transform([tp2, 0]))

        line = matplotlib.lines.Line2D((coord1[0], coord2[0]),
                                       (coord1[1], coord2[1]),
                                       transform=fig.transFigure,
                                       color='r')
        lines.append(line)

    fig.lines = lines
    plt.savefig(path)
    # uzavirani referenci a volani garbage collectoru by melo zlepsit pracovani s pameti
    plt.close('all')
    fig.clf()
    gc.collect()


print('*start*')
if not os.path.exists('data/sync'):
    os.makedirs('data/sync')
path = 'data/'
stngs = wavs.Settings(sr=22050, frame_length=2048, hop_length=2048)
# pokud zobrazujeme pouze cast nahravky musime nastavit hranice
# rng = list([1.1, 11.1])
# ref_beat = beats.BeatsData()
# ref_beat.load_beats(path+'_ref/ref.csv', stngs)
# ref_beat.trim(ref_beat.index4label(rng[0]), ref_beat.index4label(rng[1]))
# ref_wav = wavs.WavData(stngs, path+'_ref/ref.wav')
# ref_wav.load_wav([ref_beat.time[0], ref_beat.time[-1]])

# zobrazujeme celou nahravku
ref_beat = beats.BeatsData()
ref_beat.load_beats(path+'_ref/ref.csv', stngs)
ref_wav = wavs.WavData(stngs, path+'_ref/ref.wav')
ref_wav.load_wav()

print('*reference nactena*')
names = pt.listdir('data/_wavs')
# z duvodu pametove narocnosti dochazi k zastaveni procesu
# done se nastavi na True pokud jsou ulozeny vsechny grafy
done = False
while done == False:
    try:
        for name in names:
            print('*zpracovani ' + name + '*')
            # pokud je nahravka uz zpracovana, preskoci se
            if os.path.exists(path+pt.sync(name)):
                print('*skip*')
            else:
                new_beat = beats.BeatsData()
                new_beat.load_beats(path+pt.beat(name), stngs)
                # pokud zobrazujeme pouze cast nahravky musime nastavit hranice
                # new_beat.trim(ref_beat.index4label(rng[0]), ref_beat.index4label(rng[1]))
                # new_wav = wavs.WavData(stngs, path+pt.wav(name))
                # new_wav.load_wav([new_beat.time[0], new_beat.time[-1]])
                # zobrazujeme celou nahravku
                new_wav = wavs.WavData(stngs, path+pt.wav(name))
                new_wav.load_wav()
                plot_sync(path+pt.sync(name), ref_wav.data, new_wav.data, ref_beat.frame, new_beat.frame, stngs)
                print('*graf ulozen*')
            if name == names[-1]:
                done = True
    except:
        continue
print('*hotovo*')



