import numpy as np
import librosa
import pandas as pd
from dataclasses import dataclass

# datova trida pro uchovani anotaci a praci s nimi
# time - casova znacka
# frame - hodnota casove znacky prepoctena na ramce
# label - oznaceni casovek znacky
@dataclass
class BeatsData:
    time: np.ndarray = 0
    frame: np.ndarray = 0
    label: np.ndarray = 0

    # vyriznuti casti anotace
    def trim(self, start, end):
        self.time = self.time[start:end+1:]
        self.frame = self.frame[start:end+1:]
        self.label = self.label[start:end+1:]

    # hleda oznaceni doby v seznamu oznaceni a vraci indexu pole
    def index4label(self, beat):
        return np.where(self.label == beat)[0][0]

    # hleda oznaceni doby v seznamu oznaceni a vraci hodnotu casove znacky
    def time4label(self, beat):
        return self.time[np.where(self.label == beat)[0][0]]

    # prepocitava hodnotu casoveho ramce na casovou znacku podle nastaveni nahravky
    def frame2time(self, stngs):
        self.time = librosa.frames_to_time(self.frame, sr=stngs.sr, hop_length=stngs.hop_length)

    # prepocitava casovou zncku na hodnotu casoveho ramce podle nastaveni nahravky
    def time2frame(self, stngs):
        self.frame = librosa.time_to_frames(self.time, sr=stngs.sr, hop_length=stngs.hop_length)

    # nacte casove znacky a oznaceni dob
    # prepocte casove znacky na ramce podle nastaveni nahravky
    def load_beats(self, path_beats, stngs):
        csv = pd.read_csv(path_beats)
        self.time = csv['TIME'].to_numpy()
        self.time2frame(stngs)
        self.label = csv['LABEL'].to_numpy()

    # ulozi casove znacky a oznaceni dob
    def save_beats(self, path_beats, stngs):
        self.frame2time(stngs)
        df = pd.DataFrame({'TIME': self.time, 'LABEL': self.label})
        df.to_csv(path_beats, index=False)


# vypocita casove ramce dob pomoci DTW
# chroma_ref - chromatogram mreferencni nahravky
# chroma_new - chromatogram analyzovane nahravky
# beats_frame_ref - anotace v casovych ramcich
def dtw_beats(chroma_ref, chroma_new, beats_frame_ref):
    D, wp = librosa.sequence.dtw(X=chroma_ref, Y=chroma_new, metric='cosine')
    beats_frame_new = np.array([], int)
    # matice wp ma v prvnim sloupci hodnoty referencni nahravky
    # ramec se musi vyhledat
    # priradi se odpovidajici hodnote ve druhem sloupci
    wp = np.flipud(wp)
    for beat in beats_frame_ref:
        beats_frame_new = np.append(beats_frame_new, wp[np.where(wp[:, 0] == beat)[0][0], 1])
    return beats_frame_new
