MODULE sponge_mod
  USE icosa

  PRIVATE
  
  REAL,SAVE :: tau_sponge !inverse of charactericstic relaxation time scale at the topmost layer (Hz)
  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
!$OMP THREADPRIVATE(tau_sponge,iflag_sponge,mode_sponge)
  REAL,ALLOCATABLE,SAVE :: rdamp(:) ! quenching coefficient
  REAL,ALLOCATABLE,SAVE:: lambda(:) ! inverse or quenching time scale (Hz)

  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
    ENDIF

!$OMP MASTER
    ALLOCATE(rdamp(llm))
    ALLOCATE(lambda(llm))

    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
    ENDIF

! 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)
      ENDDO
      PRINT*,'rdamp = '
      DO l=ll_begin,ll_end
        PRINT*,l,rdamp(l)
      ENDDO
    ENDIF
!$OMP END MASTER
!$OMP BARRIER

  END SUBROUTINE init_sponge

  SUBROUTINE sponge(f_ue,f_due,f_theta_rhodz,f_dtheta_rhodz)
  USE icosa
  USE theta2theta_rhodz_mod
  USE pression_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)

        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
          ENDDO
        ENDDO
      END DO
    ELSE
      PRINT*,'Bad selector for variable mode_sponge : <',mode_sponge,             &
            '> options 2 and 3 not available for the moment!'
      STOP
    ENDIF

   CALL trace_end("sponge")

!$OMP BARRIER
  END SUBROUTINE sponge

END MODULE sponge_mod