import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
from scipy.io import wavfile
import librosa
import librosa.display
from scipy.constants import pi
import scipy

gamma = 0.3
k_max = 15 # k_max+1 kerrosta

# read in the audio file
fs,x = wavfile.read('/Users/mac/OneDrive - TUNI.fi/SGN/Audio/Projekti/rhythm_birdland.wav')
# x = x[0:fs]
f, t, F = signal.stft(x, fs, nperseg=1000)
amp = 2 * np.sqrt(2)

# print(np.zeros((2,3)))=[0 0 0;0 0 0]

plt.figure()
plt.pcolormesh(t, f, np.abs(F), vmin=0, vmax=amp*10)
plt.title('STFT Magnitude')
plt.ylabel('Frequency [Hz]')
plt.xlabel('Time [sec]')
plt.colorbar()
plt.show()

W = (np.abs(F))**(2*gamma)

H = 0.5*W
P = H

alpha = 0.3 # (np.var(P))/(np.var(H)+np.var(P)) # 0.3

delta = np.zeros((len(f-2),len(t-2))) # Delta tulee yhden jaljessa -> deltaa on k_max kerrosta

print(" ")
print(np.size(x), fs)
print(F.shape)
print(delta.shape)
print(" ")

for k in range(0,k_max):

    H_im1 = np.c_[ np.zeros(len(f)), H[:,1:len(t)] ]
    H_ip1 = np.c_[ H[:,0:len(t)-1], np.zeros(len(f))]
    # print(H_im1)

    P_hm1 = np.vstack([ np.transpose(np.zeros(len(t))), P[0:len(f)-1] ])
    P_hp1 = np.vstack([ P[1:len(f)], np.transpose(np.zeros(len(t))) ])
    #print(P_hm1)

    delta = alpha*(H_im1-2*H+H_ip1)/4 - (1-alpha)*(P_hm1-2*P+P_hp1)/4

    if k == k_max-1: # For luupin viimeisella kerralla vaihetta 7 varten
        H_km1 = H   # Tallennetaa H ja P apumuuttujiin
        P_km1 = P

    # Kaava 26
    H = H + delta
    H = np.maximum(H, np.zeros(H.shape))
    H = np.minimum(H, W)

    # Kaava 27
    P = W - H

# Kaava 28: Rumasti tehtyna
for h in range(0, len(f)-1):
    for i in range(0, len(t)-1):
        if H[h][i] < P[h][i]:
            H[h][i] = 0
            P[h][i] = W[h][i]
        else:
            P[h][i] = 0
            H[h][i] = W[h][i]

#if H_km1 < P_km1:
 #   H = 0
  #  P = W

#else:
 #   P = 0
  #  H = W


h = scipy.signal.istft(H**(1/(2*gamma))*np.exp(1.j*np.angle(F)))
print("Eka ISTFT tehty. Koko h = ", np.size(h))
p = scipy.signal.istft(P**(1/(2*gamma))*np.exp(1.j*np.angle(F)))
print("Tonen ISTFT tehty. Koko p = ", np.size(p))


# s = audio["samples"]
# s_q = np.square(s)
# y = h+p

# SNR_h = 10*np.log*()

plt.figure()
plt.pcolormesh(t, f, np.abs(H), vmin=0, vmax=amp*10)
plt.title('Harmonic')
plt.ylabel('Frequency [Hz]')
plt.xlabel('Time [sec]')
plt.colorbar()


plt.figure()
plt.pcolormesh(t, f, np.abs(P), vmin=0, vmax=amp*10)
plt.title('Percussive')
plt.ylabel('Frequency [Hz]')
plt.xlabel('Time [sec]')
plt.colorbar()
plt.show()


# Record the audio file
audio_name = 'Harmonic'
h = h / np.max(np.abs(h))
wavfile.write(audio_name, fs, h.T)
audio_name = 'Percussive'
p = p / np.max(np.abs(p))
wavfile.write(audio_name, fs, p.T)