import numpy as np
import scipy.fftpack as fftpack
import librosa
import scipy
import matplotlib.pyplot as plt


def librosaSpec(signal, frame_size, hop_size, n_bands, sr):  #mel STFT
    power_spec = (2 * np.abs(librosa.stft(signal, n_fft=frame_size, hop_length=hop_size))) ** 2
    filterbank = librosa.filters.mel(sr, n_fft=frame_size, n_mels=n_bands, fmin=30., fmax=17000., norm=1)
    filt_spec = np.log(np.dot(power_spec.T, filterbank.T) + 1.0) #% +1.0
    filt_spec = np.transpose(filt_spec)
    return filt_spec


def median_spec(melgram, frame_size): #median STFT
    a = []
    kernel = frame_size/100
    kernel = int(round_up_to_odd(kernel))
    for i in range(int(melgram.shape[0])):
        b = scipy.signal.medfilt(melgram[i, :], kernel)
        a.append(b)
    diff = melgram - a
    median_spec = (np.abs(diff) + diff) / 2.0
    return median_spec


def tkeo2(a): # #1 tkeo var
    l = 1
    p = 2
    q = 0
    s = 3
    aTkeo = a[l:-p] * a[p:-l] - a[q:-s] * a[s:]
    return aTkeo


def tkeo(a): # #2tkeo var
    i = a[1:-1]*a[1:-1]
    j = a[2:]*a[:-2]
    aTkeo = i-j
    return aTkeo


def round_up_to_odd(f):
    return np.ceil(f) // 2 * 2 + 1


def sec_to_frame(num_frames, anotace):   # transform beat times for NN
    frame_pos = np.zeros(num_frames)
    for idx, val in enumerate(frame_pos):
        if idx in anotace:
            frame_pos[idx] = 1
            if idx + 1 < num_frames:
                frame_pos[idx + 1] = 0.5
                frame_pos[idx - 1] = 0.5
    labels = np.array(frame_pos)
    return labels


def anot_to_vec(anotace, frame_rate=100):  #transform txt to vector
    label = list()
    with open(anotace) as anot:
        for line in anot:
            label.append(line.strip())
    label = np.array(label, dtype=np.float32)
    label = label * frame_rate
    label = np.ceil(label)
    return label


def show_spec(input):  #show spectogram
    plt.imshow(input, origin='lower', aspect='auto', cmap='jet')
    plt.colorbar(format='%.1f')
    plt.xlabel('Časový segment [-]')
    plt.ylabel('Řadové číslo spektrogramu [-]')
    plt.yticks(ticks=[0, 20, 40, 60, 80, 100, 120], labels=[0, 1, 2, 3, 4, 5, 6])
    plt.tight_layout()


def get_filename(index, data):
    file = data.at[index, 'filename']
    return file