from scipy import io
import scipy.signal
import matplotlib.pyplot as plt
import numpy as np
import scipy.fftpack
import scipy.integrate
import pandas as pd
from IPython.display import display
import time
import matplotlib.colors as colors
import random
import glob
import os


lot_of_plots = 0

path ='C:\\Users\\Marci\\Documents\\transientwaves\\imputdata\\in\\'#pc
#path ='C:\\Users\\win10\\Downloads\\transient-waves\\inputdata\\in\\'#laptop

for filename in glob.glob(os.path.join(path, '*.sav')):
    sav1 = io.readsav(filename, python_dict = True)
    print('Loaded: '+filename)
    print('\n')


shot_i=30
s=np.array(sav1['data'][shot_i])#
x=np.array(sav1['timeax'])
sxr_name=str(sav1['channels'][shot_i])[2:4]+'R'+str(sav1['channels'][shot_i])[5:-1]
shot=os.path.basename(filename)[:4]+' '+os.path.basename(filename)[5:10]+', '+str(sav1['channels'][shot_i])[2:-1]


s = s.astype(float)
x = x.astype(float)
samples=s.size
t=max(x)-min(x)
nyquist=samples/2/t
frangemax=nyquist#250000#
xf=np.linspace(0.0, frangemax, int(frangemax*t+1))
s_rate=x.size/t
print("nyquist:  ",nyquist)
print("time len: ",x.size)

#print(sav1['data'].shape)
ne = np.zeros((100, sav1['data'].shape[1]))
freq = np.zeros((100, sav1['data'].shape[1]))
Bmod = np.zeros((100, sav1['data'].shape[1]))
#print(ne[:, 10].size)

for i in range(100):
    ne[i] = np.array(sav1['data'][i])
    freq[i] = np.array(sav1['data'][shot_i + 100])
    Bmod[i] = np.array(sav1['data'][shot_i + 200])

if lot_of_plots:
    print("aaaa")
    plt.figure(figsize=(10, 5))
    plt.imshow(ne, aspect='auto', cmap='plasma', extent=[min(x), max(x), 0, 100])
    plt.title(shot + ' szonda jele')
    plt.xlabel('Idő (ms)')
    plt.ylabel('Radiális távolság')
    cax = plt.axes([0.92, 0.15, 0.03, 0.7])
    plt.colorbar(cax=cax).set_label('')
    # plt.subplots_adjust(right=0.95)
    # plt.savefig(shot+' fullprofile.pdf')
    #plt.show()

    plt.figure(figsize=(10, 5))
    plt.plot(ne[:, 200], '.-')
    plt.title(shot + ' szonda jele')
    plt.ylabel('Sűrűség')
    plt.xlabel('Radiális távolság')
    # plt.savefig(shot+' profile 200.pdf')
    #plt.show()

    plt.figure(figsize=(10, 5))
    plt.plot(x, ne[30, :], '.-')
    plt.title(shot + ' szonda jele')
    plt.xlabel('Idő (ms)')
    plt.ylabel('Sűrűség')
    # plt.savefig(shot+' at freq 30.pdf')
    #plt.show()


###Le Fourier transform
t0=time.time()
win_r = 500 #1000
d_win=5#random.randint(200,300)
std=45 #150
win_base = scipy.signal.gaussian(2*win_r+1, std=std)#np.wining(2*win_r+1)

z_padded_s = np.zeros (samples+2*win_r)
z_0=z_padded_s
z_padded_s[win_r:-win_r] = s
z_padded_s[0:win_r]=np.average(z_padded_s[win_r+1:2*win_r])
z_padded_s[-win_r:]=np.average(z_padded_s[-2*win_r:-win_r-1])
x_win=np.linspace(0.0, t+(2*win_r/s_rate), samples+2*win_r)

win = np.zeros((int(samples/d_win+1), 2*win_r+1))
for i in range(int(samples/d_win+1)):
    win [i]=win_base


z_win = np.zeros((int(samples/d_win+0.9), win_r*2+1))
for i in range(int(samples/d_win+0.9)):
    z_win[i] = win[i] * (z_padded_s [d_win*i:d_win*i+win_r+win_r+1]-np.average(z_padded_s [d_win*i:d_win*i+win_r+win_r+1]))
    #z_win[i] = win[i] * (z_padded_s [d_win*i:d_win*i+win_r+win_r+1])

if lot_of_plots:
    plt.figure(figsize=(9,9))
    plt.imshow(z_win,aspect='auto',cmap='bwr')#, extent=(min(x)*s_rate/d_win,max(x)*s_rate/d_win,max(x)*s_rate,min(x)*s_rate)
    plt.title('Gauss ablak + eredeti jel')
    plt.xlabel('Adatpont')
    plt.ylabel('Ablak elcsúsztatása')
    plt.subplots_adjust(right=0.88)
    cax = plt.axes([0.90, 0.23, 0.03, 0.54])
    plt.colorbar(cax=cax).set_label('Amplitúdó')
    plt.savefig('z_win.pdf')
    #plt.show()

z_win_fft = np.zeros((int(samples/d_win+0.9), win_r*2+1))
for i in range(int(samples/d_win+0.9)):
    z_win_fft[i] = np.abs(scipy.fftpack.fft(z_win[i]))
t1=time.time()


z_win_fft_cutted = np.zeros((int(frangemax*(2*win_r+1)/s_rate), int(samples/d_win)))
for i in range(int(samples/d_win)):
    for j in range(int(frangemax*(2*win_r +1)/s_rate)): #int(frangemax*max(x_hann))
        z_win_fft_cutted [j][i] = z_win_fft [i][j]**2


colornorm=colors.LogNorm(vmin=z_win_fft_cutted.min()*10000, vmax=z_win_fft_cutted.max()/1)
if lot_of_plots:
    logax = np.linspace(int(np.log10(z_win_fft_cutted.min())),int(np.log10(z_win_fft_cutted.max())),int(np.log10(z_win_fft_cutted.max()))-int(np.log10(z_win_fft_cutted.min()))+1)
    logax = 10**logax
    fig, ax = plt.subplots(figsize=(10,5))
    #extent=[0,t-(2*hann_r/s_rate),0,frangemax/t],
    #plt.figure(figsize=(15,10))
    #ax.plot(x,np.abs(full_f),color='g',linewidth=3,linestyle='dashed',alpha=0.6)#img = ax.
    img = plt.imshow(z_win_fft_cutted,cmap='plasma',  aspect='auto', extent=[min(x),max(x),frangemax/1000,0], norm=colornorm)#interpolation='gaussian',*30000000/4
    plt.title(shot+' szonda saját felüláteresztős szűrt STFT-je')
    plt.ylabel('Frekvencia (kHz)')
    plt.xlabel('idő (s)')#('time (s)')
    #plt.xlim(0.642,0.658)
    #plt.ylim(90,30)
    plt.subplots_adjust(right=0.90)
    plt.gca().invert_yaxis()
    cbar=plt.colorbar(img,cax=plt.axes([0.91, 0.2, 0.02, 0.6]),ticks=None)
    cbar.set_label('Fourier-amplitúdó (J)')
    #plt.axis([2.13,2.15,min(s),20000])
    #cbar.set_ticks(logax)
    #plt.savefig(shot+' stft.pdf')#34924_40
    #plt.show()



ridge_d=2000#50#
ridge_d2=800
ridge_start=4000#2#
ridge = np.zeros(int(samples/d_win+0.9))
ridge_h = np.zeros(int(samples/d_win+0.9))
maxi=z_win_fft.max()/40
print(maxi)

for i in range(int(samples/d_win+0.9)):
    ridgeval=0  #max(z_spy_stft[i][:])
    if (i==0) or (ridge[i-1]==0):
        for j in range(int(ridge_d*(2*win_r+1)/(s_rate-0.9))):
            j2=j+int((ridge_start)*(2*win_r+1)/(s_rate-0.9))
            if ridgeval < z_win_fft[i][j2]:
                ridgeval=z_win_fft[i][j2]
                ridge[i]= j2
    else:
        for j in range(int(ridge_d2*(2*win_r+1)/(s_rate-0.9))):
            j2=int(j+ridge[i-1]-(ridge_d2/2)*(2*win_r+1)/(s_rate-0.9))
            if ridgeval < z_win_fft[i][j2]:
                ridgeval=z_win_fft[i][j2]
                ridge[i]= j2
                #print(j2, z_spy_stft[j2][i])
    if ridgeval<(maxi):
        ridgeval =0
        ridge [i]=0
    ridge_h[i]=ridgeval**(1/2)
    #print(i, ridge [i])
    #print (i,"  ",ridge[i],"\t", ridgeval,"\t", z_spy_stft[i][int(ridge[i])])

ridge=ridge*(s_rate/(2*win_r+1))

#proc_x=np.linspace(min(x),max(x),np.size(ridge))
ridge[ ridge==0 ] = np.nan
if lot_of_plots:
    plt.figure(figsize=(14,3.5))
    plt.title(shot+' szonda maximum frekvenciaértékei')
    plt.plot(np.linspace(min(x),max(x),np.size(ridge)),ridge/1000,'.-')
    plt.grid(True)
    #plt.legend(('Kiszámolt', 'Eredeti'))
    #yaxismax=1/2**(view.value/12)
    plt.xlim(min(x),max(x))
    #plt.ylim(ridge_start/1100,)
    #plt.ylim(0,200)
    #plt.xlim(0.89,0.91)
    plt.ylabel('Frekvencia (kHz)')
    plt.xlabel('idő (s)')
    #plt.savefig(shot+' f_max.pdf')
    #plt.show()


logax = np.linspace(int(np.log10(z_win_fft_cutted.min())),int(np.log10(z_win_fft_cutted.max())),int(np.log10(z_win_fft_cutted.max()))-int(np.log10(z_win_fft_cutted.min()))+1)
logax = 10**logax
fig, ax = plt.subplots(figsize=(14,5))
#extent=[0,t-(2*hann_r/s_rate),0,frangemax/t],
#plt.figure(figsize=(15,10))
ax.plot(np.linspace(min(x),max(x),np.size(ridge)),ridge/1000,'.',color='g',alpha=0.6)#,linewidth=3,linestyle='dashed')#img = ax.
img = plt.imshow(z_win_fft_cutted,cmap='plasma', interpolation='gaussian', aspect='auto', extent=[min(x),max(x),frangemax/1000,0], norm=colornorm)#
#plt.title('My Short-time Fourier transform')
plt.title(shot+'-es szonda STFT-je, adaptív maximumkereséssel')
plt.ylabel('Frekvencia (kHz)')
plt.xlabel('idő (s)')#('time (s)')
#plt.xlim(0.652,0.66)
#plt.ylim(90,40)
plt.subplots_adjust(right=0.90)
plt.gca().invert_yaxis()
cbar=plt.colorbar(img,cax=plt.axes([0.91, 0.2, 0.02, 0.6]),ticks=None)
cbar.set_label('Fourier-érték (J)')
#plt.axis([2.13,2.15,min(s),20000])
cbar.set_ticks(logax)
#plt.savefig(shot+' stft_max.pdf')
#plt.show()


ridge_h[ ridge_h==0 ] = np.nan
print(maxi**(1/2))
plt.figure(figsize=(10,3.5))
plt.title('Az amplitúdó a maximum frekvenciaértékeknél, '+shot)#, AUGD 34924 J051-es szonda
plt.plot(np.linspace(min(x),max(x),np.size(ridge_h)),ridge_h,'-',markersize=1)#,x,3*abs(np.sin(2*np.pi*x))+2,
plt.grid(True)
plt.xlim(min(x),max(x))
#plt.xlim(0.89,0.91)
#plt.ylim(maxi**(1/2)/2,max(ridge_h)*1.05)
plt.ylabel('Amplitúdó')
plt.xlabel('idő (s)')
#plt.yscale("log")
#plt.savefig(shot+'_amp.pdf')
plt.show()