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


class THC1(pyOM):
   """ a THC model with salinity
   """
   def set_parameter(self):
     """set main parameter
     """
     RESOLVE=0.25; vRESOLVE=0.5
     M=self.fortran.pyom_module   
     M.nx=64 *RESOLVE
     M.ny=64 *RESOLVE
     M.nz=30 *vRESOLVE

     M.dx=80e3   /RESOLVE
     M.dt=2400.0 /RESOLVE  
     M.dz=50.0   /vRESOLVE

     M.a_h=5.0e4 /RESOLVE**2 # u_t = A u_xx,  1/T = A /L^2 
     M.a_v=1e-4
     M.k_v=1e-4
     M.k_h=0000.0 
     M.cdbot = 0e-5#*RESOLVE
     M.cdint = 2e-6
     M.eps2d_sor = 1e-9

     M.enable_interior_stress      = 1
     M.enable_hydrostatic          = 1
     M.enable_quicker_advection    = 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   
     rad  = 6370.0e3   # earth radius
     om   = pi/43082.0 # earth rotation freq. in 1/s
     phi0 = 0./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 initial conditions
     """
     M=self.fortran.pyom_module   
     self.salt_flux = self.get_salt_flux()

   def boundary_conditions(self):
     """ setup all forcing
         surface and bottom boundary conditions
         might be variable in time, called every time step
     """
     
     M=self.fortran.pyom_module   
     Ly = M.yt[-1] -M.yt[0]
     t_rest=30*86400.
     k=M.nz-2
     for j in range(M.ny): # sig = 0.8 S - 0.2T ,  T = (0.8S -sig)/0.2
       sig= M.b[:,j,k,M.taum1-1]*1024.0/9.81
       sss= M.tr[:,j,k,M.taum1-1,0]
       sst = (0.8* sss  -sig)/0.2
       t_star=15*(1-exp(-(M.yt[j]-M.yt[-2])**2/(Ly/8.)**2))
       M.surface_flux[:,j]=M.dz/t_rest*(t_star-sst) 
       M.surface_flux[:,j]=(0.8*self.salt_flux[j]-0.2*M.surface_flux[:,j])/1024.0*9.81
     return
  
   def tracer_sources(self):
     """ setup tracer source. Here a simple age tracer
     """
     M=self.fortran.pyom_module   
     k=M.nz-2
     for j in range(M.ny): 
       M.tr[:,j,k,M.taup1-1,0]=M.tr[:,j,k,M.taup1-1,0]+M.maskt[:,j,k]*2*M.dt*self.salt_flux[j]/M.dz
     return

   def get_salt_flux(self):
     M=self.fortran.pyom_module   
     Ly = M.yt[-1] -M.yt[0]
     salt_flux = M.yt*0.0
     A=0.5/86400.
     for j in range(M.ny): 
       salt_flux[j]=-A*exp(-(M.yt[j]-M.yt[-2])**2/(Ly/8.)**2 )
     fxa= sum(salt_flux )
     fxa=fxa/(M.ny -2.)
     return salt_flux-fxa
  
   def make_plot(self):
     """ diagnose the model variables, could be replaced by other version
     """
     if hasattr(self,'figure'):
         M=self.fortran.pyom_module         # fortran module with model variables
         i=int(M.nx/2)
         j=int(3*M.ny/4)
         k=int(M.nz-2)
         x=M.xt[1:-1] /1e3
         y=M.yt[1:-1] /1e3
         z=M.zt[1:-1] 
         
         self.figure.clf()
         T =  ( 0.8* M.tr[:,:,:,M.tau-1,0] -  M.b[:,:,:,M.tau-1]*1024.0/9.81 )/0.2
         ax=self.figure.add_subplot(221)
         a=T[1:-1,1:-1,1:-1]*M.maskt[1:-1,1:-1,1:-1]
         b=M.maskt[1:-1,1:-1,1:-1]
         a=sum(a,axis=0)/sum(b,axis=0)
         co=ax.contourf(y,z,a.transpose())
         self.figure.colorbar(co)
         ax.axis('tight')
         ax.set_title('Temp')
         
         ax=self.figure.add_subplot(222)
         a=M.tr[1:-1,1:-1,1:-1,M.tau-1,0]*M.maskt[1:-1,1:-1,1:-1]
         b=M.maskt[1:-1,1:-1,1:-1]
         a=sum(a,axis=0)/sum(b,axis=0)
         co=ax.contourf(y,z,a.transpose())
         self.figure.colorbar(co)
         ax.axis('tight')
         ax.set_title('Salt')
         
         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__":

   #model  = THC1()
   #model.run(runlen = 86400*365., snapint  = model.fortran.pyom_module.dt) 
   THC1(snapint  = 1 ).mainloop()
