MODULE caldyn_gcm_mod
   USE icosa
   USE transfert_mod
   USE omp_para
   USE mpipara
   IMPLICIT NONE
   PRIVATE

   PUBLIC init_caldyn, caldyn_BC, caldyn

CONTAINS

   SUBROUTINE init_caldyn
      USE abort_mod, ONLY: dynamico_abort, abort_acc
      USE getin_mod, ONLY: getin_c, getin_r
      USE compute_caldyn_mod, ONLY: compute_planetvel
      USE thermodynamics_mod, ONLY: init_thermodynamics
      USE caldyn_vars_mod
      CHARACTER(len=255) :: def
      INTEGER            :: ind
      REAL(rstd), POINTER :: planetvel(:)

      boussinesq = .FALSE.
      hydrostatic = .TRUE.
      def = "hydrostatic"
      CALL getin("equations", def)
      SELECT CASE (TRIM(def))
      CASE ("hydrostatic")
         IF (is_mpi_root) PRINT *, 'Solving traditional, shallow-atmosphere hydrostatic primitive equations'
      CASE ("nonhydrostatic")
         IF (is_mpi_root) PRINT *, 'Solving traditional, shallow-atmosphere fully-compressible equations'
         hydrostatic = .FALSE.
      CASE ("boussinesq")
         IF (is_mpi_root) PRINT *, 'Solving multi-layer thermal shallow water equations'
         boussinesq = .TRUE.
      CASE DEFAULT
         IF (is_mpi_root) PRINT *, "Bad selector for variable 'equations' : <", &
            TRIM(def), '> options are <hydrostatic>, <nonhydrostatic>, <boussinesq>'
         STOP
      END SELECT

      IF (.NOT. hydrostatic) THEN
         CALL abort_acc("hydrostatic /= .TRUE.")
      END IF

      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)
      dysl_caldyn_vert = dysl
      CALL getin("dysl_caldyn_vert", dysl_caldyn_vert)

      def = 'advective'
      CALL getin('caldyn_vert', def)
      SELECT CASE (TRIM(def))
      CASE ('advective')
         caldyn_vert_variant = caldyn_vert_noncons
      CASE ('conservative')
         IF (.NOT. dysl_caldyn_vert) THEN
            IF (is_mpi_root) PRINT *, 'caldyn_vert=cons requires dysl_caldyn_vert=.TRUE.'
            STOP
         END IF
         caldyn_vert_variant = caldyn_vert_cons
      CASE DEFAULT
         IF (is_mpi_root) PRINT *, 'Bad selector for variable caldyn_vert : <', &
            TRIM(def), '> options are <conservative>, <advective>'
         STOP
      END SELECT
      IF (is_master) PRINT *, 'caldyn_vert=', def

      ! dysl is not supported with openACC
      IF (dysl_geopot .OR. dysl_caldyn_fast .OR. dysl_slow_hydro .OR. dysl_pvort_only .OR. dysl_caldyn_coriolis) THEN
         CALL abort_acc("dysl not supported")
      END IF

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

      def = 'trisk'
      CALL getin('caldyn_kinetic', def)
      SELECT CASE (TRIM(def))
      CASE ('trisk')
         caldyn_kinetic = kinetic_trisk
      CASE ('consistent')
         caldyn_kinetic = kinetic_consistent
      CASE DEFAULT
         IF (is_mpi_root) PRINT *, 'Bad selector for variable caldyn_kinetic : <', &
            TRIM(def), '> options are <trisk>, <consistent>'
         STOP
      END SELECT
      IF (is_master) PRINT *, 'caldyn_kinetic=', def

      IF (.TRUE.) CALL init_caldyn_thermo_obsolete ! we should be able to skip that, but for the moment it breaks

      CALL init_thermodynamics(thermo, getin_c, getin_r, dynamico_abort)
      nqdyn = thermo%nqdyn
      IF (is_master) PRINT *, 'nqdyn =', nqdyn
      CALL allocate_caldyn

      SELECT CASE (grid_type)
      CASE (grid_ico)
         ! NB : unused now, would be useful for non-traditional / deep-atmosphere
         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
      CASE (grid_unst)
      END SELECT

   END SUBROUTINE init_caldyn

   SUBROUTINE init_caldyn_thermo_obsolete
      USE caldyn_vars_mod, ONLY: hydrostatic
      CHARACTER(len=255) :: 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', 'binarygas', 'entropy')
         caldyn_thermo = thermo_theta
         physics_thermo = thermo_dry
      CASE ('variable_Cp') ! copié de ce qui est fait pour l'entropy
         caldyn_thermo = thermo_theta
         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>, <binarygas>, <variable_Cp>'
         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
            PRINT *, 'Internal error in init_caldyn : illegal value of caldyn_thermo = ', caldyn_thermo
            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

   END SUBROUTINE init_caldyn_thermo_obsolete

   SUBROUTINE allocate_caldyn
      USE caldyn_vars_mod, ONLY: f_out_u, f_qu, f_qv, f_pk, f_wwuu, f_berni, f_Ftheta, f_planetvel, f_Fel, f_gradPhi2, f_wil, f_Wetadot
      USE caldyn_vars_mod, ONLY: hydrostatic
      CALL allocate_field(f_out_u, field_u, type_real, llm)
      CALL allocate_field(f_qu, field_u, type_real, llm, ondevice=.TRUE.)
      CALL allocate_field(f_qv, field_z, type_real, llm, ondevice=.TRUE.)
      CALL allocate_field(f_pk, field_t, type_real, llm, nqdyn, name='pk', ondevice=.TRUE.)
      CALL allocate_field(f_wwuu, field_u, type_real, llm + 1, name='wwuu', ondevice=.TRUE.)
      CALL allocate_field(f_Ftheta, field_u, type_real, llm, name='Ftheta', ondevice=.TRUE.)
      CALL allocate_field(f_berni, field_t, type_real, llm, name='berni', ondevice=.TRUE.)
      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 caldyn_vars_mod, ONLY: f_wwuu
      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, ij

      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)

            SELECT CASE (grid_type)
            CASE (grid_unst)
               ! FIXME OMP
               DO ij = 1, primal_num
                  geopot(1, ij) = phis(ij)
                  wflux(1, ij) = 0.
                  wflux(llm + 1, ij) = 0.
               END DO
               DO ij = 1, edge_num
                  wwuu(1, ij) = 0.
                  wwuu(llm + 1, ij) = 0.
               END DO
            CASE (grid_ico)
               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
               END DO
            END SELECT
         END DO
      END IF

      !$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 observable_mod
      USE disvert_mod, ONLY: caldyn_eta, eta_mass, bp
      USE caldyn_vars_mod, ONLY: f_qu, f_qv, f_pk, f_berni, f_wwuu, f_planetvel !, f_Fel, f_gradPhi2, f_wil, f_Wetadot
      USE caldyn_vars_mod, ONLY: caldyn_conserv, conserv_energy, conserv_gassmann, conserv_enstrophy, hydrostatic
      USE caldyn_vars_mod, ONLY: req_ps, req_mass, req_theta_rhodz, req_u, req_qu
      USE compute_caldyn_mod, ONLY: compute_caldyn_vert, compute_pvort, compute_caldyn_horiz, compute_geopot
      USE trace
      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(:, :)
      REAL(rstd), POINTER :: berni(:, :)

! 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
      END IF

      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 (conserv_energy, conserv_gassmann) ! 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)
            berni = f_berni(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(:,:,1), qu, qv)
            CALL compute_geopot(mass, theta, ps, berni, pk, geopot)
         END DO

         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(:,:,1), pk(:,:,1), 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
         END DO

      CASE (conserv_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(:,:,1), qu, qv)
            berni = f_berni(ind)
            pk = f_pk(ind)
            geopot = f_geopot(ind)
            CALL compute_geopot(mass, theta, ps, berni, 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(:,:,1), pk(:,:,1), 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
         END DO

      CASE DEFAULT
         PRINT *, 'Invalid value of caldyn_conserv : ', caldyn_conserv
         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)
      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
               END IF
            END DO
         END DO

      END DO
      IF (is_mpi_root) PRINT *, "mass_tot ", mass_tot, "      dmass_tot ", dmass_tot

   END SUBROUTINE check_mass_conservation

END MODULE caldyn_gcm_mod
