
import sys; sys.path.append('../py_src')

from pyOM_cdf import pyOM_cdf as pyOM
from numpy import *


HRESOLVE = 0.5
VRESOLVE = 0.5
N_0     = 0.004
M_0     = sqrt(1e-5*0.1/1024.0*9.801)
spg_width = int(3*HRESOLVE)
t_rest=1./(5.*86400)

class jets1(pyOM):
   """ a wide channel model with zonal jets
   """
   def set_parameter(self):
     """set main parameter
     """
     M=self.fortran.pyom_module   
     M.nx  = 128*HRESOLVE 
     M.ny  = 128*HRESOLVE
     M.nz  = 20 *VRESOLVE;

     radius = 6370.0e3        # Earth radius in m
     degtom = radius/180.0*pi # conversion degrees latitude to meters
     M.dx    = 1./3.*degtom*cos(30./180.*pi)/HRESOLVE
     M.dz    = 100.0 /VRESOLVE
     M.dt    = 2400.0/HRESOLVE
     M.k_v=1.0e-4/VRESOLVE
     M.a_v=10.e-4/VRESOLVE
     M.eps2d_sor = 1e-9
     M.enable_hydrostatic          = 1
     M.enable_noslip               = 1
     M.enable_cyclic_x             = 1
     M.enable_interior_stress      = 1
     M.cdint = 1.0e-7
     M.enable_biharmonic_friction  = 1
     M.ahbi  = 5e11/HRESOLVE**2
     M.enable_biharmonic_diffusion = 1
     M.khbi  = M.ahbi
     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   
     rad  = 6370.0e3   # earth radius
     om   = pi/43082.0 # earth rotation freq. in 1/s
     phi0 = 10./180. *pi
     M.coriolis_t[:]   =  2*om*sin(phi0)+2*om*cos(phi0)/rad*M.yt[:]
     M.coriolis_hor[:] =  0.
     return  

   def initial_conditions(self):
     """ setup all initial conditions
     """
     M=self.fortran.pyom_module   
     self.ini_b  = zeros((M.nx,M.ny,M.nz),'d')
     self.ini_bs = zeros((M.nx,M.ny,M.nz),'d')
     self.T_rest = zeros((M.nx,M.ny,M.nz),'d')
     for j in range(M.ny):
       for k in range(M.nz):
         B0 = M_0**2*M.yt[j]
         fxa = 0.5e-3*sin(M.xt*8.5/M.xt[-1]*pi)
         self.ini_b[:,j,k]  = (     M_0**2*M.yt[j] - N_0**2*M.zt[k])*M.maskt[:,j,k]
         self.ini_bs[:,j,k] = (fxa+ M_0**2*M.yt[j] - N_0**2*M.zt[k])*M.maskt[:,j,k]
     for k in range(3): M.b[:,:,:,k] = self.ini_bs
     for j in range(1,spg_width+1):
       self.T_rest[:,j,:]  =  t_rest/j*M.maskt[:,j,:]*2*M.dt
     for j in range(M.ny-1,M.ny-spg_width-1,-1):
       self.T_rest[:,j-1,:]  =  t_rest/(-1.*(j-M.ny)) *M.maskt[:,j-1,:]*2*M.dt
     return
  
   def restoring_zones(self):
     """  add here restoring zones
     """
     M=self.fortran.pyom_module   
     M.b[:,:,:,M.taup1-1]=M.b[:,:,:,M.taup1-1] + self.T_rest*(self.ini_b-M.b[:,:,:,M.taum1-1])
     return



   def make_plot(self):
     """ diagnose the model variables, could be replaced by other version
     """
     M=self.fortran.pyom_module         # fortran module with model variables
     k=M.nz*3/4
     i=int(M.nx/2)
     self.figure.clf()
     ax=self.figure.add_subplot(221)
     ax.contourf(M.xt[1:-1]/1e3,M.yt[1:-1]/1e3,M.b[1:-1,1:-1,k,M.tau-1].transpose()*1e2)
     ax.set_title('buoyancy')
     ax.set_ylabel('y [km]')
     
     ax.axis('tight')
     ax=self.figure.add_subplot(222)
     ax.contourf(M.yt[1:-1]/1e3,M.zt[1:-1],M.v[i,1:-1,1:-1,M.tau-1].transpose())
     ax.set_title('meridional velocity')
     ax.axis('tight')
     
     ax=self.figure.add_subplot(223)
     ax.contourf(M.yt[1:-1]/1e3,M.zt[1:-1],M.u[i,1:-1,1:-1,M.tau-1].transpose())
     ax.set_title('zonal velocity')
     ax.set_ylabel('z [km]')
     ax.set_xlabel('y [km]')
     ax.axis('tight')
     
     ax=self.figure.add_subplot(224)
     ax.contourf(M.yt[1:-1]/1e3,M.zt[1:-1],M.b[i,1:-1,1:-1,M.tau-1].transpose())
     ax.set_title('buoyancy')
     ax.set_xlabel('y [km]')
     ax.axis('tight')
     return

  
if __name__ == "__main__":
    model=jets1()
    model.run(snapint=3*86400.0,runlen=365*86400.)
    #jets1(snapint=72).mainloop()


