#include "options.inc"


#ifdef enable_smagorinsky_friction

c=======================================================================
c      Smagornisky mixing
c      linked in the code in momentum.F only
c=======================================================================

      module smagorinsky_module
      real, parameter :: smag = ( 4./3.14 )**2
      real,allocatable :: tension_Z(:,:,:), strain_Z(:,:,:)
      real,allocatable :: tension_T(:,:,:), strain_T(:,:,:)
      real,allocatable :: diff_T(:,:,:),    diff_Z(:,:,:)
      end module smagorinsky_module



      subroutine smagorinsky
c=======================================================================
c=======================================================================
      use cpflame_module
      use smagorinsky_module
      implicit none
      integer :: i,j,k,js,je
      logical, save :: first= .true.
#include "netcdf.inc"
      integer :: iret,ncid,npe,itdimid,ilen,itimeid
      type(time_type) :: time
      integer :: corner(4), edges(4),varid
      real :: a(imt,js_pe:je_pe,km),fxa

      if (first) call smagorinsky_init

      call u_x_minus_v_y(tension_T) ! at tracer point
      call border_exchg3D(tension_T,1)
      call setcyclic3D(tension_T)
      call u_y_plus_v_x(strain_Z)   ! at zeta point
      call border_exchg3D(strain_Z,1)
      call setcyclic3D(strain_Z)

      js=max(2,js_pe); je = min(je_pe,jmt-1)
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         strain_T(i,j,k)=0.25*(strain_Z(i,j  ,k)+strain_Z(i-1,j  ,k)
     &                        +strain_Z(i,j-1,k)+strain_Z(i-1,j-1,k))
         tension_Z(i,j,k)=0.25*(tension_T(i,j,k)+tension_T(i+1,j,k)
     &                       +tension_T(i,j+1,k)+tension_T(i+1,j+1,k))
         fxa = strain_T(i,j,k)**2+tension_T(i,j,k)**2
         diff_T(i,j,k)=smag*dx**2*sqrt(fxa)
         fxa = strain_Z(i,j,k)**2+tension_Z(i,j,k)**2
         diff_Z(i,j,k)=smag*dx**2*sqrt(fxa)
        enddo
       enddo
      enddo
      call border_exchg3D(diff_Z,1)
      call setcyclic3D(diff_Z)
      call border_exchg3D(diff_T,1)
      call setcyclic3D(diff_T)


      if (snapshot_time_step.or.initial_time==current_time) then

       do npe=0,n_pes-1
        if (my_pe==npe) then
         iret=nf_open('smagorinsky.cdf',NF_WRITE,ncid)
         iret=nf_set_fill(ncid, NF_NOFILL, iret)
         iret=nf_inq_dimid(ncid,'Time',itdimid)
         iret=nf_inq_dimlen(ncid, itdimid,ilen)
         iret=nf_inq_varid(ncid,'Time',itimeid)
         if (my_pe==0) then
          ilen=ilen+1
          time = current_time-initial_time
          fxa = time%days + time%seconds/86400.
          iret= nf_put_vara_double(ncid,itimeid,ilen,1,fxa)
         endif

         Corner = (/1,js_pe,1,ilen/); 
         edges  = (/imt,je_pe-js_pe+1,km,1/)
         iret=nf_inq_varid(ncid,'Strain',varid)
         a=strain_T(:,js_pe:je_pe,:)
         where( maskT(:,js_pe:je_pe,:) == 0.) a = spval
         iret= nf_put_vara_double(ncid,varid,corner,edges,a)

         iret=nf_inq_varid(ncid,'Tension',varid)
         a=tension_T(:,js_pe:je_pe,:)
         where( maskT(:,js_pe:je_pe,:) == 0.) a = spval
         iret= nf_put_vara_double(ncid,varid,corner,edges,a)

         iret=nf_inq_varid(ncid,'Diff',varid)
         a=Diff_T(:,js_pe:je_pe,:)
         where( maskT(:,js_pe:je_pe,:) == 0.) a = spval
         iret= nf_put_vara_double(ncid,varid,corner,edges,a)

         call ncclos (ncid, iret)
        endif
        call barrier
       enddo
      endif
      first=.false.
      end subroutine smagorinsky



      subroutine smagorinsky_init
c=======================================================================
c=======================================================================
      use cpflame_module
      use smagorinsky_module
      implicit none
#include "netcdf.inc"
      integer :: iret,ncid,dims(4),vid
      integer :: lon_tdim,lat_tdim,z_tdim,z_udim,itimedim
      character (len=80) :: name, unit

      allocate(tension_T(imt,jmt,km),strain_T(imt,jmt,km))
      tension_T=0.; strain_T=0.
      allocate(tension_Z(imt,jmt,km),strain_Z(imt,jmt,km))
      tension_Z=0.; strain_Z=0.
      allocate(Diff_Z(imt,jmt,km),Diff_T(imt,jmt,km))
      Diff_Z=0.; Diff_T=0.

      if (my_pe==0) then
       call def_grid_cdf('smagorinsky.cdf')
       iret=nf_open('smagorinsky.cdf',NF_WRITE,ncid)
       iret=nf_set_fill(ncid, NF_NOFILL, iret)
       call ncredf(ncid, iret)
       iret=nf_inq_dimid(ncid,'xt',lon_tdim)
       iret=nf_inq_dimid(ncid,'yt',lat_tdim)
       iret=nf_inq_dimid(ncid,'zt',z_tdim)
       iret=nf_inq_dimid(ncid,'zu',z_udim)
       iret=nf_inq_dimid(ncid,'Time',itimedim)
       dims = (/Lon_tdim,lat_tdim, z_tdim, iTimedim/)
       vid = ncvdef (ncid,'Strain', NCFLOAT,4,dims,iret)
       name = 'Strain (u_y+v_x)'; unit = '1/s'
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)
       vid = ncvdef (ncid,'Tension', NCFLOAT,4,dims,iret)
       name = 'Tension (u_x-v_y)'; unit = '1/s'
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)
       vid = ncvdef (ncid,'Diff', NCFLOAT,4,dims,iret)
       name = 'Viscosity'; unit = 'm^2/s'
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)
       call ncclos (ncid, iret)
      endif
      end subroutine smagorinsky_init



      subroutine smagorinsky_fric_u(diff_fe,diff_fn)
c=======================================================================
c       horizontal harmonic friction for u  
c       diff. fluxes are stored in diff_fe and diff_fn
c       account for no slip boundary condition if requested
c=======================================================================
      use cpflame_module
      use smagorinsky_module
      implicit none
      integer :: i,j,k,js,je
      real :: diff_fn(imt,jmt,km), diff_fe(imt,jmt,km),diff

      js=max(2,js_pe); je = min(je_pe,jmt-1)

      do k=2,km-1
       do j=js,je
        do i=1,imt-1
         diff = diff_T(i+1,j,k)
         diff_fe(i,j,k)=diff*(u(i+1,j,k,1,taum1)-u(i,j,k,1,taum1))/dx
        enddo
       enddo
      enddo
      call setcyclic3D(diff_fe)
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         diff = diff_Z(i,j,k)
         diff_fn(i,j,k)=diff*(u(i,j+1,k,1,taum1)-u(i,j,k,1,taum1))/dx
     &                   *maskU(i,j+1,k)*maskU(i,j,k)
        enddo
       enddo
      enddo
      if (enable_noslip) then
       do j=js-1,je
        diff_fn(:,j,:)=diff_fn(:,j,:)-2*A_h*u(:,j,:,1,taum1)/dx
     &                   *(1-maskU(:,j+1,:))*maskU(:,j,:)
        diff_fn(:,j,:)=diff_fn(:,j,:)+2*A_h*u(:,j+1,:,1,taum1)/dx
     &                   *(1-maskU(:,j,:))*maskU(:,j+1,:)
       enddo
      endif
      call border_exchg3D(diff_fn,1)
      call setcyclic3D(diff_fn)
      end subroutine smagorinsky_fric_u





      subroutine smagorinsky_fric_v(diff_fe,diff_fn)
c=======================================================================
c       horizontal harmonic friction for v  
c       diff. fluxes are stored in diff_fe and diff_fn
c       account for no slip boundary condition if requested
c=======================================================================
      use cpflame_module
      use smagorinsky_module
      implicit none
      integer :: i,j,k,js,je
      real :: diff_fn(imt,jmt,km), diff_fe(imt,jmt,km),diff

      js=max(2,js_pe); je = min(je_pe,jmt-1)

      do k=2,km-1
       do j=js,je
        do i=1,imt-1
         diff = diff_Z(i,j,k)
         diff_fe(i,j,k)=diff*(u(i+1,j,k,2,taum1)-u(i,j,k,2,taum1))/dx
     &                       *maskV(i+1,j,k)*maskV(i,j,k)
        enddo
       enddo
      enddo
      if (enable_noslip) then
       do i=1,imt-1
         diff_fe(i,:,:)=diff_fe(i,:,:)-2*A_h*u(i,:,:,2,taum1)/dx
     &                   *(1-maskV(i+1,:,:))*maskV(i,:,:)
         diff_fe(i,:,:)=diff_fe(i,:,:)+2*A_h*u(i+1,:,:,2,taum1)/dx
     &                   *(1-maskV(i,:,:))*maskV(i+1,:,:)
       enddo
      endif
      call setcyclic3D(diff_fe)
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         diff = diff_T(i,j+1,k)
         diff_fn(i,j,k)=diff*(u(i,j+1,k,2,taum1)-u(i,j,k,2,taum1) )/dx
        enddo
       enddo
      enddo
      call border_exchg3D(diff_fn,1)
      call setcyclic3D(diff_fn)
      end subroutine smagorinsky_fric_v




      subroutine u_y_plus_v_x(tension)
c=======================================================================
c       calculate d/dy u plus d/dx v on zeta points
c=======================================================================
      use cpflame_module
      implicit none
      integer :: i,j,k,js,je
      real :: tension(imt,jmt,km)
      js=max(2,js_pe); je = min(je_pe,jmt-1)
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         tension(i,j,k)=(u(i,j+1,k,1,tau)-u(i,j,k,1,tau))/dx
     &                   *maskU(i,j,k)*maskU(i,j+1,k)
     &                 +(u(i+1,j,k,2,tau)-u(i,j,k,2,tau))/dx
     &                   *maskV(i,j,k)*maskV(i+1,j,k)
        enddo
       enddo
      enddo
      end subroutine u_y_plus_v_x



      subroutine v_x_minus_u_y(omega_Z)
c=======================================================================
c       calculate relative vorticity on zeta points
c=======================================================================
      use cpflame_module
      implicit none
      integer :: i,j,k,js,je
      real :: omega_Z(imt,jmt,km)
      js=max(2,js_pe); je = min(je_pe,jmt-1)
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         omega_Z(i,j,k)=-(u(i,j+1,k,1,tau)-u(i,j,k,1,tau))/dx
     &                   *maskU(i,j,k)*maskU(i,j+1,k)
     &                 +(u(i+1,j,k,2,tau)-u(i,j,k,2,tau))/dx
     &                   *maskV(i,j,k)*maskV(i+1,j,k)
        enddo
       enddo
      enddo
      end subroutine v_x_minus_u_y



      subroutine u_x_minus_v_y(strain)
c=======================================================================
c       calculate d/dx u minus d/dy v on tracer points
c=======================================================================
      use cpflame_module
      implicit none
      integer :: i,j,k,js,je
      real :: strain(imt,jmt,km)
      js=max(2,js_pe); je = min(je_pe,jmt-1)
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         strain(i,j,k)=(u(i,j,k,1,tau)-u(i-1,j,k,1,tau))/dx
     &                 *maskU(i,j,k)*maskU(i-1,j,k)
     &                -(u(i,j,k,2,tau)-u(i,j-1,k,2,tau))/dx
     &                 *maskV(i,j,k)*maskV(i,j-1,k)
        enddo
       enddo
      enddo
      end subroutine u_x_minus_v_y

#else
      subroutine viscosity_dummy
      end
#endif
