
 module fcontrol_module
!=======================================================================
!     simple control of integration length and configuration
!=======================================================================
      implicit none
      integer :: itt,enditt
      real*8  :: runlen,snapint
      logical :: enable_diag_zonalave  = .false.
      logical :: enable_diag_timeave   = .false.
      logical :: enable_diag_particle  = .false.
      logical :: enable_bolus_velocity = .false.
      logical :: enable_pv_mixing      = .false.
      logical :: enable_diag_barbi     = .false.
      logical :: enable_diag_quasi_stokes  = .false.
      character*80 :: snap_file = 'pyOM.cdf'
 end module fcontrol_module

 program main
!=======================================================================
!      Top level driver for fortran
!=======================================================================
      use pyOM_module   
      use fcontrol_module
      use timing_module
      implicit none
      integer :: ierr,otaum1,n,io
!---------------------------------------------------------------------------------
!       Initialize MPI setting
!---------------------------------------------------------------------------------
      print '(a,f5.3)' ,'here is pyOM version ',version
      call mpi_init(ierr)
      call my_mpi_init(my_comm,ierr)
!---------------------------------------------------------------------------------
!      Initialize model
!---------------------------------------------------------------------------------
      call tic('setup')
      itt=0
      call setup
      if (enable_pv_mixing)      call init_pv_mixing
      if (enable_bolus_velocity) call init_bolus
      if (enable_diag_particle)  call init_particle
      if (enable_diag_zonalave)  call init_zonal_averages
      if (enable_diag_timeave)   call init_time_averages
      if (enable_diag_barbi)     call init_diag_barbi
      if (enable_diag_quasi_stokes)  call init_quasi_stokes
!---------------------------------------------------------------------------------
!---------------------------------------------------------------------------------
      enditt = itt+runlen/dt
      if (my_pe==0) then
        print'(a,e8.2,a)',' integration for ',runlen,' s'
        print'(a,i10,a,i10)',' from time step ',itt,' to ',enditt
        print'(a,e8.2,a)',' snapshot intervall is ',snapint,' s'
        print'(a,i8,a)',' this is any ',int(snapint/dt),' time steps'
      endif
      call toc('setup')

      call tic('main loop')
!---------------------------------------------------------------------------------
!      Begin main model loop
!---------------------------------------------------------------------------------
 10   continue
      call boundary_conditions
      if (enable_hydrostatic)       call vertical_velocity(ierr)

      call tic('buoyancy')
      call integrate_buoyancy(ierr)

      call restoring_zones
      if (enable_back_state)        call background_state_buoyancy(ierr)
      if (enable_hydrostatic)       call convection(ierr) 
      if (enable_diag_tracer)  then
                 call integrate_tracer(ierr)
                 if (enable_isopycnal_diffusion) call isopycnal_diffusion(ierr)
                 call tracer_sources
      endif
      call toc('buoyancy')

      call tic('momentum')
      call momentum_tendency(ierr)
      if (enable_pv_mixing)         call pv_mixing_main(ierr)
      if (enable_bolus_velocity)    call bolus_velocity
      if (enable_vert_friction_trm) call vert_friction_trm(ierr)
      call momentum_restoring_zones()       
      if (enable_back_state)        call background_state_momentum(ierr)
      call toc('momentum')

      call tic('pressure')
      call solve_pressure(ierr)
      call toc('pressure')

      call integrate(ierr)

      call tic('diagnose')
      call diagnose
      if (enable_diag_tracer)   call diag_tracer
      if (enable_diag_particle) call integrate_particle
      if (enable_diag_zonalave) call zonal_averages
      if (enable_diag_timeave)  call time_averages
      if (enable_diag_barbi)    call diag_barbi
      if (enable_diag_quasi_stokes)  call quasi_stokes
      if (enable_pv_mixing)     call pv_mixing_diag
      call toc('diagnose')

      otaum1=taum1; taum1= tau; tau  = taup1; taup1= otaum1
      if (itt < enditt) then
         itt=itt+1
         goto 10
      endif
      call toc('main loop')
!---------------------------------------------------------------------------------
!       End main model loop
!---------------------------------------------------------------------------------
      call write_restart(itt,'restart.dta')
      if (enable_diag_tracer)   call tracer_write_restart('tracer_restart.dta')
      if (enable_diag_particle) call particle_write_restart(ierr)
      if (enable_diag_quasi_stokes)  then 
        call quasi_stokes_write_restart()
        call quasi_stokes_write_cdf()
      endif
      if (enable_diag_timeave) then
       call write_time_averages_restart(ierr)
       call time_averages_calc
       call time_averages_write
      endif

      if (my_pe==0) then
       call get_free_iounit(n,ierr)
       open(n,file='ritt',form='formatted',status='unknown')
       write(n,*) itt
       close(n)
      endif
!--------------------------------------------------------------
!     show timing results here
!--------------------------------------------------------------
      do n = 0,n_pes
       call fortran_barrier
       if (my_pe == n) then
        print*,''
        print*,'Timing summary for PE #',my_pe 
        print*,' costs for measuring      = ',timing_secs('tictoc'),' s'
        print*,' setup time summary       = ',timing_secs('setup'),' s'
        print*,' main loop time summary   = ',timing_secs('main loop') ,' s'
        print*,' buoyancy time summary    = ',timing_secs('buoyancy') ,' s'
        print*,' momentum time summary    = ',timing_secs('momentum') ,' s'
        print*,' pressure time summary    = ',timing_secs('pressure') ,' s'
        print*,' diagnostics time summary = ',timing_secs('diagnose') ,' s'
       endif
      enddo

   call mpi_finalize(ierr)
 end program main



 subroutine setup
!=======================================================================
!      setup everything
!=======================================================================
      use pyOM_module   
      use fcontrol_module
      implicit none
      integer :: i,j,k,n,ierr
      logical :: file_exists = .false.
! ----------------------------------
!       Set main parameter
! ----------------------------------
      call set_parameter
! ----------------------------------
!       Some more or less useful output
! ----------------------------------
      if (my_pe==0) then
       print*,''
       print'(a,i4,a,i4,a,i4)',' grid size    : nx=',nx,' ny=',ny,' nz=',nz
       print'(a,e8.2,a,e8.2,a)',' grid spacing : Delta x =',dx,'m  Delta z =',dz,'m'
       print'(a,e8.2,a,e8.2,a,e8.2,a)',' domain size  : ',nx*dx,'m X ',ny*dx,'m X ',nz*dz,'m'
       print'(a,e8.2,a)',' time step    : ',dt,'s'
       print'(a,e8.2,a)',' lateral  diffusivity : K_h=',K_h,' m^2/s'
       print'(a,e8.2,a)',' vertical diffusivity : K_v=',K_v,' m^2/s'
       print'(a,e8.2,a)',' lateral  viscosity   : A_h=',A_h,' m^2/s'
       print'(a,e8.2,a)',' vertical viscosity   : A_v=',A_v,' m^2/s'
      endif

       c2dt=dt
! ----------------------------------
!      time step splitting for explicit free surface
! ----------------------------------
       dtex=dx/sqrt(9.81*(nz-2)*dz ) /4.
       n=nint(dt/dtex)
       dtex=dt/n
       if (my_pe==0.and.enable_expl_free_surf) then
        print*,' free surface time step : ',dtex,' s'
       endif
! ----------------------------------
!      domain decomposition for each PE
! ----------------------------------
      if (n_pes>1) then
       n_pes_j = n_pes
       j_blk = (ny-1)/n_pes_j + 1    ! j-extent of each block
       my_blk_j = my_pe + 1 ! number of PE in j-dir.
       js_pe = (my_blk_j-1)*j_blk + 1
       je_pe = min(my_blk_j*j_blk,ny)
! ----------------------------------
!      check for incorrect domain decomposition
! ----------------------------------
       if (my_blk_j==n_pes_j .and. js_pe>=je_pe-2) then
        print*,' ERROR: on PE: ', my_pe
        print*,' domain decompositon impossible in j-direction'
        print*,' choose other number of PEs in j-direction'
        call halt_stop(' ... ')
       endif
      else
       n_pes_j = n_pes
       j_blk = ny
       my_blk_j = 1 
       js_pe = 1
       je_pe = ny
      endif
! ----------------------------------
!      print out the PE decomposition, let all PEs talk
! ----------------------------------
      if (my_pe==0) print*,' Domain decomposition:'
      do n=0,n_pes-1
       if (my_pe==n) then
         print*,''
         print*,' sub domain for PE #',n
         print*,' my_blk_j=',my_blk_j
         print*,' js_pe=',js_pe,' je_pe=',je_pe
         print*,''
       endif
      enddo
! ----------------------------------
!      allocate work space
! ----------------------------------
      call allocate_work_space
! ----------------------------------
!      setup grid
! ----------------------------------
      xt(1)=0.0
      xu(1)=xt(1)+dx/2.0
      do i=2,nx
        xt(i)=xt(i-1)+dx
        xu(i)=xu(i-1)+dx
      enddo
      yt(1)=0.0
      yu(1)=yt(1)+dx/2.0
      do i=2,ny
        yt(i)=yt(i-1)+dx
        yu(i)=yu(i-1)+dx
      enddo
      zw(1)=0
      zt(1)=zw(1)-dz/2.0
      do k=2,nz
        zt(k)=zt(k-1)+dz
        zw(k)=zw(k-1)+dz
      enddo
      zt=zt-dz*(nz-2)
      zw=zw-dz*(nz-2)
! ----------------------------------
!      set coriolis parameter
! ----------------------------------
      call set_coriolis()
! ----------------------------------
!      Land mask
! ----------------------------------
      maskT=0.0
      maskT(2:nx-1,2:ny-1,2:nz-1)=1.0
      call topography
      if (enable_cyclic_x) maskT(1,:,:)=0.0
      if (enable_cyclic_x) maskT(nx,:,:)=0.0
      if (enable_cyclic_y) maskT(:,1,:)=0.0
      if (enable_cyclic_y) maskT(:,ny,:)=0.0
      call setcyclic3D(nx,ny,nz,maskT)
      maskU=maskT
      do i=1,nx-1
       maskU(i,:,:)=min(maskT(i,:,:),maskT(i+1,:,:))
      enddo
      call setcyclic3D(nx,ny,nz,maskU)
      maskV=maskT
      do j=1,ny-1
       maskV(:,j,:)=min(maskT(:,j,:),maskT(:,j+1,:))
      enddo
      call setcyclic3D(nx,ny,nz,maskV)
      maskW=maskT
      do k=1,nz-1
       maskW(:,:,k)=min(maskT(:,:,k),maskT(:,:,k+1))
      enddo
      call setcyclic3D(nx,ny,nz,maskW)
! ----------------------------------
!     depth 
! ----------------------------------
      ht=0.0;hu=0.0;hv=0.0 
      do k=1,nz
       ht=ht+maskT(:,:,k)*dz
       hu=hu+masku(:,:,k)*dz
       hv=hv+maskv(:,:,k)*dz
      enddo
      do k=nz,1,-1
       where (maskU(:,:,k)>0) k_bottom_u=k
       where (maskV(:,:,k)>0) k_bottom_v=k
      enddo
! ----------------------------------
!     coefficients for poisson solver
! ----------------------------------
      call make_coef2d(ierr)
      if (.not.enable_hydrostatic) call make_coef3d(ierr)
! ----------------------------------
!   initial conditions 
! ----------------------------------
      if (my_pe==0) print*,' using initial conditions '
      call initial_conditions
      inquire( FILE='restart.dta', EXIST=file_exists )
      if (file_exists) then 
        call read_restart(itt,'restart.dta')
      else
       call exchg_prog3Dfield_complete(nx,ny,nz,b)
       call exchg_prog3Dfield_complete(nx,ny,nz,u)
       call exchg_prog3Dfield_complete(nx,ny,nz,v)
       call exchg_prog3Dfield_complete(nx,ny,nz,w)
       call exchg_prog3Dfield_complete(nx,ny,nz,p_full)
       call exchg_prog2Dfield_complete(nx,ny,eta)
       call exchg_3Dfield_complete(nx,ny,nz,psi)
       if (enable_expl_free_surf) then
         call exchg_prog2Dfield_complete(nx,ny,bu)
         call exchg_prog2Dfield_complete(nx,ny,bv)
       endif
      endif
! ----------------------------------
!   initialize diagnostic       
! ----------------------------------
      if (my_pe==0) call init_snap_cdf()
! ----------------------------------
!  passive tracer stuff
! ----------------------------------
      if (my_pe==0 .and. enable_diag_tracer) call init_tracer_diag()
      inquire(FILE='tracer_restart.dta', EXIST=file_exists )
      if (file_exists) call tracer_read_restart('tracer_restart.dta')
  end subroutine setup



 subroutine allocate_work_space
!=======================================================================
!      Allocate memory for all arrays in main module
!=======================================================================
      use pyOM_module   
      implicit none
      allocate(zt(nz),zw(nz),xt(nx ),xu(nx ),yt(ny ),yu(ny ) )
      taum1=1; tau=2; taup1=3
      allocate( u(nx ,ny ,nz,3) ); u=0.
      allocate( v(nx ,ny ,nz,3) ); v=0.
      allocate( w(nx ,ny ,nz,3) ); w=0.
      allocate( b(nx ,ny ,nz,3)   ); b=0.
      allocate( K_b(nx ,ny ,nz)  );    K_b=0.
      allocate( p_hydro(nx ,ny ,nz) ); p_hydro=0.
      allocate( p_full(nx ,ny ,nz,3)  ); p_full=0.
      allocate( p_surf(nx ,ny )     ); p_surf=0.
      allocate( psi(nx ,ny ,nz)     ); psi=0.
      allocate( eta(nx ,ny ,3)    ); eta=0.
      allocate( maskT(nx ,ny ,nz)   ); maskT=0.
      allocate( maskU(nx ,ny ,nz)   ); maskU=0.
      allocate( maskV(nx ,ny ,nz)   ); maskV=0.
      allocate( maskW(nx ,ny ,nz)   ); maskW=0.
      allocate( fu(nx ,ny ,nz)      ); fu=0.
      allocate( fv(nx ,ny ,nz)      ); fv=0.
      allocate( fw(nx ,ny ,nz)      ); fw=0.
      allocate( surface_flux(nx ,ny )  ); surface_flux=0.
      allocate( bottom_flux(nx ,ny )  ); bottom_flux=0.
      allocate( surface_taux(nx ,ny ) ); surface_taux=0.
      allocate( surface_tauy(nx ,ny ) ); surface_tauy=0.
      allocate( bottom_taux(nx ,ny ) ); bottom_taux=0.
      allocate( bottom_tauy(nx ,ny ) ); bottom_tauy=0.
      allocate(coriolis_t(ny ), coriolis_hor(ny ) )
      coriolis_t=0.;coriolis_hor=0.
      allocate( hu(nx ,ny ),hv(nx ,ny ),ht(nx ,ny ))
      hu=0.;hv=0.;ht=0.
      allocate( cf2D(nx ,ny ,3,3) )
      if (.not.enable_hydrostatic) then
        allocate( cf3D(nx ,ny ,nz,3,3,3))
      endif
      if (enable_expl_free_surf) then
       allocate( bu(nx ,ny ,3)    ); bu=0.
       allocate( bv(nx ,ny ,3)    ); bv=0.
      endif
      allocate( k_bottom_u(nx,ny) ); k_bottom_u=0
      allocate( k_bottom_v(nx,ny) ); k_bottom_v=0
      if (enable_vert_friction_trm) then
        allocate( a_trm(nx,ny,nz) ); a_trm=0.
      endif
      if  (enable_back_state) then
        allocate( back(nx,ny,nz,3) ); back=0.
        allocate( u0(nx,ny,nz) ); u0=0.
      endif
      if (enable_diag_tracer) then
         allocate( tr(nx,ny,nz,3,nt) ); tr=0.0
      endif
 end subroutine allocate_work_space

