MODULE caldyn_gcm_mod
  USE icosa
  USE transfert_mod
  USE caldyn_kernels_hevi_mod
  USE caldyn_kernels_base_mod
  USE caldyn_kernels_mod
  IMPLICIT NONE
  PRIVATE

  PUBLIC init_caldyn, caldyn_BC, caldyn
  
CONTAINS
  
  SUBROUTINE init_caldyn
    USE icosa
    USE observable_mod
    USE mpipara
    USE omp_para
    IMPLICIT NONE
    CHARACTER(len=255) :: def
    INTEGER            :: ind
    REAL(rstd),POINTER :: planetvel(:)

    hydrostatic=.TRUE.
    CALL getin("hydrostatic",hydrostatic)
  
    dysl=.NOT.hydrostatic ! dysl defaults to .FALSE. if hydrostatic, .TRUE. if NH                                              
    CALL getin("dysl",dysl)
    dysl_geopot=dysl
    CALL getin("dysl_geopot",dysl_geopot)
    dysl_caldyn_fast=dysl
    CALL getin("dysl_caldyn_fast",dysl_caldyn_fast)
    dysl_slow_hydro=dysl
    CALL getin("dysl_slow_hydro",dysl_slow_hydro)
    dysl_pvort_only=dysl
    CALL getin("dysl_pvort_only",dysl_pvort_only)
    dysl_caldyn_coriolis=dysl
    CALL getin("dysl_caldyn_coriolis",dysl_caldyn_coriolis)

    def='energy'
    CALL getin('caldyn_conserv',def)
    SELECT CASE(TRIM(def))
    CASE('energy')
       caldyn_conserv=energy
    CASE('enstrophy')
       caldyn_conserv=enstrophy
    CASE DEFAULT
       IF (is_mpi_root) PRINT *,'Bad selector for variable caldyn_conserv : <', &
            TRIM(def),'> options are <energy>, <enstrophy>'
       STOP
    END SELECT
    IF (is_master) PRINT *, 'caldyn_conserv=',def

    nqdyn=1 ! default value
    physics_thermo = thermo_none

    def='theta'
    CALL getin('thermo',def)
    SELECT CASE(TRIM(def))
    CASE('boussinesq')
       boussinesq=.TRUE.
       caldyn_thermo=thermo_boussinesq
       IF(.NOT. hydrostatic) THEN
          PRINT *, 'thermo=boussinesq and hydrostatic=.FALSE. : Non-hydrostatic boussinesq equations are not supported'
          STOP
       END IF
    CASE('theta')
       caldyn_thermo=thermo_theta
       physics_thermo=thermo_dry
    CASE('entropy')
       caldyn_thermo=thermo_entropy
       physics_thermo=thermo_dry
    CASE('theta_fake_moist')
       caldyn_thermo=thermo_theta
       physics_thermo=thermo_fake_moist
    CASE('entropy_fake_moist')
       caldyn_thermo=thermo_entropy
       physics_thermo=thermo_fake_moist
    CASE('moist')
       caldyn_thermo=thermo_moist_debug
       physics_thermo=thermo_moist
       nqdyn = 2
    CASE DEFAULT
       IF (is_mpi_root) PRINT *,'Bad selector for variable caldyn_thermo : <', &
            TRIM(def),'> options are <theta>, <entropy>'
       STOP
    END SELECT

    IF(is_master) THEN
       SELECT CASE(caldyn_thermo)
       CASE(thermo_theta)
          PRINT *, 'caldyn_thermo = thermo_theta'
       CASE(thermo_entropy)
          PRINT *, 'caldyn_thermo = thermo_entropy'
       CASE(thermo_moist_debug)
          PRINT *, 'caldyn_thermo = thermo_moist_debug'
       CASE DEFAULT
          STOP
       END SELECT

       SELECT CASE(physics_thermo)
       CASE(thermo_dry)
          PRINT *, 'physics_thermo = thermo_dry'
       CASE(thermo_fake_moist)
          PRINT *, 'physics_thermo = thermo_fake_moist'
       CASE(thermo_moist)
          PRINT *, 'physics_thermo = thermo_moist'
       END SELECT

       PRINT *, 'nqdyn =', nqdyn
    END IF

    CALL allocate_caldyn

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       planetvel=f_planetvel(ind)
       CALL compute_planetvel(planetvel)
    END DO

  END SUBROUTINE init_caldyn

  SUBROUTINE allocate_caldyn
  USE icosa
  IMPLICIT NONE

    CALL allocate_field(f_out_u,field_u,type_real,llm) 
    CALL allocate_field(f_qu,field_u,type_real,llm) 
    CALL allocate_field(f_qv,field_z,type_real,llm) 
    CALL allocate_field(f_pk,    field_t,type_real,llm,  name='pk')
    CALL allocate_field(f_wwuu,  field_u,type_real,llm+1,name='wwuu')    
    CALL allocate_field(f_planetvel,  field_u,type_real, name='planetvel') ! planetary velocity at r=a
    IF(.NOT.hydrostatic) THEN
       CALL allocate_field(f_Fel,      field_u,type_real,llm+1,name='F_el')
       CALL allocate_field(f_gradPhi2, field_t,type_real,llm+1,name='gradPhi2')
       CALL allocate_field(f_wil,      field_t,type_real,llm+1,name='w_il')
       CALL allocate_field(f_Wetadot,  field_t,type_real,llm,name='W_etadot')
    END IF
  END SUBROUTINE allocate_caldyn

  SUBROUTINE caldyn_BC(f_phis, f_geopot, f_wflux)
    USE icosa
    USE mpipara
    USE omp_para
    TYPE(t_field),POINTER :: f_phis(:)
    TYPE(t_field),POINTER :: f_geopot(:)
    TYPE(t_field),POINTER :: f_wflux(:)
    REAL(rstd),POINTER  :: phis(:)
    REAL(rstd),POINTER  :: wflux(:,:)
    REAL(rstd),POINTER  :: geopot(:,:)
    REAL(rstd),POINTER  :: wwuu(:,:)

    INTEGER :: ind,i,j,ij,l

    IF (is_omp_first_level) THEN
       DO ind=1,ndomain
          IF (.NOT. assigned_domain(ind)) CYCLE
          CALL swap_dimensions(ind)
          CALL swap_geometry(ind)
          geopot=f_geopot(ind)
          phis=f_phis(ind)
          wflux=f_wflux(ind)
          wwuu=f_wwuu(ind)
          
          DO ij=ij_begin_ext,ij_end_ext
              ! lower BCs : geopot=phis, wflux=0, wwuu=0
              geopot(ij,1) = phis(ij)
              wflux(ij,1) = 0.
              wwuu(ij+u_right,1)=0   
              wwuu(ij+u_lup,1)=0   
              wwuu(ij+u_ldown,1)=0
              ! top BCs : wflux=0, wwuu=0
              wflux(ij,llm+1)  = 0.
              wwuu(ij+u_right,llm+1)=0   
              wwuu(ij+u_lup,llm+1)=0   
              wwuu(ij+u_ldown,llm+1)=0
          ENDDO
       END DO
    ENDIF

    !$OMP BARRIER
  END SUBROUTINE caldyn_BC
   
  SUBROUTINE caldyn(write_out,f_phis, f_ps, f_mass, f_theta_rhodz, f_u, f_q, &
       f_geopot, f_hflux, f_wflux, f_dps, f_dmass, f_dtheta_rhodz, f_du)
    USE icosa
    USE observable_mod
    USE disvert_mod, ONLY : caldyn_eta, eta_mass
    USE vorticity_mod
    USE kinetic_mod
    USE theta2theta_rhodz_mod
    USE wind_mod
    USE mpipara
    USE trace
    USE omp_para
    USE output_field_mod
    USE checksum_mod
    IMPLICIT NONE
    LOGICAL,INTENT(IN)    :: write_out
    TYPE(t_field),POINTER :: f_phis(:)
    TYPE(t_field),POINTER :: f_ps(:)
    TYPE(t_field),POINTER :: f_mass(:)
    TYPE(t_field),POINTER :: f_theta_rhodz(:)
    TYPE(t_field),POINTER :: f_u(:)
    TYPE(t_field),POINTER :: f_q(:)
    TYPE(t_field),POINTER :: f_geopot(:)
    TYPE(t_field),POINTER :: f_hflux(:), f_wflux(:)
    TYPE(t_field) :: f_dps(:)
    TYPE(t_field) :: f_dmass(:)
    TYPE(t_field) :: f_dtheta_rhodz(:)
    TYPE(t_field) :: f_du(:)
    
    REAL(rstd),POINTER :: ps(:), dps(:)
    REAL(rstd),POINTER :: mass(:,:), theta_rhodz(:,:,:), dtheta_rhodz(:,:,:)
    REAL(rstd),POINTER :: u(:,:), du(:,:), hflux(:,:), wflux(:,:)
    REAL(rstd),POINTER :: qu(:,:)
    REAL(rstd),POINTER :: qv(:,:)

! temporary shared variable
    REAL(rstd),POINTER  :: theta(:,:,:)  
    REAL(rstd),POINTER  :: pk(:,:)
    REAL(rstd),POINTER  :: geopot(:,:)
    REAL(rstd),POINTER  :: convm(:,:) 
    REAL(rstd),POINTER  :: wwuu(:,:)
        
    INTEGER :: ind
    LOGICAL,SAVE :: first=.TRUE.
!$OMP THREADPRIVATE(first)
    
    IF (first) THEN 
      first=.FALSE.
      IF(caldyn_eta==eta_mass) THEN
         CALL init_message(f_ps,req_i1,req_ps)
      ELSE
         CALL init_message(f_mass,req_i1,req_mass)
      END IF
      CALL init_message(f_theta_rhodz,req_i1,req_theta_rhodz)
      CALL init_message(f_u,req_e1_vect,req_u)
      CALL init_message(f_qu,req_e1_scal,req_qu)
      ! Overlapping com/compute (deactivated) : MPI messages need to be sent at first call to caldyn
      ! This is needed only once : the next ones will be sent by timeloop
!      IF(caldyn_eta==eta_mass) THEN
!         CALL send_message(f_ps,req_ps) 
!         CALL wait_message(req_ps)  
!      ELSE
!         CALL send_message(f_mass,req_mass) 
!         CALL wait_message(req_mass)  
!      END IF
    ENDIF
    
    CALL trace_start("caldyn")

    IF(caldyn_eta==eta_mass) THEN
       CALL send_message(f_ps,req_ps) ! COM00
       CALL wait_message(req_ps) ! COM00
    ELSE
       CALL send_message(f_mass,req_mass) ! COM00
       CALL wait_message(req_mass) ! COM00
    END IF
    
    CALL send_message(f_theta_rhodz,req_theta_rhodz) ! COM01
    CALL wait_message(req_theta_rhodz) ! COM01 Moved from caldyn_pvort
    CALL send_message(f_u,req_u) ! COM02
    CALL wait_message(req_u) ! COM02

    IF(.NOT.hydrostatic) THEN
       STOP 'caldyn_gcm may not be used yet when non-hydrostatic'
    END IF

    SELECT CASE(caldyn_conserv)
    CASE(energy) ! energy-conserving
       DO ind=1,ndomain
          IF (.NOT. assigned_domain(ind)) CYCLE
          CALL swap_dimensions(ind)
          CALL swap_geometry(ind)
          ps=f_ps(ind)
          u=f_u(ind)
          theta_rhodz = f_theta_rhodz(ind)
          mass=f_mass(ind)
          theta = f_theta(ind)
          qu=f_qu(ind)
          qv=f_qv(ind)
          pk = f_pk(ind)
          geopot = f_geopot(ind)  
          hflux=f_hflux(ind)
          convm = f_dmass(ind)
          dtheta_rhodz=f_dtheta_rhodz(ind)
          du=f_du(ind)
          CALL compute_pvort(ps,u,theta_rhodz(:,:,1), mass,theta,qu,qv) ! COM00 COM01 COM02
!          CALL compute_theta(ps,theta_rhodz, mass,theta)
!          CALL compute_pvort_only(u,mass,qu,qv)

          CALL compute_geopot(mass,theta, ps,pk,geopot)
!          du(:,:)=0.
!          CALL compute_caldyn_fast(0.,u,mass,theta,pk,geopot,du)
       ENDDO       

       CALL send_message(f_u,req_u) ! COM02
       CALL wait_message(req_u) ! COM02
       CALL send_message(f_qu,req_qu) ! COM03
       CALL wait_message(req_qu) ! COM03

       DO ind=1,ndomain
          IF (.NOT. assigned_domain(ind)) CYCLE
          CALL swap_dimensions(ind)
          CALL swap_geometry(ind)
          ps=f_ps(ind)
          u=f_u(ind)
          theta_rhodz = f_theta_rhodz(ind)
          mass=f_mass(ind)
          theta = f_theta(ind)
          qu=f_qu(ind)
          qv=f_qv(ind)
          pk = f_pk(ind)
          geopot = f_geopot(ind)  
          hflux=f_hflux(ind)
          convm = f_dmass(ind)
          dtheta_rhodz=f_dtheta_rhodz(ind)
          du=f_du(ind)

          CALL compute_caldyn_horiz(u,mass,qu,theta,pk,geopot, hflux,convm,dtheta_rhodz(:,:,1),du)
!          CALL compute_caldyn_slow_hydro(u,mass,hflux,du, .FALSE.) ! FIXME
!          CALL compute_caldyn_Coriolis(hflux,theta,qu, convm,dtheta_rhodz,du) 
          IF(caldyn_eta==eta_mass) THEN
             wflux=f_wflux(ind)
             wwuu=f_wwuu(ind)
             dps=f_dps(ind)
             CALL compute_caldyn_vert(u,theta,mass,convm, wflux,wwuu, dps, dtheta_rhodz(:,:,1), du)
          END IF
       ENDDO       
   
    CASE(enstrophy) ! enstrophy-conserving
       DO ind=1,ndomain
          IF (.NOT. assigned_domain(ind)) CYCLE
          CALL swap_dimensions(ind)
          CALL swap_geometry(ind)
          ps=f_ps(ind)
          u=f_u(ind)
          theta_rhodz=f_theta_rhodz(ind)
          mass=f_mass(ind)
          theta = f_theta(ind)
          qu=f_qu(ind)
          qv=f_qv(ind)
          CALL compute_pvort(ps,u,theta_rhodz(:,:,1), mass,theta,qu,qv)
          pk = f_pk(ind)
          geopot = f_geopot(ind)  
          CALL compute_geopot(ps,mass,theta, pk,geopot)
          hflux=f_hflux(ind)
          convm = f_dmass(ind)
          dtheta_rhodz=f_dtheta_rhodz(ind)
          du=f_du(ind)
          CALL compute_caldyn_horiz(u,mass,qu,theta,pk,geopot, hflux,convm,dtheta_rhodz(:,:,1),du)
          IF(caldyn_eta==eta_mass) THEN
             wflux=f_wflux(ind)
             wwuu=f_wwuu(ind)
             dps=f_dps(ind)
             CALL compute_caldyn_vert(u,theta,mass,convm, wflux,wwuu, dps, dtheta_rhodz, du)
          END IF
       ENDDO
       
    CASE DEFAULT
       STOP
    END SELECT

!$OMP BARRIER
    !    CALL check_mass_conservation(f_ps,f_dps)
    CALL trace_end("caldyn")
!!$OMP BARRIER
    
END SUBROUTINE caldyn

!-------------------------------- Diagnostics ----------------------------

  SUBROUTINE check_mass_conservation(f_ps,f_dps)
  USE icosa
  USE mpipara
  IMPLICIT NONE
    TYPE(t_field),POINTER :: f_ps(:)
    TYPE(t_field),POINTER :: f_dps(:)
    REAL(rstd),POINTER :: ps(:)
    REAL(rstd),POINTER :: dps(:)
    REAL(rstd) :: mass_tot,dmass_tot
    INTEGER :: ind,i,j,ij
    
    mass_tot=0
    dmass_tot=0
    
    CALL transfert_request(f_dps,req_i1)
    CALL transfert_request(f_ps,req_i1)

    DO ind=1,ndomain
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)

      ps=f_ps(ind)
      dps=f_dps(ind)

      DO j=jj_begin,jj_end
        DO i=ii_begin,ii_end
          ij=(j-1)*iim+i
          IF (domain(ind)%own(i,j)) THEN
            mass_tot=mass_tot+ps(ij)*Ai(ij)/g
            dmass_tot=dmass_tot+dps(ij)*Ai(ij)/g
          ENDIF
        ENDDO
      ENDDO
   
    ENDDO
    IF (is_mpi_root) PRINT*, "mass_tot ", mass_tot,"      dmass_tot ",dmass_tot        

  END SUBROUTINE check_mass_conservation  
  
END MODULE caldyn_gcm_mod