MODULE sponge_mod
   USE icosa
   USE sponge_zonal_mod

   PRIVATE

   INTEGER, SAVE :: itau_sponge  ! Iterations frequency between sponge call
   REAL, SAVE :: tau_sponge !inverse of charactericstic relaxation time scale at the topmost layer (Hz)
   INTEGER, SAVE :: itau_sponge_zonal_update ! number of sponge iterations between two updates of the zonal mean

   INTEGER, SAVE :: iflag_sponge !0 --> for no sponge
   !1 --> for sponge over 4 topmost layers
   !2 --> for sponge from top to ~1% of top layer pressure
   INTEGER, SAVE :: mode_sponge !1 --> u and v relax towards 0
   !2 --> u and v relax towards their zonal mean
   !3 --> u,v and pot. temp. relax towards their zonal mean
   !4 --> u and v relax towards their exponential moving average
   !5 --> u and v and pot. temp. relax towards their exponential moving average
   REAL, SAVE    :: sponge_alpha  ! parameter alpha for exponential moving average :

   !$OMP THREADPRIVATE(tau_sponge,iflag_sponge,mode_sponge,itau_sponge,itau_sponge_zonal_update)
   REAL, ALLOCATABLE, SAVE :: rdamp(:) ! quenching coefficient
   !$OMP THREADPRIVATE(rdamp)
   REAL, ALLOCATABLE, SAVE:: lambda(:) ! inverse or quenching time scale (Hz)
   !$OMP THREADPRIVATE(lambda)

   PUBLIC sponge, init_sponge, iflag_sponge

CONTAINS

   SUBROUTINE init_sponge
      USE icosa
      USE disvert_mod
      USE omp_para
      USE mpipara, ONLY: is_mpi_master
      IMPLICIT NONE
      INTEGER :: l

      tau_sponge = 1.e-4
      CALL getin("tau_sponge", tau_sponge)
      PRINT *, 'tau_sponge = ', tau_sponge

      iflag_sponge = 0
      CALL getin("iflag_sponge", iflag_sponge)
      PRINT *, 'iflag_sponge = ', iflag_sponge

      mode_sponge = 1
      CALL getin("mode_sponge", mode_sponge)
      PRINT *, 'mode_sponge = ', mode_sponge

      IF (iflag_sponge == 0) THEN
         PRINT *, 'init_sponge: no sponge'
         RETURN
      END IF

!$OMP MASTER
      ALLOCATE (rdamp(llm))
      ALLOCATE (lambda(llm))
      !$acc enter data create(rdamp(:),lambda(:)) async

      IF (iflag_sponge == 1) THEN
! sponge quenching over the topmost 4 atmospheric layers
         lambda(:) = 0.
         lambda(llm) = tau_sponge
         lambda(llm - 1) = tau_sponge/2.
         lambda(llm - 2) = tau_sponge/4.
         lambda(llm - 3) = tau_sponge/8.
      ELSE IF (iflag_sponge == 2) THEN
! sponge quenching over topmost layers down to pressures which are
! higher than 100 times the topmost layer pressure
         lambda(:) = tau_sponge*max(presnivs(llm)/presnivs(:) - 0.01, 0.)
      ELSE
         PRINT *, 'Bad selector for variable iflag_sponge : <', iflag_sponge, &
            '> options are 0,1,2'
         STOP
      END IF

! quenching coefficient rdamp(:)
!         rdamp(:)=dt*lambda(:) ! Explicit Euler approx.
!      rdamp(:)=1.-exp(-lambda(:)*dt)
      rdamp(:) = itau_dissip*lambda(:)

      IF (is_mpi_master) THEN
         PRINT *, 'lambda = '
         DO l = ll_begin, ll_end
            PRINT *, l, lambda(l)
         END DO
         PRINT *, 'rdamp = '
         DO l = ll_begin, ll_end
            PRINT *, l, rdamp(l)
         END DO
      END IF
!$OMP END MASTER
!$OMP BARRIER
      !$acc update device(rdamp(:), lambda(:)) async

   END SUBROUTINE init_sponge

   SUBROUTINE sponge(f_ue, f_due, f_theta_rhodz, f_dtheta_rhodz)
      USE icosa
      USE theta2theta_rhodz_mod
      USE exner_mod
      USE geopotential_mod
      USE trace
      USE time_mod
      USE omp_para
      IMPLICIT NONE
      TYPE(t_field), POINTER :: f_ue(:)
      TYPE(t_field), POINTER :: f_due(:)
      TYPE(t_field), POINTER :: f_theta_rhodz(:)
      TYPE(t_field), POINTER :: f_dtheta_rhodz(:)

      REAL(rstd), POINTER         :: due(:, :)!,due_sponge(:,:)
      REAL(rstd), POINTER         :: ue(:, :)
      REAL(rstd), POINTER         :: theta_rhodz(:, :, :)
      REAL(rstd), POINTER         :: dtheta_rhodz(:, :, :)!,dtheta_sponge(:,:)

      INTEGER :: ind
      INTEGER :: l, ij

!$OMP BARRIER

      CALL trace_start("sponge")

      IF (mode_sponge == 1) THEN
         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)
            theta_rhodz = f_theta_rhodz(ind)
            dtheta_rhodz = f_dtheta_rhodz(ind)

            !$acc parallel loop collapse(2) present(ue(:,:), due(:,:), rdamp(:), dtheta_rhodz)  async
            DO l = ll_begin, ll_end
!$SIMD
               DO ij = ij_begin, ij_end

                  due(ij + u_right, l) = -rdamp(l)*ue(ij + u_right, l)
                  due(ij + u_lup, l) = -rdamp(l)*ue(ij + u_lup, l)
                  due(ij + u_ldown, l) = -rdamp(l)*ue(ij + u_ldown, l)

                  dtheta_rhodz(ij, l, :) = 0.0
               END DO
            END DO
            !$acc end parallel loop
         END DO
      ELSE
         PRINT *, 'Bad selector for variable mode_sponge : <', mode_sponge, &
            '> options 2 and 3 not available for the moment!'
         STOP
      END IF

      CALL trace_end("sponge")

!$OMP BARRIER
   END SUBROUTINE sponge

END MODULE sponge_mod
