#!/usr/bin/env python
 
# import usefull modules
import matplotlib 
import numpy as nm
from pylab import *
import os.path
from scipy import signal
from phys import radabs
                                                                                         
#figure(figsize=(6, 4.5))

#pathobs = os.path.expanduser('/media/victor/main/Files/main/observdata/vilui15/')
pathobs = os.path.expanduser('~/Files/main/observdata/vilui15/')
#pathobs = '../results/Kuivajarvi2013_base/'

nmeas = 1 #Number of files


files = ['']*nmeas
#file =  ['KJdata2013.dat'] #Kuivajarvi
file =  ['all_trunc.dat'] #Vilui2015
auxfile =  'rad1m.dat' #Vilui2015
#file =  ['T_ML_mod.csv']
#files = ['Gill/S/01/','Gill/N/01/', 'Gill/C/tor_down/01/', 'Gill/C/ars_up/01/', \
#         'Gill/windsonicwest/','Gill/windsoniccenter/','Gill/windsoniceast/']
delim = ','
#paths = []*nmeas
paths = [pathobs]

# 3D component Gill
#n30 = 0
#n31 = 2
# 2D component Gill
#n20 = 3
#n21 = 5

#for i in range(n31-n30+1):
#        file.append('22150201_checked.CSV') 
#for i in range(n21-n20+1):
#        file.append('20150201_checked.CSV') 

fid_data = []
for i in range(nmeas):
        files[i] = paths[i] + file[i]
        fid_data.append(open(files[i],'r'))
if 'auxfile' in locals():
        fid_aux = open(pathobs+auxfile,'r')

#Column numbers of variables for spectrum calculation
#colu = [2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17]
#colu = [13,14,15] #Kuivajarvi
colu = [8,9]#[5,6,7,8,9] #Vilui2015
#colu = [1]#[5,6,7,8,9] 
#Depths
#z = [0.2,0.5,1,1.5,2,2.5,3,3.5,4,4.5,5,6,7,8,10,12] 
#z = [6,7,8] #Kuivajarvi
#z = [0.2,1.15,2.1,3.05,4] #Vilui2015
#legend = ['6 m','7 m','8 m']#Kuivajarvi
legend = ['4 m'] #['0.2 m','1.15 m','2.1 m','3.05 m', '4 m']#Vilui2015
depth = 4. #depth of measurements

nvars = len(colu) #Number of variables for spectrum calculation in each file

#Measurement data
U = np.zeros((nvars,nmeas))

# Number of periods for spectra calculation, the resulting spectrum is an averaged one
nstart = 1 #The first data item to start spectrum calculation
ntime = 1 #The number of time slots for spectrum calculation, for subsequent averaging between spectra
#tsteprat = 20
taver = 10*24*60 #Length of spectrum calculation period, data items, Vilui2015
#taver = 180*48 #Length of spectrum calculation period, data items, Kuivajarvi2013
dt = 60 #sec, high-frequency datam Vilui2015
#dt = 60*30 #sec, high-frequency data Kuivajarvi
tskip = 0 #the number of values to be skipped in high-frequency data
dirout = paths

# Frequency scale for spectra
taverm1 = taver - 1
frec3d = np.linspace(2*pi/(taverm1*dt),2*pi*(taverm1-1)/(taverm1*dt),taverm1-1) #Frequencies, s
dfrec = 2*pi/(taverm1*dt)
per3d = 2*pi/frec3d/3600. #Periods, hours
#frec2d = np.linspace(2*pi/(taverm1*dt),2*pi*(taver/tsteprat-2)/(taverm1*dt),taver/tsteprat-2)

tstep = [taver]*nmeas
#for i in range(n31-n30+1):
#        tstep.append(taver) #averaging time for 3D Gill data
#for i in range(n21-n20+1):
#        tstep.append(tstep[n30]/tsteprat) #ratio of frequencies of 3D to 2D Gill measurements


#xlabel('$x$')
#ylabel('$y$')
#axis('image')
u = np.zeros((taver,nvars))
u_aux = np.zeros(taver)
# Time-averaged spectra
ucsum = np.zeros((taver,nvars,nmeas))
#vcsum = np.zeros((taver,nmeas))
#Kolmogoroff scaling
#e53 = 1.E-11*frec3d**(-5./3.)
#e2 = 1.E-11*frec3d**(-2.)
fig = matplotlib.pyplot.figure()
for i in range(ntime):
        for j in range(nmeas):
                #Jumping to the first data item, to start calculation with
                if i == 0:
                        for k in range(nstart):
                                fid_data[j].readline()
                Usum = nm.zeros(nvars)
                n = 0
                for k in range(tstep[j]):
                        work = fid_data[j].readline()
                        #Skipping data only for high-frequency data
                        if tskip > 0:
                                for nn in range(tskip-1):
                                        work = fid_data[j].readline()
                        work = work.split(delim)
                        #print work
                        for m in range(nvars):
                                if work[colu[m]-1] != 'NaN':
                                        u[n,m] = float(work[colu[m]-1])
                                else:
                                        u[n,m] = u[n-1,m]

                        Usum[0:nvars] += u[n,0:nvars]

                        #Assuming nmeas = 1
                        if 'auxfile' in locals():
                                u_aux[n] = float(fid_aux.readline())

                        n += 1

                                

                #Calculating averages
                #U[0:nvars,j] = Usum[0:nvars]/n
                #Calculating deviations
                #u[:n-1] = u[:n-1] - U[j]
                #v[:n-1] = v[:n-1] - V[j]

		#Preprocessing the data
		#Taking time derivative
                #dTdz = 2.6 #temperature gradient in the pycnocline
                dz = 0.95
		u_ = u
		for k in range(n-1):
                        dTdz = (u_[k,nvars-1] - u_[k,nvars-2])/dz
			u[k,nvars-1] = ((u_[k+1,nvars-1] - u_[k,nvars-1])/dt - radabs(u_aux[k],depth))/dTdz			
			#u[k,nvars-1] = ((u_[k+1,nvars-1] - u_[k,nvars-1])/dt)/dTdz			
                dTdz = (u_[n-2,nvars-1] - u_[n-2,nvars-2])/dz
		u[n-1,nvars-1] = ((u_[n-1,nvars-1] - u_[n-2,nvars-1])/dt - radabs(u_aux[n-2],depth))/dTdz
		#u[n-1,nvars-1] = ((u_[n-1,nvars-1] - u_[n-2,nvars-1])/dt)/dTdz

                #Detrending the series
                for k in range(nvars):
                        u[:n-1,k] = signal.detrend(u[:n-1,k], type='linear')
                #Creating periodic series
                for m in range(n):
                        u[m,0:nvars] = u[m,0:nvars] - \
                        u[0,0:nvars] - (u[n-1,0:nvars] - u[0,0:nvars])/(n-1)*m
                #Fast Fourier transform
                nm1 = n - 1
                uc = nm.zeros((nm1,nvars))
                for k in range(nvars):
                        uc[:,k] = np.fft.fft(u[:,k],nm1)
                #Calculating power spectra
                for m in range(nvars):
                        for k in range(nm1):
                                uc[k,m] = abs(uc[k,m])**2/nm1**2
                if nm1%2 == 0: #even length
                        for k in range(nm1/2-1):
                                uc[k,0:nvars] = uc[k,0:nvars] + uc[nm1-1-k,0:nvars]
                        uc[nm1/2:,0:nvars] = 0.
                        nsp = nm1/2-1
                else: #odd length
                        for k in range(int((nm1-1)/2-1)):
                                uc[k,0:nvars] = uc[k,0:nvars] + uc[nm1-1-k,0:nvars]
                        uc[(nm1-1)/2+1:,0:nvars] = 0.
                        nsp = (nm1-1)/2
                #print sum(u*u)-sum(uc), sum(v*v)-sum(vc)
                u[:] = 0
                if j == 0: #not displaying the 0-th Fourier coefficient
                        #for k in range(nvars):
                        k = nvars-1
                        uc[1:,k] = uc[1:,k]/dfrec
                        #plot = matplotlib.pyplot.plot(per3d,uc[1:,k],lw=2)#frec3d
                        plot = matplotlib.pyplot.plot(frec3d[0:nsp],uc[1:nsp+1,k],lw=2)
                else:
                        k = nvars-1
                        #for k in range(nvars):
                        uc[1:,k] = uc[1:,k]/dfrec
                        #plot = matplotlib.pyplot.plot(per3d,uc[1:,k],lw=2)#frec3d
                        plot = matplotlib.pyplot.plot(frec3d[0:nsp],uc[1:nsp+1,k],lw=2)
                #Running mean
                #naver = 10
                #runm = np.convolve(uc[1:,k], np.ones((naver,))/naver, mode='same')
                #matplotlib.pyplot.plot(per3d,runm,lw=2,color='r')#frec3d

                # Summation of spectra for averaging
                for m in range(nvars):
                        for k in range(nm1):
                                ucsum[k,m,j] += uc[k,m]
        #-5/3 spectrum
        if 'e53' in locals():
                matplotlib.pyplot.plot(frec3d,e53,'g-',lw=2)
        #-2 spectrum
        if 'e2' in locals():
                matplotlib.pyplot.plot(frec3d,e2,'r-',lw=2)
        matplotlib.pyplot.xscale('log')
        matplotlib.pyplot.yscale('log')
        matplotlib.pyplot.xlabel('Frequency, $s^{-1}$',fontsize=20)
        #matplotlib.pyplot.xlabel('Period, hour',fontsize=16)
        matplotlib.pyplot.ylabel('Spectral density, $m^2/s$',fontsize=20)
        #plot.tick_params(axis='both', labelsize=14)
        matplotlib.pyplot.tick_params(axis='both', labelsize=18)
        #matplotlib.pyplot.legend(legend)
        for i in range(nvars): print('Average spectrum level = ', np.mean(uc[1:nsp+1,i]))
        # Calculating horizontal kinetic energy
        EH = 0.
        EV = 0.
        fcor = 1.19E-4#Coriolis parameter
        N = 0.2 #Brunt Vaisala
        sphor = np.zeros(nm1)
        for i in range(taverm1-1):
                if (frec3d[i] > fcor and frec3d[i] < N):
                        EV += uc[i,nvars-1]*dfrec
                        EH += uc[i,nvars-1]/(frec3d[i]**2*(frec3d[i]**2 - fcor**2))* \
                        ((N**2 - frec3d[i]**2)*(frec3d[i]**2 + fcor**2))*dfrec
                        sphor[i] = uc[i,nvars-1]/(frec3d[i]**2*(frec3d[i]**2 - fcor**2))* \
                        ((N**2 - frec3d[i]**2)*(frec3d[i]**2 + fcor**2))
        print((taverm1-1)*dfrec)
        print('Horizontal kinetic energy ' + str(EH))
        print('Vertical kinetic energy ' + str(EV))
        matplotlib.pyplot.plot(frec3d[0:nsp],sphor[1:nsp+1],lw=2)
        #Vertical line
        #matplotlib.pyplot.axvline(x=8.5e-5,lw=2, color='k')
        #matplotlib.pyplot.text(9.e-5,1.e-2,'$\omega=8.5*10^{-5} s^{-1}$',fontsize=16)
        #matplotlib.pyplot.axvline(x=4.5e-4,lw=2, color='k')
        #matplotlib.pyplot.text(5.e-4,0.8e-2,'$\omega=4.5*10^{-4} s^{-1}$',fontsize=16)
        show()
        #date = work[0].split('/')
        #date = date[0] + '-' + date[1] + '-' + date[2]
        #print date,work[1]
        savefig(paths[0]+'spectra.png')
        #close()

#Plotting averaged spectra
#ucsum[:,:,:] = ucsum[:,:,:]/ntime
#for j in range(nmeas):
#        if j == 0:
#                work1 = []
#                for i in range(len(frec3d)):
#                        work1.append(ucsum[i+1,m,j])
#                plotid = matplotlib.pyplot.plot(frec3d,work1)
#        else:
#                if j <= n31: #3D Gill
#                        work1 = []
#                        for i in range(len(frec3d)):
#                                work1.append(ucsum[i+1,j]+vcsum[i+1,j])
#                        matplotlib.pyplot.plot(frec3d,work1)
#                else: #2D Gill
#                        work1 = []
#                        for i in range(len(frec2d)):
#                                work1.append(ucsum[i+1,j]+vcsum[i+1,j])
#                        matplotlib.pyplot.plot(frec2d,work1)
##-5/3 law
#matplotlib.pyplot.plot(frec3d,e53,'b-',lw=2)
#matplotlib.pyplot.title('Averaged spectra')
#matplotlib.pyplot.xscale('log')
#matplotlib.pyplot.yscale('log')
#matplotlib.pyplot.xlabel('Frequency, $s^{-1}$')
#matplotlib.pyplot.ylabel('Spectral coefficients, $m^2/s^2$')
#show()
#date = work[0].split('/')
#date = date[0] + '-' + date[1] + '-' + date[2]
#savefig(dirout + 'spectra_aver_' + 'taver_' + str(taver) + '_uv_' + date+'_' + work[1] + '.png')


#U = [1,-1,1,-1,1,-1]
#V = [-1,1,-1,1,-1,1]


