def resid(w,bvf_,k,dz,omega,N):
        for i in range(N-2):
                w[i+2] = dz**2*w[i+1]*k**2*(1. - bvf_[i+1]**2/omega**2) + 2*w[i+1] - w[i]
        return w[N-1] - 0.

import numpy as nm
import pylab as pl
import phys
import matplotlib as mpl
import time 

#Solution of Sturm-Liouville problem to find seiche frequencies,
# following (Muennich et al., 1992)

expname = 'BolshoiVilui20153_base'
nyear = '2015'
nmonth = '07'
nday = '25'
#nheader = 12 #number of lines in header for files with monthly data
nheader = 7 #number of lines in header for files with daily data

h = 7. #6.4, 12.5, lake depth,m
L = 5.E+3 #2.6E+3lake length,m
L1 = 2.5E+2 #lake width, m
N = 100 #number of layers
dz = h/(N-1) #grid spacing
domega = 1.E-7 #step for scanning omega values
dw0 = 1.E-7 #threshold dw
resid_max = 1.E-10
maxiter = 100

z = nm.linspace(0.,h,N)

omega_min = 1.E-6
omega_max = 1.E-3

m = 1 #along-lake horizontal wavelength number
n = 0 #transversal horizontal wavelength number
k = ((m*nm.pi/L)**2 + (n*nm.pi/L1)**2)**0.5

T=nm.zeros(N)#water temperature,Celsius

#Linear temperature profile
#Tsurf=20.
#Tbot=4.
#T=nm.linspace(Tbot,Tsurf,N)

#fid = open('../results/' + expname + '/monthly/Profiles  1  1' + nyear + nmonth + '.dat','r')
fid = open('../results/' + expname + '/daily/Profiles  1  1' + nyear + nmonth + nday + '.dat','r')
filecont = fid.read()
filecont = filecont.split('\n')
filecont = filecont[nheader:]
ndat_ = len(filecont)#number of data
print(ndat_)

#Reading temperature data
zdat_ = nm.zeros(ndat_)#depth
Tdat_ = nm.zeros(ndat_)#temperature
Sdat_ = nm.zeros(ndat_)#salinity
ndat = 0
for i in range(ndat_):
        x = filecont[i].split()
        if len(x) >= 2:
                zdat_[i] = h + float(x[0])
                Tdat_[i] =     float(x[1])
                Sdat_[i] =     float(x[2])
                #print(zdat[i],Tdat[i])
                ndat = ndat + 1

zdat = nm.zeros(ndat)
Tdat = nm.zeros(ndat)
Sdat = nm.zeros(ndat)
for i in range(ndat):
        zdat[i] = zdat_[ndat-1-i]
        Tdat[i] = Tdat_[ndat-1-i]
        Sdat[i] = Sdat_[ndat-1-i]


# Data from Muennich et al. 1992
#ndat = 6
#zdat = nm.linspace(0.,h,ndat)
#Tdat = nm.zeros(ndat)
#Tdat[0] = 6.
#Tdat[1] = 7.
#Tdat[2] = 9.
#Tdat[3] = 10.5
#Tdat[4] = 18.5
#Tdat[5] = 19.8

#Interpolation to fine grid
T = nm.interp(z,zdat[:ndat],Tdat[:ndat])
S = nm.interp(z,zdat[:ndat],Sdat[:ndat])
#print(zdat[:ndat],Tdat[:ndat])
#print(z,T)

bvf_ = nm.zeros(N)
for i in range(N-2):
        bvf_[i+1] = phys.bvf(T[i+2],T[i],S[i+2],S[i],2.*dz)

w = nm.zeros(N)#vertical velocity amplitude

w[0] = 0#bottom b.c.
dwdz0 = 1.E+3
w[1] = w[0] + dwdz0*dz

#Solving initial value problem
omega = omega_min
omegas = []
Ts = []
dw = 1.
while omega < omega_max:
        dw_new = resid(w,bvf_,k,dz,omega,N)
        #print(w)
        #time.sleep(1)
        #print(omega,w[N-1])
        if dw_new*dw < 0:
                #print('Res changed sign, starting chorde method',omega,dw)
                iter = 0
                dw3 = 1.
                dw1 = dw
                dw2 = dw_new
                omega1 = omega - domega
                omega2 = omega
                while (abs(dw3) > resid_max):
                        iter = iter + 1
                        if (dw1 != dw3):
                                dw1 = resid(w,bvf_,k,dz,omega1,N) 
                        if (dw2 != dw3):
                                dw2 = resid(w,bvf_,k,dz,omega2,N) 
                        omega3 = (omega1*dw2 - omega2*dw1)/(dw2 - dw1)
                        dw3 = resid(w,bvf_,k,dz,omega3,N)
                        if (dw1 * dw3 < 0.):
                                omega2 = omega3
                                dw2 = dw3
                        elif (dw2 * dw3 < 0.):
                                omega1 = omega3
                                dw1 = dw3
                        if (iter > maxiter):
                                print('Didnt converge',omega3,dw3)
                                break
                if (abs(dw3) < resid_max):
                        omegas.append(omega3)
                        Ts.append(2*nm.pi/omega3/3600.)
                        print('CONVERGED!:','omega=',omega3,'T=',2*nm.pi/omega3/3600.,dw3)
                        pl.plot(w)
        dw = dw_new
        omega = omega + domega

pl.legend(Ts)
pl.show()
#print(omegas)
