MODULE dissip_gcm_mod
   USE icosa
   USE abort_mod
   USE omp_para
   USE trace
   IMPLICIT NONE
   PRIVATE

   TYPE(t_field), POINTER, SAVE :: f_due_gradiv(:)
   TYPE(t_field), POINTER, SAVE :: f_due_gradrot(:)

   TYPE(t_field), POINTER, SAVE :: f_dtheta_diss(:)
   TYPE(t_field), POINTER, SAVE :: f_dtheta_rhodz_diss(:)
   TYPE(t_message), SAVE :: req_due_gradiv, req_due_gradrot, req_dtheta

   INTEGER, SAVE :: nitergdiv = 1
!$OMP THREADPRIVATE(nitergdiv)
   INTEGER, SAVE :: nitergrot = 1
!$OMP THREADPRIVATE(nitergrot)
   INTEGER, SAVE :: niterdivgrad = 1
!$OMP THREADPRIVATE(niterdivgrad)

   REAL, ALLOCATABLE, SAVE :: tau_graddiv(:)
!$OMP THREADPRIVATE(tau_graddiv)
   REAL, ALLOCATABLE, SAVE :: tau_gradrot(:)
!$OMP THREADPRIVATE(tau_gradrot)
   REAL, ALLOCATABLE, SAVE :: tau_divgrad(:)
!$OMP THREADPRIVATE(tau_divgrad)

   REAL, SAVE :: cgraddiv
!$OMP THREADPRIVATE(cgraddiv)
   REAL, SAVE :: cgradrot
!$OMP THREADPRIVATE(cgradrot)
   REAL, SAVE :: cdivgrad
!$OMP THREADPRIVATE(cdivgrad)

   INTEGER, SAVE :: rayleigh_friction_type, rayleigh_shear
!$OMP THREADPRIVATE(rayleigh_friction_type)
!$OMP THREADPRIVATE(rayleigh_shear)
   REAL, SAVE    :: rayleigh_tau, rayleigh_limlat
!$OMP THREADPRIVATE(rayleigh_tau)
!$OMP THREADPRIVATE(rayleigh_limlat)

   REAL, SAVE    :: dtdissip
!$OMP THREADPRIVATE(dtdissip)

   PUBLIC init_dissip, dissip

   INTERFACE
      SUBROUTINE plugin_vert_prof_dissip(llm, presnivs, preff, scale_height, zvert)
         INTEGER, INTENT(IN) :: llm
         REAL, INTENT(IN) :: presnivs(:)
         REAL, INTENT(IN) :: preff
         REAL, INTENT(IN) :: scale_height
         REAL, INTENT(OUT) :: zvert(:)
      END SUBROUTINE
   END INTERFACE

   PROCEDURE(plugin_vert_prof_dissip), POINTER :: vert_prof_dissip_plugin => NULL()

   PUBLIC vert_prof_dissip_plugin

CONTAINS

   SUBROUTINE allocate_dissip
      CALL allocate_field(f_due_gradiv, field_u, type_real, llm)
      CALL allocate_field(f_due_gradrot, field_u, type_real, llm)
      CALL allocate_field(f_dtheta_diss, field_t, type_real, llm)
      CALL allocate_field(f_dtheta_rhodz_diss, field_t, type_real, llm)
      ALLOCATE (tau_graddiv(llm))
      ALLOCATE (tau_gradrot(llm))
      ALLOCATE (tau_divgrad(llm))
   END SUBROUTINE allocate_dissip

   SUBROUTINE init_dissip
      USE disvert_mod
      USE mpi_mod
      USE mpipara
      USE transfert_mod
      USE time_mod
      USE transfert_omp_mod
      REAL(rstd)            :: tau
      CHARACTER(len=255)    :: rayleigh_friction_key

      ! @getin_doc rayleigh_friction_type realm=dissipation type=String default=none
      !  Selects vertical profile for Rayleigh friction. Valid values :
      !  * `dcmip2_schaer_noshear` : relaxes toward the non-sheared wind profile of DCMIP2008 test case 2.1
      !  * `dcmip2_schaer_shear`   : relaxes toward the sheared wind profile of DCMIP2008 test case 2.2
      rayleigh_friction_key = 'none'
      CALL getin("rayleigh_friction_type", rayleigh_friction_key)
      SELECT CASE (TRIM(rayleigh_friction_key))
      CASE ('none')
         rayleigh_friction_type = 0
         IF (is_master) PRINT *, 'No Rayleigh friction'
      CASE ('dcmip2_schaer_noshear')
         rayleigh_friction_type = 1
         rayleigh_shear = 0
         IF (is_master) PRINT *, 'Rayleigh friction : Schaer-like mountain without shear DCMIP2.1'
      CASE ('dcmip2_schaer_shear')
         rayleigh_shear = 1
         rayleigh_friction_type = 2
         IF (is_master) PRINT *, 'Rayleigh friction : Schaer-like mountain with shear DCMIP2.2'
      CASE ('giant_liu_schneider')
         rayleigh_friction_type = 99
         IF (is_master) PRINT *, 'Rayleigh friction : giant planets Liu Schneider 2010'
      CASE DEFAULT
         IF (is_master) PRINT *, 'Bad selector : rayleigh_friction_type =', TRIM(rayleigh_friction_key), &
            ' in dissip_gcm.f90/init_dissip'
         STOP
      END SELECT

      ! @getin_doc rayleigh_friction_tau realm=dissipation type=Real default=0
      !   Relaxation time for Rayleigh friction
      IF (rayleigh_friction_type > 0) THEN
         rayleigh_tau = 0.
         CALL getin("rayleigh_friction_tau", rayleigh_tau)
         rayleigh_tau = rayleigh_tau/scale_factor
         IF (rayleigh_tau <= 0) THEN
            IF (is_master) PRINT *, 'Forbidden : negative value for rayleigh_friction_tau =', rayleigh_tau
            STOP
         END IF
         IF (rayleigh_friction_type == 99) THEN
            rayleigh_limlat = 0.
            CALL getin("rayleigh_limlat", rayleigh_limlat)
            rayleigh_limlat = rayleigh_limlat*3.14159/180.
         END IF
      END IF

      CALL allocate_dissip

      CALL init_message(f_due_gradiv, req_e1_vect, req_due_gradiv)
      CALL init_message(f_due_gradrot, req_e1_vect, req_due_gradrot)
      CALL init_message(f_dtheta_rhodz_diss, req_i1, req_dtheta)

      tau_graddiv(:) = 5000
      CALL getin("tau_graddiv", tau)
      tau_graddiv(:) = tau/scale_factor

      CALL getin("nitergdiv", nitergdiv)

      tau_gradrot(:) = 5000
      CALL getin("tau_gradrot", tau)
      tau_gradrot(:) = tau/scale_factor

      CALL getin("nitergrot", nitergrot)

      tau_divgrad(:) = 5000
      CALL getin("tau_divgrad", tau)
      tau_divgrad(:) = tau/scale_factor

      CALL getin("niterdivgrad", niterdivgrad)

      IF (grid_type == grid_ico) THEN
         CALL dissip_constants
         CALL dissip_profile
         CALL dissip_timescale
      ELSE ! FIXME unstructured
         itau_dissip = 1000000000
      END IF

   END SUBROUTINE init_dissip

   SUBROUTINE dissip_constants
      ! The SAVE attribute is required here, otherwise allocate_field will not work with OpenMP
      TYPE(t_field), POINTER, SAVE  :: f_u(:)
      TYPE(t_field), POINTER, SAVE  :: f_du(:)
      TYPE(t_field), POINTER, SAVE  :: f_theta(:)
      TYPE(t_field), POINTER, SAVE  :: f_dtheta(:)
      REAL(rstd), POINTER :: u(:)
      REAL(rstd), POINTER :: du(:)
      REAL(rstd), POINTER :: theta(:)
      REAL(rstd), POINTER :: dtheta(:)
      REAL(rstd) :: dumax, dthetamax
      INTEGER :: it, iter, ind
      CALL allocate_field(f_u, field_u, type_real)
      CALL allocate_field(f_du, field_u, type_real)
      CALL allocate_field(f_theta, field_t, type_real)
      CALL allocate_field(f_dtheta, field_t, type_real)

      PRINT *, '=========', SHAPE(f_u)

      !--------------------------- Compute cgraddiv ----------------------------
      cgraddiv = -1.
      CALL random_vel(f_u)
      DO it = 1, 20
         DO iter = 1, nitergdiv
            CALL transfert_request(f_u, req_e1_vect)
            DO ind = 1, ndomain
               IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
               CALL swap_dimensions(ind)
               CALL swap_geometry(ind)
               u = f_u(ind)
               du = f_du(ind)
               CALL compute_gradiv(u, du, 1, 1)
               u = du
            END DO
         END DO
         CALL transfert_request(f_du, req_e1_vect)
         dumax = max_vel(f_du)
         CALL rescale_field(1./dumax, f_du, f_u)
         IF (is_master) PRINT *, "gradiv : it :", it, ": dumax", dumax
      END DO

      cgraddiv = dumax**(-1./nitergdiv)
      IF (is_master) THEN
         PRINT *, "gradiv : dumax", dumax
         PRINT *, 'mean T-cell edge size (km)', 1.45*radius/iim_glo/1000., &
            'effective T-cell half-edge size (km)', dumax**(-.5/nitergdiv)/1000
         PRINT *, 'Max. time step assuming c=340 m/s and Courant number=3 (ARK2.3) :', &
            3./340.*dumax**(-.5/nitergdiv)
         PRINT *, "cgraddiv : ", cgraddiv
      END IF

      !----------------- Compute cgradrot --------------------
      cgradrot = -1.
      CALL random_vel(f_u)
      DO it = 1, 20
         DO iter = 1, nitergrot
            CALL transfert_request(f_u, req_e1_vect)
            DO ind = 1, ndomain
               IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
               CALL swap_dimensions(ind)
               CALL swap_geometry(ind)
               u = f_u(ind)
               du = f_du(ind)
               CALL compute_gradrot(u, du, 1, 1)
               u = du
            END DO
         END DO
         CALL transfert_request(f_du, req_e1_vect)
         dumax = max_vel(f_du)
         CALL rescale_field(1./dumax, f_du, f_u)
         IF (is_master) PRINT *, "gradrot : it :", it, ": dumax", dumax
      END DO

      cgradrot = dumax**(-1./nitergrot)
      IF (is_master) PRINT *, "gradrot : dumax", dumax
      IF (is_master) PRINT *, "cgradrot : ", cgradrot

      !----------------- Compute cdivgrad --------------------

      cdivgrad = -1.
      CALL random_scalar(f_theta)
      DO it = 1, 20
         DO iter = 1, niterdivgrad
            CALL transfert_request(f_theta, req_i1)
            DO ind = 1, ndomain
               IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
               CALL swap_dimensions(ind)
               CALL swap_geometry(ind)
               theta = f_theta(ind)
               dtheta = f_dtheta(ind)
               CALL compute_divgrad(theta, dtheta, 1, 1)
               theta = dtheta
            END DO
         END DO
         CALL transfert_request(f_dtheta, req_i1)
         dthetamax = max_scalar(f_dtheta)
         IF (is_master) PRINT *, "divgrad : it :", it, ": dthetamax", dthetamax
         CALL rescale_field(1./dthetamax, f_dtheta, f_theta)
      END DO

      cdivgrad = dthetamax**(-1./niterdivgrad)
      IF (is_master) PRINT *, "divgrad : divgrad", dthetamax
      IF (is_master) PRINT *, "cdivgrad : ", cdivgrad

   END SUBROUTINE dissip_constants

   SUBROUTINE dissip_profile
      USE disvert_mod
      ! parameters used by the various profiles
      ! IF planet_type == earth
      !    IF dissip_vert_prof == 0
      !      none
      !    IF dissip_vert_prof == 1
      !      dissip_zref, dissip_deltaz, dissip_factz
      ! IF planet_type == other
      !    IF dissip_vert_prof == 0
      !      dissip_fac_mid
      !      + dissip_deltaz, dissip_hdelta, dissip_fac_up, dissip_pupstart IF ok_strato
      !    IF dissip_vert_prof == 1
      !      fac_mid, fac_up, startalt, delta => middle (hardcoded), scaleheight
      !

      REAL(rstd), PARAMETER :: fact = 2., &
                               fac_mid = 3., & ! coefficient for lower/middle atmosphere
                               fac_up = 30., & ! coefficient for upper atmosphere
                               startalt = 70., & ! altitude (in km) for the transition from middle to upper atm.
                               delta = 30., & ! Size (in km) of the transition region between middle/upper atm.
                               middle = startalt + delta/2.
      REAL(rstd) :: dissip_zref, dissip_deltaz, dissip_factz, &
                    dissip_hdelta, dissip_fac_up, dissip_fac_mid, dissip_pupstart, &
                    scaleheight, &
                    zz, pseudoz, pup, &
                    sigma(llm), zvert(llm)
      CHARACTER(LEN=255) :: planet_type ! earth, other, other_strato ; other_strato = other + ok_strato
      LOGICAL :: ok_strato
      INTEGER    :: l, dissip_vert_prof

      ! select vertical profile of horizontal dissipation coefficients, see also [781]

      planet_type = 'earth'
      CALL getin('dissip_planet_type', planet_type)
      SELECT CASE (planet_type)
      CASE ('earth', 'other')
         ok_strato = .FALSE.
      CASE ('other_strato')
         planet_type = 'other'
         ok_strato = .TRUE.
      CASE DEFAULT
         STOP "Invalid value of dissip_planet_type, valid values are <earth>, <other>, <other_strato>"
      END SELECT

      dissip_vert_prof = 0
      CALL getin('dissip_vert_prof', dissip_vert_prof)

      SELECT CASE (dissip_vert_prof)
      CASE (0)
         IF (TRIM(planet_type) == 'other') THEN
            ! Default values given below are for a Venus-like planet
            CALL getin('dissip_fac_mid', dissip_fac_mid)
            dissip_fac_mid = 2.
            IF (ok_strato) THEN
               dissip_fac_up = 50.
               ! deltaz et hdelta in km
               dissip_deltaz = 30.
               dissip_hdelta = 5.
               ! pupstart in Pa
               dissip_pupstart = 1.e4
               CALL getin('dissip_deltaz', dissip_deltaz)
               CALL getin('dissip_hdelta', dissip_hdelta)
               CALL getin('dissip_fac_up', dissip_fac_up)
               CALL getin('dissip_pupstart', dissip_pupstart)
            END IF
         END IF
      CASE (1)
         IF (TRIM(planet_type) == 'earth') THEN
            dissip_zref = 30.
            dissip_deltaz = 10.
            dissip_factz = 4.
            CALL getin('dissip_zref', dissip_zref)
            CALL getin('dissip_deltaz', dissip_deltaz)
            CALL getin('dissip_factz', dissip_factz)
         ELSE
            ! fac_mid, fac_up, startalt, delta => middle are hardcoded
            CALL getin('dissip_scaleheight', scaleheight)
         END IF
      CASE DEFAULT
         STOP 'Invalid value of dissip_vert_prof : valid values are 0,1'
      END SELECT

      IF (ap_bp_present) THEN
         sigma(:) = preff/presnivs(:)
      ELSE
         sigma(:) = 1.
      END IF

      SELECT CASE (TRIM(planet_type))
      CASE ('earth')
         DO l = 1, llm
            IF (dissip_vert_prof == 1) THEN
               pseudoz = 8.*LOG(sigma(l))
               zvert(l) = 1 + &
                          (TANH((pseudoz - dissip_zref)/dissip_deltaz) + 1.)/2. &
                          *(dissip_factz - 1.)
            ELSE
               zz = 1.-sigma(l)
               zvert(l) = fact - (fact - 1.)/(1.+zz*zz)
            END IF
         END DO
      CASE ('other')
         SELECT CASE (dissip_vert_prof)
         CASE (0)
            ! First step: going from 1 to dissip_fac_mid (in gcm.def)
            !============
            DO l = 1, llm
               zz = 1.-sigma(l)
               zvert(l) = dissip_fac_mid - (dissip_fac_mid - 1.)/(1.+zz*zz)
            END DO
            !
            ! Second step if ok_strato:  from dissip_fac_mid to dissip_fac_up (in gcm.def)
            !==========================
            ! Utilisation de la fonction tangente hyperbolique pour augmenter
            ! arbitrairement la dissipation et donc la stabilite du modele a
            ! partir d'une certaine altitude.
            !
            !   Le facteur multiplicatif de basse atmosphere etant deja pris
            !   en compte, il faut diviser le facteur multiplicatif de haute
            !   atmosphere par celui-ci.

            IF (ok_strato) THEN
               Pup = dissip_pupstart*exp(-0.5*dissip_deltaz/dissip_hdelta)
               DO l = 1, llm
                  zvert(l) = zvert(l)*(1.0 + ((dissip_fac_up/dissip_fac_mid - 1) &
                                              *(1.-(0.5*(1 + tanh(-6./dissip_deltaz* &
                                                                  (-dissip_hdelta*log(presnivs(l)/Pup))))))))
               END DO
            END IF
         CASE (1)
            DO l = 1, llm
               zz = 1.-sigma(l)
               zvert(l) = fac_mid - (fac_mid - 1.)/(1.+zz*zz)
               zvert(l) = zvert(l)*(1.0 + ((fac_up/fac_mid - 1)* &
                                           (1.-(0.5*(1 + tanh(-6./ &
                                                              delta*(scaleheight*(-log(presnivs(l)/preff)) - middle))))) &
                                           ))
            END DO
         END SELECT
      END SELECT

      IF (is_master) PRINT *, 'vertical profile of horizontal dissipation : ', zvert(:)

      DO l = 1, llm
         tau_graddiv(l) = tau_graddiv(l)/zvert(l)
         tau_gradrot(l) = tau_gradrot(l)/zvert(l)
         tau_divgrad(l) = tau_divgrad(l)/zvert(l)
      END DO
   END SUBROUTINE dissip_profile

   SUBROUTINE dissip_timescale
      INTEGER :: l
      REAL(rstd) :: mintau
      mintau = tau_graddiv(1)
      DO l = 1, llm
         mintau = MIN(mintau, tau_graddiv(l))
         mintau = MIN(mintau, tau_gradrot(l))
         mintau = MIN(mintau, tau_divgrad(l))
      END DO

      IF (rayleigh_friction_type > 0) mintau = MIN(mintau, rayleigh_tau)

      IF (mintau > 0) THEN
         IF (itau_dissip == 0) THEN
            IF (is_master) PRINT *, "init_dissip: Automatic computation of itau_dissip..."
            itau_dissip = INT(mintau/dt)
         END IF
      ELSE
         IF (is_master) PRINT *, "init_dissip: No dissipation time set, setting itau_dissip to 1000000000"
         itau_dissip = 100000000
      END IF
      itau_dissip = MAX(1, itau_dissip)
      dtdissip = itau_dissip*dt

      IF (is_master) THEN
         PRINT *, "init_dissip: rayleigh_tau", rayleigh_tau, "mintau ", mintau
         PRINT *, "init_dissip: itau_dissip", itau_dissip, " dtdissip ", dtdissip
      END IF

      IF (dtdissip > 2.*mintau) THEN
         IF (is_master) PRINT *, 'The CFL condition for dissipation dtdissip<2*mintau is violated : dtdissip, mintau ', &
            dtdissip, mintau
         STOP
      END IF

   END SUBROUTINE dissip_timescale

   SUBROUTINE dissip(f_ps, f_mass, f_phis, f_geopot, f_theta_rhodz, f_ue, f_dtheta_rhodz, f_due)
      TYPE(t_field), POINTER :: f_ps(:), f_mass(:), f_phis(:), f_geopot(:)
      TYPE(t_field), POINTER :: f_theta_rhodz(:), f_dtheta_rhodz(:)
      TYPE(t_field), POINTER :: f_ue(:), f_due(:)

      REAL(rstd), POINTER         :: due(:, :)
      REAL(rstd), POINTER         :: phi(:, :), ue(:, :)
      REAL(rstd), POINTER         :: due_diss1(:, :)
      REAL(rstd), POINTER         :: due_diss2(:, :)
      REAL(rstd), POINTER         :: dtheta_rhodz(:, :, :)
      REAL(rstd), POINTER         :: dtheta_rhodz_diss(:, :)

      INTEGER :: ind, l, ij, nn

!$OMP BARRIER

      CALL trace_start("dissip")
      CALL gradiv(f_ue, f_due_gradiv)
      CALL gradrot(f_ue, f_due_gradrot)
      CALL divgrad_theta_rhodz(f_mass, f_theta_rhodz, f_dtheta_rhodz_diss)

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         due = f_due(ind)
         due_diss1 = f_due_gradiv(ind)
         due_diss2 = f_due_gradrot(ind)
         dtheta_rhodz = f_dtheta_rhodz(ind)
         dtheta_rhodz_diss = f_dtheta_rhodz_diss(ind)

         DO l = ll_begin, ll_end
!DIR$ SIMD
            DO ij = ij_begin, ij_end
               due(ij + u_right, l) = -0.5*(due_diss1(ij + u_right, l)/tau_graddiv(l) &
                                            + due_diss2(ij + u_right, l)/tau_gradrot(l))*itau_dissip
               due(ij + u_lup, l) = -0.5*(due_diss1(ij + u_lup, l)/tau_graddiv(l) &
                                          + due_diss2(ij + u_lup, l)/tau_gradrot(l))*itau_dissip
               due(ij + u_ldown, l) = -0.5*(due_diss1(ij + u_ldown, l)/tau_graddiv(l) &
                                            + due_diss2(ij + u_ldown, l)/tau_gradrot(l))*itau_dissip
               dtheta_rhodz(ij, l, 1) = -0.5*dtheta_rhodz_diss(ij, l)/tau_divgrad(l)*itau_dissip
            END DO
         END DO

         IF (rayleigh_friction_type > 0) THEN
            IF (rayleigh_friction_type < 99) THEN
               phi = f_geopot(ind)
               ue = f_ue(ind)
               DO l = ll_begin, ll_end
                  DO ij = ij_begin, ij_end
                     CALL relax(t_right, u_right)
                     CALL relax(t_lup, u_lup)
                     CALL relax(t_ldown, u_ldown)
                  END DO
               END DO
            ELSE
               ue = f_ue(ind)
               DO ij = ij_begin, ij_end
                  nn = ij + u_right
                  IF (ABS(lat_e(nn)) .gt. rayleigh_limlat) THEN
                     !print*, "latitude", lat_e(nn)*180./3.14159
                     due(nn, ll_begin:ll_begin + 1) = due(nn, ll_begin:ll_begin + 1) - (ue(nn, ll_begin:ll_begin + 1)/rayleigh_tau)
                  END IF
                  nn = ij + u_lup
                  IF (ABS(lat_e(nn)) .gt. rayleigh_limlat) THEN
                     due(nn, ll_begin:ll_begin + 1) = due(nn, ll_begin:ll_begin + 1) - (ue(nn, ll_begin:ll_begin + 1)/rayleigh_tau)
                  END IF
                  nn = ij + u_ldown
                  IF (ABS(lat_e(nn)) .gt. rayleigh_limlat) THEN
                     due(nn, ll_begin:ll_begin + 1) = due(nn, ll_begin:ll_begin + 1) - (ue(nn, ll_begin:ll_begin + 1)/rayleigh_tau)
                  END IF
               END DO
            END IF
         END IF
      END DO

      CALL trace_end("dissip")

      CALL write_dissip_tendencies
!$OMP BARRIER

   CONTAINS

      SUBROUTINE relax(shift_t, shift_u)
         USE dcmip_initial_conditions_test_1_2_3
         REAL(rstd) :: z, ulon, ulat, lon, lat, & ! input to test2_schaer_mountain
                       p, hyam, hybm, w, t, phis, ps, rho, q, &   ! unused input/output to test2_schaer_mountain
                       fz, u3d(3), uref
         REAL(rstd), PARAMETER :: zh = 2e4, ztop = 3e4  ! DCMIP values
         LOGICAL :: hybrid_eta
         INTEGER :: shift_u, shift_t, zcoords, nn
         z = (phi(ij, l) + phi(ij + shift_t, l))/(2.*g)
         IF (z > zh) THEN  ! relax only in the sponge layer z>zh
            nn = ij + shift_u
            zcoords = 1; hybrid_eta = .FALSE. ! use z instead of p or hyam/hybm
            CALL test2_schaer_mountain(lon_e(nn), lat_e(nn), p, z, zcoords, hybrid_eta, &
                                       hyam, hybm, rayleigh_shear, ulon, ulat, w, t, phis, ps, rho, q)
            u3d = ulon*elon_e(nn, :) ! ulat=0
            uref = sum(u3d*ep_e(nn, :))

            fz = sin((pi/2)*(z - zh)/(ztop - zh))
            fz = fz*fz/rayleigh_tau
            due(nn, l) = due(nn, l) - itau_dissip*fz*(ue(nn, l) - uref)
         END IF
      END SUBROUTINE relax

      SUBROUTINE write_dissip_tendencies
         USE observable_mod, ONLY: f_buf_ulon, f_buf_ulat
         USE wind_to_lonlat_mod
         USE output_field_mod

         CALL transfert_request(f_due_gradiv, req_e1_vect)
         CALL un2ulonlat(f_due_gradiv, f_buf_ulon, f_buf_ulat, (1./(tau_graddiv(1))))
         CALL output_field("dulon_diss1", f_buf_ulon)
         CALL output_field("dulat_diss1", f_buf_ulat)
         !
         CALL transfert_request(f_due_gradrot, req_e1_vect)
         CALL un2ulonlat(f_due_gradrot, f_buf_ulon, f_buf_ulat, (1./(tau_graddiv(1))))
         CALL output_field("dulon_diss2", f_buf_ulon)
         CALL output_field("dulat_diss2", f_buf_ulat)
      END SUBROUTINE write_dissip_tendencies

   END SUBROUTINE dissip

   SUBROUTINE gradiv(f_ue, f_due)
      TYPE(t_field), POINTER :: f_ue(:)
      TYPE(t_field), POINTER :: f_due(:)
      REAL(rstd), POINTER :: ue(:, :)
      REAL(rstd), POINTER :: due(:, :)
      INTEGER :: ind
      INTEGER :: it, l, ij

      CALL trace_start("gradiv")

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ue = f_ue(ind)
         due = f_due(ind)
         DO l = ll_begin, ll_end
            !DIR$ SIMD
            DO ij = ij_begin, ij_end
               due(ij + u_right, l) = ue(ij + u_right, l)
               due(ij + u_lup, l) = ue(ij + u_lup, l)
               due(ij + u_ldown, l) = ue(ij + u_ldown, l)
            END DO
         END DO
      END DO

      DO it = 1, nitergdiv
         CALL send_message(f_due, req_due_gradiv)
         CALL wait_message(req_due_gradiv)
         DO ind = 1, ndomain
            IF (.NOT. assigned_domain(ind)) CYCLE
            CALL swap_dimensions(ind)
            CALL swap_geometry(ind)
            due = f_due(ind)
            CALL compute_gradiv(due, due, ll_begin, ll_end)
         END DO
      END DO

      CALL trace_end("gradiv")
   END SUBROUTINE gradiv

   SUBROUTINE gradrot(f_ue, f_due)
      TYPE(t_field), POINTER :: f_ue(:)
      TYPE(t_field), POINTER :: f_due(:)
      REAL(rstd), POINTER :: ue(:, :)
      REAL(rstd), POINTER :: due(:, :)
      INTEGER :: ind, it, l, ij
      CALL trace_start("gradrot")

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ue = f_ue(ind)
         due = f_due(ind)
         DO l = ll_begin, ll_end
!DIR$ SIMD
            DO ij = ij_begin, ij_end
               due(ij + u_right, l) = ue(ij + u_right, l)
               due(ij + u_lup, l) = ue(ij + u_lup, l)
               due(ij + u_ldown, l) = ue(ij + u_ldown, l)
            END DO
         END DO
      END DO

      DO it = 1, nitergrot
         CALL send_message(f_due, req_due_gradrot)
         CALL wait_message(req_due_gradrot)
         DO ind = 1, ndomain
            IF (.NOT. assigned_domain(ind)) CYCLE
            CALL swap_dimensions(ind)
            CALL swap_geometry(ind)
            due = f_due(ind)
            CALL compute_gradrot(due, due, ll_begin, ll_end)
         END DO
      END DO

      CALL trace_end("gradrot")
   END SUBROUTINE gradrot

   SUBROUTINE divgrad(f_theta, f_dtheta)
      TYPE(t_field), POINTER :: f_theta(:)
      TYPE(t_field), POINTER :: f_dtheta(:)
      REAL(rstd), POINTER :: theta(:, :)
      REAL(rstd), POINTER :: dtheta(:, :)
      INTEGER :: ind, it

      CALL trace_start("divgrad")

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         theta = f_theta(ind)
         dtheta = f_dtheta(ind)
         dtheta = theta
      END DO

      DO it = 1, niterdivgrad
         CALL transfert_request(f_dtheta, req_i1)
         DO ind = 1, ndomain
            IF (.NOT. assigned_domain(ind)) CYCLE
            CALL swap_dimensions(ind)
            CALL swap_geometry(ind)
            dtheta = f_dtheta(ind)
            CALL compute_divgrad(dtheta, dtheta, ll_begin, ll_end)
         END DO
      END DO

      CALL trace_end("divgrad")
   END SUBROUTINE divgrad

   SUBROUTINE divgrad_theta_rhodz(f_mass, f_theta_rhodz, f_dtheta_rhodz)
      TYPE(t_field), POINTER :: f_mass(:)
      TYPE(t_field), POINTER :: f_theta_rhodz(:)
      TYPE(t_field), POINTER :: f_dtheta_rhodz(:)

      REAL(rstd), POINTER :: mass(:, :)
      REAL(rstd), POINTER :: theta_rhodz(:, :, :)
      REAL(rstd), POINTER :: dtheta_rhodz(:, :)

      INTEGER :: ind
      INTEGER :: it, l, ij

      CALL trace_start("divgrad")

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         mass = f_mass(ind)
         theta_rhodz = f_theta_rhodz(ind)
         dtheta_rhodz = f_dtheta_rhodz(ind)
         DO l = ll_begin, ll_end
!DIR$ SIMD
            DO ij = ij_begin, ij_end
               dtheta_rhodz(ij, l) = theta_rhodz(ij, l, 1)/mass(ij, l)
            END DO
         END DO
      END DO

      DO it = 1, niterdivgrad

         CALL send_message(f_dtheta_rhodz, req_dtheta)
         CALL wait_message(req_dtheta)
         DO ind = 1, ndomain
            IF (.NOT. assigned_domain(ind)) CYCLE
            CALL swap_dimensions(ind)
            CALL swap_geometry(ind)
            dtheta_rhodz = f_dtheta_rhodz(ind)
            CALL compute_divgrad(dtheta_rhodz, dtheta_rhodz, ll_begin, ll_end)
         END DO

      END DO

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         dtheta_rhodz = f_dtheta_rhodz(ind)
         mass = f_mass(ind)

         DO l = ll_begin, ll_end
!DIR$ SIMD
            DO ij = ij_begin, ij_end
               dtheta_rhodz(ij, l) = dtheta_rhodz(ij, l)*mass(ij, l)
            END DO
         END DO
      END DO

      CALL trace_end("divgrad")
   END SUBROUTINE divgrad_theta_rhodz

   SUBROUTINE compute_gradiv(ue, gradivu_e, llb, lle)
      INTEGER, INTENT(IN)     :: llb
      INTEGER, INTENT(IN)     :: lle
      REAL(rstd), INTENT(IN)  :: ue(iim*3*jjm, llm)
      REAL(rstd), INTENT(OUT) :: gradivu_e(iim*3*jjm, llm)
      REAL(rstd) :: divu_i(iim*jjm, llb:lle)

      INTEGER :: ij, l

      DO l = llb, lle
!DIR$ SIMD
         DO ij = ij_begin, ij_end
            divu_i(ij, l) = 1./Ai(ij)*(ne(ij, right)*ue(ij + u_right, l)*le(ij + u_right) + &
                                       ne(ij, rup)*ue(ij + u_rup, l)*le(ij + u_rup) + &
                                       ne(ij, lup)*ue(ij + u_lup, l)*le(ij + u_lup) + &
                                       ne(ij, left)*ue(ij + u_left, l)*le(ij + u_left) + &
                                       ne(ij, ldown)*ue(ij + u_ldown, l)*le(ij + u_ldown) + &
                                       ne(ij, rdown)*ue(ij + u_rdown, l)*le(ij + u_rdown))
         END DO
      END DO

      DO l = llb, lle
!DIR$ SIMD
         DO ij = ij_begin, ij_end
            gradivu_e(ij + u_right, l) = -1/de(ij + u_right)*(ne(ij, right)*divu_i(ij, l) + &
                                                              ne(ij + t_right, left)*divu_i(ij + t_right, l))
            gradivu_e(ij + u_lup, l) = -1/de(ij + u_lup)*(ne(ij, lup)*divu_i(ij, l) + &
                                                          ne(ij + t_lup, rdown)*divu_i(ij + t_lup, l))
            gradivu_e(ij + u_ldown, l) = -1/de(ij + u_ldown)*(ne(ij, ldown)*divu_i(ij, l) + &
                                                              ne(ij + t_ldown, rup)*divu_i(ij + t_ldown, l))
         END DO
      END DO

      DO l = llb, lle
!DIR$ SIMD
         DO ij = ij_begin, ij_end
            gradivu_e(ij + u_right, l) = -gradivu_e(ij + u_right, l)*cgraddiv
            gradivu_e(ij + u_lup, l) = -gradivu_e(ij + u_lup, l)*cgraddiv
            gradivu_e(ij + u_ldown, l) = -gradivu_e(ij + u_ldown, l)*cgraddiv
         END DO
      END DO

   END SUBROUTINE compute_gradiv

   SUBROUTINE compute_divgrad(theta, divgrad_i, llb, lle)
      INTEGER, INTENT(IN)     :: llb
      INTEGER, INTENT(IN)     :: lle
      REAL(rstd), INTENT(IN)  :: theta(iim*jjm, llm)
      REAL(rstd), INTENT(OUT) :: divgrad_i(iim*jjm, llm)
      REAL(rstd)  :: grad_e(3*iim*jjm, llb:lle)
      INTEGER :: ij, l
      DO l = llb, lle
!DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            grad_e(ij + u_right, l) = -1/de(ij + u_right)*(ne(ij, right)*theta(ij, l) + &
                                                           ne(ij + t_right, left)*theta(ij + t_right, l))
            grad_e(ij + u_lup, l) = -1/de(ij + u_lup)*(ne(ij, lup)*theta(ij, l) + &
                                                       ne(ij + t_lup, rdown)*theta(ij + t_lup, l))
            grad_e(ij + u_ldown, l) = -1/de(ij + u_ldown)*(ne(ij, ldown)*theta(ij, l) + &
                                                           ne(ij + t_ldown, rup)*theta(ij + t_ldown, l))
         END DO
      END DO

      DO l = llb, lle
!DIR$ SIMD
         DO ij = ij_begin, ij_end
            divgrad_i(ij, l) = 1./Ai(ij)*(ne(ij, right)*grad_e(ij + u_right, l)*le(ij + u_right) + &
                                          ne(ij, rup)*grad_e(ij + u_rup, l)*le(ij + u_rup) + &
                                          ne(ij, lup)*grad_e(ij + u_lup, l)*le(ij + u_lup) + &
                                          ne(ij, left)*grad_e(ij + u_left, l)*le(ij + u_left) + &
                                          ne(ij, ldown)*grad_e(ij + u_ldown, l)*le(ij + u_ldown) + &
                                          ne(ij, rdown)*grad_e(ij + u_rdown, l)*le(ij + u_rdown))
         END DO
      END DO

      DO l = llb, lle
         DO ij = ij_begin, ij_end
            divgrad_i(ij, l) = -divgrad_i(ij, l)*cdivgrad
         END DO
      END DO

   END SUBROUTINE compute_divgrad

   SUBROUTINE compute_gradrot(ue, gradrot_e, llb, lle)
      INTEGER, INTENT(IN)     :: llb
      INTEGER, INTENT(IN)     :: lle
      REAL(rstd), INTENT(IN)  :: ue(iim*3*jjm, llm)
      REAL(rstd), INTENT(OUT) :: gradrot_e(iim*3*jjm, llm)
      REAL(rstd) :: rot_v(2*iim*jjm, llb:lle)

      INTEGER :: ij, l

      DO l = llb, lle
!DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext

            rot_v(ij + z_up, l) = 1./Av(ij + z_up)* &
                                  (ne(ij, rup)*ue(ij + u_rup, l)*de(ij + u_rup) + &
                                   ne(ij + t_rup, left)*ue(ij + t_rup + u_left, l)*de(ij + t_rup + u_left) - &
                                   ne(ij, lup)*ue(ij + u_lup, l)*de(ij + u_lup))

            rot_v(ij + z_down, l) = 1./Av(ij + z_down)* &
                                    (ne(ij, ldown)*ue(ij + u_ldown, l)*de(ij + u_ldown) + &
                                     ne(ij + t_ldown, right)*ue(ij + t_ldown + u_right, l)*de(ij + t_ldown + u_right) - &
                                     ne(ij, rdown)*ue(ij + u_rdown, l)*de(ij + u_rdown))

         END DO
      END DO

      DO l = llb, lle
!DIR$ SIMD
         DO ij = ij_begin, ij_end
            gradrot_e(ij + u_right, l) = 1/le(ij + u_right)*ne(ij, right)*(rot_v(ij + z_rdown, l) - rot_v(ij + z_rup, l))
            gradrot_e(ij + u_lup, l) = 1/le(ij + u_lup)*ne(ij, lup)*(rot_v(ij + z_up, l) - rot_v(ij + z_lup, l))
            gradrot_e(ij + u_ldown, l) = 1/le(ij + u_ldown)*ne(ij, ldown)*(rot_v(ij + z_ldown, l) - rot_v(ij + z_down, l))
         END DO
      END DO

      DO l = llb, lle
!DIR$ SIMD
         DO ij = ij_begin, ij_end
            gradrot_e(ij + u_right, l) = -gradrot_e(ij + u_right, l)*cgradrot
            gradrot_e(ij + u_lup, l) = -gradrot_e(ij + u_lup, l)*cgradrot
            gradrot_e(ij + u_ldown, l) = -gradrot_e(ij + u_ldown, l)*cgradrot
         END DO
      END DO

   END SUBROUTINE compute_gradrot

!----------------------- Utility routines ------------------

   SUBROUTINE global_max(dumax)
      USE mpi_mod
      USE mpipara
      USE transfert_omp_mod
      REAL(rstd) :: dumax, dumax1
      IF (using_mpi) THEN
         CALL reduce_max_omp(dumax, dumax1)
         !$OMP MASTER
         CALL MPI_ALLREDUCE(dumax1, dumax, 1, MPI_REAL8, MPI_MAX, comm_icosa, ierr)
         !$OMP END MASTER
         CALL bcast_omp(dumax)
      ELSE
         CALL allreduce_max_omp(dumax, dumax1)
         dumax = dumax1
      END IF
   END SUBROUTINE global_max

   FUNCTION max_vel(f_du)
      TYPE(t_field)       :: f_du(:)
      REAL(rstd), POINTER  :: du(:)
      INTEGER :: ind, i, j, ij
      REAL(rstd) :: max_vel, dumax

      dumax = 0.
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         du = f_du(ind)

         DO j = jj_begin, jj_end
            DO i = ii_begin, ii_end
               ij = (j - 1)*iim + i
               if (le(ij + u_right) > 1e-100) dumax = MAX(dumax, ABS(du(ij + u_right)))
               if (le(ij + u_lup) > 1e-100) dumax = MAX(dumax, ABS(du(ij + u_lup)))
               if (le(ij + u_ldown) > 1e-100) dumax = MAX(dumax, ABS(du(ij + u_ldown)))
            END DO
         END DO
      END DO
      CALL global_max(dumax)
      max_vel = dumax
   END FUNCTION max_vel

   FUNCTION max_scalar(f_dtheta)
      TYPE(t_field)       :: f_dtheta(:)
      REAL(rstd), POINTER  :: dtheta(:)
      INTEGER :: ind, i, j, ij
      REAL(rstd) :: max_scalar, dthetamax
      dthetamax = 0.
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
         CALL swap_dimensions(ind)
         dtheta = f_dtheta(ind)
         DO j = jj_begin, jj_end
            DO i = ii_begin, ii_end
               ij = (j - 1)*iim + i
               dthetamax = MAX(dthetamax, ABS(dtheta(ij)))
            END DO
         END DO
      END DO
      CALL global_max(dthetamax)
      max_scalar = dthetamax
   END FUNCTION max_scalar

   SUBROUTINE random_vel(f_u)
      TYPE(t_field)      :: f_u(:)
      REAL(rstd), POINTER :: u(:)
      INTEGER :: M, ind, i, j, ij
      REAL(rstd) :: r

      !$OMP BARRIER
      !$OMP MASTER
      DO ind = 1, ndomain
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         u = f_u(ind)

         ! set random seed to get reproductibility when using a different number of process
         CALL RANDOM_SEED(size=M)
         CALL RANDOM_SEED(put=(/(i, i=1, M)/))

         DO j = jj_begin, jj_end
            DO i = ii_begin, ii_end
               ij = (j - 1)*iim + i
               CALL RANDOM_NUMBER(r)
               u(ij + u_right) = r - 0.5
               CALL RANDOM_NUMBER(r)
               u(ij + u_lup) = r - 0.5
               CALL RANDOM_NUMBER(r)
               u(ij + u_ldown) = r - 0.5
            END DO
         END DO
      END DO
      !$OMP END MASTER
      !$OMP BARRIER

   END SUBROUTINE random_vel

   SUBROUTINE random_scalar(f_theta)
      TYPE(t_field)      :: f_theta(:)
      REAL(rstd), POINTER :: theta(:)
      INTEGER :: M, ind, i, j, ij
      REAL(rstd) :: r
      !$OMP BARRIER
      !$OMP MASTER
      DO ind = 1, ndomain
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         theta = f_theta(ind)
         ! set random seed to get reproductibility when using a different number of process
         CALL RANDOM_SEED(size=M)
         CALL RANDOM_SEED(put=(/(i, i=1, M)/))

         DO j = jj_begin, jj_end
            DO i = ii_begin, ii_end
               ij = (j - 1)*iim + i
               CALL RANDOM_NUMBER(r)
               theta(ij) = r - 0.5
            END DO
         END DO
      END DO
      !$OMP END MASTER
      !$OMP BARRIER
   END SUBROUTINE random_scalar

   SUBROUTINE rescale_field(scal, f_du, f_u)
      TYPE(t_field)      :: f_du(:), f_u(:)
      REAL(rstd), POINTER :: du(:), u(:)
      REAL(rstd) :: scal
      INTEGER :: ind
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
         CALL swap_dimensions(ind)
         u = f_u(ind)
         du = f_du(ind)
         u = scal*du
      END DO
   END SUBROUTINE rescale_field

END MODULE dissip_gcm_mod
