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


# datova trida pro extrakci a uchovani parametru
@dataclass
class FeatsData:
    rms: np.ndarray = 0
    tempo: np.ndarray = 0
    spec_cent: np.ndarray = 0
    spec_band: np.ndarray = 0
    spec_roll: np.ndarray = 0
    zero: np.ndarray = 0
    plp: np.ndarray = 0
    pyin: np.ndarray = 0
    mfcc: np.ndarray = 0

    # vola funkce na extrakci parametru
    # y - analyzovana data
    # stngs - nastaveni analyzovanych dat
    # target - zrychleni extrakce omezenim poctu parametru
    def get_feats(self, y, stngs, target=None):
        if target == 'FourNotes':
            self.rms = rms(y, stngs)
            self.pyin = pyin(y, stngs)
        if target == 'tempo_rms':
            self.rms = rms(y, stngs)
            self.tempo = tempo(y, stngs)
        else:
            self.rms = rms(y, stngs)
            self.tempo = tempo(y, stngs)
            self.spec_cent = spec_cent(y, stngs)
            self.spec_band = spec_band(y, stngs)
            self.spec_roll = spec_roll(y, stngs)
            self.zero = zero(y, stngs)
            self.plp = plp(y, stngs)
            self.pyin = pyin(y, stngs)
            self.mfcc = mfcc(y, stngs)

    # ulozi extrahovane parametry do CSV
    # data se pred ulozenim zplostuji a transponuji
    # indexy se neukladaji
    # path - cesta k souboru
    # target - omezeni poctu parametru
    def save_feats(self, path, target=None):
        if target == 'FourNotes':
            pd.DataFrame({'RMS': self.rms.ravel().T,
                          'PYIN': self.pyin[0]
                          }).to_csv(path, index=False)

        if target == 'tempo_rms':
            pd.DataFrame({'RMS': self.rms.ravel().T,
                          'TEMPO': self.tempo.ravel().T
                          }).to_csv(path, index=False)
        else:
            pd.DataFrame({'RMS': self.rms.ravel().T,
                          'TEMPO': self.tempo.ravel().T,
                          'SP_CENT': self.spec_cent.ravel().T,
                          'SP_BAND': self.spec_band.ravel().T,
                          'SP_ROLL': self.spec_roll.ravel().T,
                          'ZERO': self.zero.ravel().T,
                          'PLP': self.plp.T,
                          'PYIN': self.pyin[0],
                          'MFCC_02': self.mfcc[1].ravel().T,
                          'MFCC_03': self.mfcc[2].ravel().T,
                          'MFCC_04': self.mfcc[3].ravel().T,
                          'MFCC_05': self.mfcc[4].ravel().T,
                          'MFCC_06': self.mfcc[5].ravel().T,
                          'MFCC_07': self.mfcc[6].ravel().T,
                          'MFCC_08': self.mfcc[7].ravel().T,
                          'MFCC_09': self.mfcc[8].ravel().T,
                          'MFCC_10': self.mfcc[9].ravel().T,
                          'MFCC_11': self.mfcc[10].ravel().T,
                          'MFCC_12': self.mfcc[11].ravel().T,
                          'MFCC_13': self.mfcc[12].ravel().T
                          }).to_csv(path, index=False)

    # nacte extrahovane parametry ze souboru
    # parametry MFCC se doplnuji nulou, protoze je neukladame, ale chceme zachovat pocet parametru
    # path - cesta k souboru
    # target - omezeni poctu parametru
    def load_feats(self, path_feats, target=None):
        csv = pd.read_csv(path_feats)
        if target == 'FourNotes':
            self.rms = csv['RMS'].to_numpy()
            self.pyin = csv['PYIN'].to_numpy()
        if target == 'tempo_rms':
            self.rms = csv['RMS'].to_numpy()
            self.tempo = csv['TEMPO'].to_numpy()
        else:
            self.rms = csv['RMS'].to_numpy()
            self.tempo = csv['TEMPO'].to_numpy()
            self.spec_cent = csv['SP_CENT'].to_numpy()
            self.spec_band = csv['SP_BAND'].to_numpy()
            self.spec_roll = csv['SP_ROLL'].to_numpy()
            self.zero = csv['ZERO'].to_numpy()
            self.plp = csv['PLP'].to_numpy()
            self.pyin = csv['PYIN'].to_numpy()
            self.mfcc = np.zeros(len(csv['MFCC_02'].to_numpy().T))
            self.mfcc = np.vstack((self.mfcc, csv['MFCC_02'].to_numpy().T))
            self.mfcc = np.vstack((self.mfcc, csv['MFCC_03'].to_numpy().T))
            self.mfcc = np.vstack((self.mfcc, csv['MFCC_04'].to_numpy().T))
            self.mfcc = np.vstack((self.mfcc, csv['MFCC_05'].to_numpy().T))
            self.mfcc = np.vstack((self.mfcc, csv['MFCC_06'].to_numpy().T))
            self.mfcc = np.vstack((self.mfcc, csv['MFCC_07'].to_numpy().T))
            self.mfcc = np.vstack((self.mfcc, csv['MFCC_08'].to_numpy().T))
            self.mfcc = np.vstack((self.mfcc, csv['MFCC_09'].to_numpy().T))
            self.mfcc = np.vstack((self.mfcc, csv['MFCC_10'].to_numpy().T))
            self.mfcc = np.vstack((self.mfcc, csv['MFCC_11'].to_numpy().T))
            self.mfcc = np.vstack((self.mfcc, csv['MFCC_12'].to_numpy().T))
            self.mfcc = np.vstack((self.mfcc, csv['MFCC_13'].to_numpy().T))
            self.mfcc = np.vstack((self.mfcc, np.zeros(len(csv['MFCC_02'].to_numpy().T))))
            self.mfcc = np.vstack((self.mfcc, np.zeros(len(csv['MFCC_02'].to_numpy().T))))
            self.mfcc = np.vstack((self.mfcc, np.zeros(len(csv['MFCC_02'].to_numpy().T))))
            self.mfcc = np.vstack((self.mfcc, np.zeros(len(csv['MFCC_02'].to_numpy().T))))
            self.mfcc = np.vstack((self.mfcc, np.zeros(len(csv['MFCC_02'].to_numpy().T))))
            self.mfcc = np.vstack((self.mfcc, np.zeros(len(csv['MFCC_02'].to_numpy().T))))
            self.mfcc = np.vstack((self.mfcc, np.zeros(len(csv['MFCC_02'].to_numpy().T))))


# funkce pro extrakci parametru
# y - analyzovana data
# stngs - nastaveni analyzovanych dat
def rms(y, stngs):
    rms = librosa.feature.rms(y=y, frame_length=stngs.frame_length, hop_length=stngs.hop_length)
    return rms


def spec_cent(y, stngs):
    cent = librosa.feature.spectral_centroid(y=y, sr=stngs.sr, n_fft=stngs.frame_length, hop_length=stngs.hop_length)
    return cent


def spec_band(y, stngs):
    band = librosa.feature.spectral_bandwidth(y=y, sr=stngs.sr, n_fft=stngs.frame_length, hop_length=stngs.hop_length)
    return band


def spec_roll(y, stngs):
    roll = librosa.feature.spectral_rolloff(y=y, sr=stngs.sr, n_fft=stngs.frame_length, hop_length=stngs.hop_length)
    return roll


def chroma_stft(y, stngs):
    chroma_stft = librosa.feature.chroma_stft(y=y, sr=stngs.sr, n_fft=stngs.frame_length, hop_length=stngs.hop_length)
    return chroma_stft


def chroma_cens(y, stngs):
    chroma_cens = librosa.feature.chroma_cens(y=y, sr=stngs.sr, hop_length=stngs.hop_length, win_len_smooth=None)
    return chroma_cens


def mfcc(y, stngs):
    mfcc = librosa.feature.mfcc(y=y, sr=stngs.sr, n_fft=stngs.frame_length, hop_length=stngs.hop_length)
    return mfcc


def zero(y, stngs):
    zero = librosa.feature.zero_crossing_rate(y=y, frame_length=stngs.frame_length, hop_length=stngs.hop_length)
    return zero


def plp(y, stngs):
    plp = librosa.beat.plp(y=y, hop_length=stngs.hop_length)
    return plp


def tempogram(y, stngs):
    onset_env = librosa.onset.onset_strength(y, sr=stngs.sr)
    tempogram = librosa.feature.tempogram(onset_envelope=onset_env, sr=stngs.sr, hop_length=stngs.hop_length)
    return tempogram


# extrakce melodie
# nepouzitelne pro polyfonni nahravku
def pyin(y, stngs):
    pyin = librosa.pyin(y=y,
                        fmin=librosa.note_to_hz('G3'),
                        fmax=librosa.note_to_hz('A6'),
                        # max_transition_rate=2.5,
                        frame_length=stngs.frame_length,
                        hop_length=stngs.hop_length
                        )
    return pyin


# dynamicke tempo se vypocitava pomoci nabehove obalky
# glob - statisticka metoda ureni globalniho tempa
def tempo(y, stngs, glob=None):
    onset_env = librosa.onset.onset_strength(y, sr=stngs.sr)
    prior_lognorm = scipy.stats.lognorm(loc=np.log(120), scale=120, s=1)
    tempo = librosa.beat.tempo(onset_envelope=onset_env,
                               sr=stngs.sr,
                               hop_length=stngs.hop_length,
                               aggregate=glob,
                               prior=prior_lognorm)
    return tempo
