
import sys; sys.path.append('../py_src')
from pyOM_gui import pyOM_gui as pyOM
from numpy import *
import lin_stab

GRID_POINTS = 25       # number of grid points in lateral direction
EDDY_LENGTH_SCALES = 4 # number of eddies in domain

if 0: # meso-scale regime
  Ro  = 0.05
  Ek  = 0.001
  CFL = 0.02    # CFL number
  delta = 0.01 # aspect ratio

if 0: # quite sub-meso-scale regime
  Ro  = 0.1
  Ek  = 0.01
  CFL = 0.05   # CFL number
  delta = 0.02 # aspect ratio
if 1: # very sub-meso-scale regime
  Ro  = 0.5
  Ek  = 0.01   # without quicker mom adv, else 0.01
  CFL = 0.05   # CFL number
  delta = 0.01 # aspect ratio

H0 = 500.0
beta= 0e-11  # beta 
f0  = 1e-4   # coriolis freq.
N0=f0/delta  # stability freq.
U0 = Ro*N0*H0

class channel1(pyOM):
   """ channel with restoring zones
   """
   def set_parameter(self):
     """set main parameter
     """
     M=self.fortran.pyom_module   
     M.nx    = GRID_POINTS+2
     M.nz    = 12+2
     M.ny    = GRID_POINTS+2
     M.dz    = H0/(M.nz-2)
     M.enable_hydrostatic           = 1
     M.enable_cyclic_x              = 1
     M.enable_noslip                = 1
     M.enable_quicker_advection     = 1
     M.enable_quicker_mom_advection = 1
     M.enable_back_state            = 1
     M.enable_back_zonal_flow       = 1

     # print some numbers
     Lr    = N0*H0/f0          # Rossby radius
     Ro    = U0/(f0*Lr)       # Rossby number
     Ri    = N0**2*H0**2/U0**2 # Richardson number
     print
     print  ' L  = %f km'%(Lr/1e3)
     print  ' Ro = %f '%Ro
     print  ' Ri = %f '%Ri
     print  ' delta = %f '%delta
     print  ' ell = %f '%(Lr/6400e3)

     # solve linear stability problem first
     ksx=linspace(0,3.2,15);    kx=ksx/Lr
     ksy=linspace(-3.2,3.2,30); ky=ksy/Lr
     zw=arange(M.nz)*M.dz
     zt=zw-M.dz/2.0
     U=U0/2+U0*zt[1:-1]/H0
     V=U*0
     B=N0**2*zt[1:-1]
     om_max,om,kmax,lmax,u,v,w,b,p=lin_stab.pe(U,V,B,M.dz,\
                                                kx,ky,0.,beta,0.,f0,0.,0)
     print ' Max. growth rate %f 1/days ' % (-imag(om)*86400)
     print ' k_max = %f Lr , l_max = %f Lr' % (kmax*Lr,lmax*Lr)

     self.lin_stab_kmax = kmax
     self.lin_stab_b = b
     L = EDDY_LENGTH_SCALES*2*pi/kmax   
     M.dx  =   L/(M.nx-2) 
     M.dt    = CFL/U0*M.dx   # CFL=U*dt/dx  
     print " dx=%f km, dt= %f s "%(M.dx/1e3,M.dt)
     print " CFL  = ",U0*M.dt/M.dx
     print " CFL  = ",real(om)/kmax*M.dt/M.dx
     M.eps2d_sor = 1e-6 *(M.dx/20e3)**2

     M.a_h = Ek*f0*M.dx**2 
     M.k_h = Ek*f0*M.dx**2 
     M.a_v = Ek*f0*M.dz**2 
     M.k_v = Ek*f0*M.dz**2 
     print " A_h = %f m^2/s  Ek = %f"%(M.a_h,Ek)
     print " A_v = %f m^2/s  Ek = %f"%(M.a_v,Ek)

     return

   def set_coriolis(self):
     """ vertical and horizontal Coriolis parameter on yt grid
         routine is called after initialization of grid
     """
     M=self.fortran.pyom_module   
     M.coriolis_t[:]   =  f0+beta*M.yt[:]
     M.coriolis_hor[:] =  0.
     return

   def initial_conditions(self):
     """ setup all initial conditions
     """  
     M=self.fortran.pyom_module   
 
     # background velocity
     for k in range(M.nz):
        M.u0[:,:,k]=(U0/2+U0*M.zt[k]/H0)*M.masku[:,:,k]

     # background buoyancy
     for k in range(1,M.nz-1):
       for j in range(M.ny-1):
        fxa = (M.coriolis_t[j]+M.coriolis_t[j+1])/2.0
        uz1 = (M.u0[:,j,k+1]-M.u0[:,j,k])/M.dz*M.masku[:,j,k+1]
        uz2 = (M.u0[:,j,k]-M.u0[:,j,k-1])/M.dz*M.masku[:,j,k-1]
        uz3 = (M.u0[:,j+1,k+1]-M.u0[:,j+1,k])/M.dz*M.masku[:,j+1,k+1]
        uz4 = (M.u0[:,j+1,k]-M.u0[:,j+1,k-1])/M.dz*M.masku[:,j+1,k-1]
        uz = (uz1+uz2+uz3+uz4)/(M.masku[:,j,k+1]+M.masku[:,j,k-1] \
                          +M.masku[:,j+1,k+1]+M.masku[:,j+1,k-1]+1e-12)
        M.back[:,j+1,k,0]=M.back[:,j,k,0]+M.dx*M.maskt[:,j,k]*uz*fxa

     for k in range(M.nz):
        M.back[:,:,k,0]=(M.back[:,:,k,0]-N0**2*M.zt[k])*M.maskt[:,:,k]
     M.back[:,:,:,1] = M.back[:,:,:,0]    
     M.back[:,:,:,2] = M.back[:,:,:,0]


     # perturbation buoyancy
     for j in range(M.ny):
      for k in range(1,M.nz-1):
       ky=pi/((M.ny-2)*M.dx) 
       phase = cos(self.lin_stab_kmax*M.xt) \
       +complex(0,1)*sin(self.lin_stab_kmax*M.xt)
       phase = phase*sin(ky*M.yt[j])
       M.b[:,j,k,0] =0.5*real(phase*self.lin_stab_b[k-1])*M.maskt[:,j,k]

     for a in M.b,M.u,M.v,M.w,M.p_full:
       a[:,:,:,1] = a[:,:,:,0]    
       a[:,:,:,2] = a[:,:,:,0]
     return

   def restoring_zones(self):
     """  add here restoring zones
     """
     M=self.fortran.pyom_module
     if 0:  # restoring at side walls
      t0 = 2./500.0 
      j=3
      M.b[:,:j,:,M.taup1-1]=M.b[:,:j,:,M.taup1-1] \
           - t0*M.b[:,:j,:,M.taum1-1]*M.maskt[:,:j,:]
      M.b[:,-j:,:,M.taup1-1]=M.b[:,-j:,:,M.taup1-1] \
           - t0*M.b[:,-j:,:,M.taum1-1]*M.maskt[:,-j:,:]
     if 1:  # interior restoring
      t0 = 2./1000.0 
      M.b[:,:,:,M.taup1-1]=M.b[:,:,:,M.taup1-1] \
            - t0*M.b[:,:,:,M.taum1-1]*M.maskt[:,:,:]
     if 0:  # damping at side walls
      t0 = 2./50.0 
      j=4
      M.u[:,:j,:,M.taup1-1]=M.u[:,:j,:,M.taup1-1] \
           - t0*M.u[:,:j,:,M.taum1-1]*M.masku[:,:j,:]
      M.u[:,-j:,:,M.taup1-1]=M.u[:,-j:,:,M.taup1-1] \
           - t0*M.u[:,-j:,:,M.taum1-1]*M.masku[:,-j:,:]
      M.v[:,:j,:,M.taup1-1]=M.v[:,:j,:,M.taup1-1] \
           - t0*M.v[:,:j,:,M.taum1-1]*M.maskv[:,:j,:]
      M.v[:,-j:,:,M.taup1-1]=M.v[:,-j:,:,M.taup1-1] \
           - t0*M.v[:,-j:,:,M.taum1-1]*M.maskv[:,-j:,:]
     return


   def make_plot(self):
     """ make a plot using methods of self.figure
     """
     if hasattr(self,'figure'):
       M=self.fortran.pyom_module         # fortran module with model variables
       k=M.nz*3/4
       i=int(M.nx/2)
       j=int(M.nx/2)
       x=M.xt[1:-1]/1e3
       y=M.yt[1:-1]/1e3
       z=M.zt[1:-1]

       self.figure.clf()
       if 1:
         ax=self.figure.add_subplot(221)
         a=M.b[i,1:-1,1:-1,M.tau-1] + M.back[i,1:-1,1:-1,M.tau-1]
         co=ax.contourf(y,z,a.transpose())
         a=M.u[i,1:-1,1:-1,M.tau-1] + M.u0[i,1:-1,1:-1]
         ax.contour(y,z,a.transpose(),10,colors='black')
         ax.set_ylabel('y [km]')
         self.figure.colorbar(co)
         ax.set_title('total buoyancy');  
         ax.axis('tight') 
         
       ax=self.figure.add_subplot(222)
       co=ax.contourf(x,y,M.b[1:-1,1:-1,k,M.tau-1].transpose()*1e2)
       a=M.b[1:-1,1:-1,k,M.tau-1] + M.back[1:-1,1:-1,k,M.tau-1]
       ax.contour(x,y,a.transpose(),10,colors='black')
       ax.set_title('buoyancy perturbation'); 
       self.figure.colorbar(co)
       ax.axis('tight')

       if hasattr(self,'mean_itt'):  # plot means
         ax=self.figure.add_subplot(223)
         a=self.mean_vb[1:-1,1:-1] 
         if self.mean_itt >0: a=a/self.mean_itt 
         co=ax.contourf(y,z,a.transpose())
         ax.set_ylabel('y [km]')
         self.figure.colorbar(co)
         ax.set_title('mean buoyancy');
         ax.axis('tight')

       ax=self.figure.add_subplot(224)
       a=M.b[1:-1,1:-1,1:-1,M.tau-1] + M.back[1:-1,1:-1,1:-1,M.tau-1]
       b=M.maskt[1:-1,1:-1,1:-1]
       a=sum(a*b,axis=0)/sum(b,axis=0)
       co=ax.contourf(M.yt[1:-1]/1e3,M.zt[1:-1],a.transpose())
       self.figure.colorbar(co)
       ax.set_title('zonal average');
       ax.axis('tight')
     return

   def diagnose(self):
     pyOM.diagnose(self)
     M=self.fortran.pyom_module   
     if 1: # calculate means
        if hasattr(self,'mean_itt'):
          self.mean_itt = self.mean_itt+1
          b=M.b[1:-1,:,:,M.tau-1]
          v=M.v[1:-1,:,:,M.tau-1]
          mt=M.maskt[1:-1,:,:]
          mv=M.maskv[1:-1,:,:]
          bm=sum(b*mt,axis=0)/sum(mt,axis=0)
          vm=sum(v*mv,axis=0)/sum(mv,axis=0)
          vb=0*v  
          for i in range(M.nx-2):
            vb[i,1:,:]= (v[i,1:,:]-vm[1:,:])*(b[i,:-1,:]-bm[:-1,:])
          vb=sum(vb*mv,axis=0)/sum(mv,axis=0)
          self.mean_b  += bm
          self.mean_v  += vm
          self.mean_vb += vb
        else:
          self.mean_itt = 0
          self.mean_b  = zeros( (M.ny,M.nz), 'd' )
          self.mean_v  = zeros( (M.ny,M.nz), 'd' )
          self.mean_vb = zeros( (M.ny,M.nz), 'd' )
          self.mean_wb = zeros( (M.ny,M.nz), 'd' ) 

     if 0:  # check some numerical criteria
        c=M.dx/M.dt 
        a=(M.u[:,:,:,M.tau-1]+M.u0[:,:,:])/c
        a=amax(amax(amax(a,axis=0),axis=0),axis=0)
        print " maximal CFL number ",a
        if M.dt > M.dx**2/(2.*max(M.k_h,1e-8)): print "K_h too large"
        if M.dt > M.dx**2/(2.*max(M.a_h,1e-8)): print "A_h too large"
        if M.dt > M.dz**2/(2.*max(M.k_v,1e-8)): print "K_v too large"
        if M.dt > M.dz**2/(2.*max(M.a_v,1e-8)): print "A_v too large"
        a=(M.u[:,:,:,M.tau-1]+M.u0[:,:,:])*M.dx/M.k_h
        a=amax(amax(amax(a,axis=0),axis=0),axis=0)
        print " maximal hor. Peclet number ",a
     

if __name__ == "__main__": 
         A=channel1(snapint = 5, plotint= 5*24)
         A.mainloop()


