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


class iso1(pyOM):
   """ 
   """
   def set_parameter(self):
     """set main parameter
     """
     M=self.fortran.pyom_module   
     M.nx    = 3
     M.nz    = 20
     M.ny    = 32
     M.dx    = 20e3
     M.dz    = 100.0 
     M.dt    = 1200.0 
     M.eps2d_sor = 1e-6
     M.a_v   = 10.e-4 
     M.k_v   = 1.0e-4 
     M.enable_hydrostatic          = 1
     M.enable_cyclic_x             = 1
     M.enable_quicker_advection    = 1
     M.enable_quicker_mom_advection= 1
     
     M.enable_vert_friction_trm    = 1

     M.enable_diag_tracer    = 1
     M.nt=1
     M.enable_isopycnal_diffusion = 1
     M.k_gm=2000.0
     
     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   
     f0 = 1e-4; beta=0e-11
     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   
     U_0 = 0.5
     N_0 = 0.004
     f=M.coriolis_t[M.ny/2]
     
     # velocity
     for k in range(M.nz):
      for j in range(M.ny):
        a=exp(-( M.yt[j]-M.yt[-1]/2.)**2/100e3**2 )
        M.u[:,j,k,0]=a*(U_0/2+a*U_0*M.zt[k]/((M.nz-2)*M.dz))*M.masku[:,j,k]
     M.u[:,:,:,1] = M.u[:,:,:,0]    
     M.u[:,:,:,2] = M.u[:,:,:,0]


     # buoyancy
     #fac = 2*0.2/1024*9.81
     #for k in range(M.nz):
     #  for j in range(M.ny):
     #    y = (M.yt[j]-M.yt[1])/ (M.yt[-2]-M.yt[1])
     #    B0=fac*0.5*(tanh((y-0.5)/0.1)+1.0)
     #    M.b[:,j,k,0]  = (B0-N_0**2*M.zt[k])*M.maskt[:,j,k]
        
     for k in range(M.nz):
       for j in range(M.ny-1):
        a1=  exp(-( M.yt[j]-M.yt[-1]/2.)**2/100e3**2 )
        a2=  exp(-( M.yt[j+1]-M.yt[-1]/2.)**2/100e3**2 )
        a=(a1+a2)/2.0
        fxa = (M.coriolis_t[j]+M.coriolis_t[j+1])/2.0
        M.b[:,j+1,k,0]=M.b[:,j,k,0]+M.dx*M.maskt[:,j,k]*a*U_0/((M.nz-2)*M.dz)*fxa
     for k in range(M.nz):
        M.b[:,:,k,0]=(M.b[:,:,k,0]-N_0**2*M.zt[k])*M.maskt[:,:,k]
     M.b[:,:,:,1] = M.b[:,:,:,0]    
     M.b[:,:,:,2] = M.b[:,:,:,0]

     for k in range(M.nz):
       for j in range(M.ny):
         M.tr[:,j,k,0,0] = exp( -( M.zt[k]-M.zt[0]/2.)**2/100.0**2 - ( M.yt[j]-M.yt[-1]/2.)**2/100e3**2 )
     M.tr[:,:,:,1,0] = M.tr[:,:,:,0,0]    
     M.tr[:,:,:,2,0] = M.tr[:,:,:,0,0]
     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()
       
       ax=self.figure.add_subplot(221)
       a=M.b[1,1:-1,1:-1,M.tau-1] 
       co=ax.contourf(y,z,a.transpose())
       self.figure.colorbar(co)
       a=M.u[1,1:-1,1:-1,M.tau-1] 
       ax.contour(y,z,a.transpose(),10,colors='k')
       ax.set_title('buoyancy'); ax.set_xlabel('y [km]'); ax.set_ylabel('z [m]')
       ax.axis('tight')
       
       ax=self.figure.add_subplot(222)
       a=M.tr[1,1:-1,1:-1,M.tau-1,0] 
       co=ax.contourf(y,z,a.transpose())
       a=M.b[1,1:-1,1:-1,M.tau-1] 
       ax.contour(y,z,a.transpose(),10,colors='k')
       ax.axis('tight')


       ax=self.figure.add_subplot(223)
       a=M.v[1:-1,1:-1,1:-1,M.tau-1]*M.maskv[1:-1,1:-1,1:-1]
       b=M.maskv[1:-1,1:-1,1:-1]
       a=sum(a,axis=0)*M.dx
       a=cumsum(a,axis=1)*M.dz  /1e6
       co=ax.contourf(y,z,a.transpose())
       self.figure.colorbar(co)
       ax.axis('tight')
       ax.set_title('MOC')

     return


if __name__ == "__main__": iso1(snapint  = 24).mainloop()
