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
  !$OMP THREADPRIVATE(rdamp)
  REAL,ALLOCATABLE,SAVE:: lambda(:) ! inverse or quenching time scale (Hz)
  !$OMP THREADPRIVATE(lambda)

  PUBLIC sponge, init_sponge, iflag_sponge, pre_init_sponge

  INTEGER,SAVE :: ni_glo
  !$OMP THREADPRIVATE(ni_glo)
  INTEGER,SAVE :: nj_glo
  !$OMP THREADPRIVATE(nj_glo)
  INTEGER,SAVE :: ni
  !$OMP THREADPRIVATE(ni)
  INTEGER,SAVE :: nj
  !$OMP THREADPRIVATE(nj)
  INTEGER,SAVE :: ibegin
  !$OMP THREADPRIVATE(ibegin)
  INTEGER,SAVE :: jbegin
  !$OMP THREADPRIVATE(jbegin)
 
  TYPE(t_field),POINTER,SAVE :: f_hflux(:)
  TYPE(t_field),POINTER,SAVE :: f_ulon(:)
  TYPE(t_field),POINTER,SAVE :: f_ulat(:)
  TYPE(t_field),POINTER,SAVE :: f_ulon_sponge(:)
  TYPE(t_field),POINTER,SAVE :: f_ulat_sponge(:)
  TYPE(t_field),POINTER,SAVE :: f_rhodz_sponge(:)
  TYPE(t_field),POINTER,SAVE :: f_theta_sponge(:)

  TYPE(t_field),POINTER,SAVE :: f_dulon(:)
  TYPE(t_field),POINTER,SAVE :: f_dulat(:)
 
  INTEGER,SAVE  :: n_sponge
  !$OMP THREADPRIVATE(n_sponge)
  INTEGER,SAVE  :: begin_sponge
  !$OMP THREADPRIVATE(begin_sponge)
  
CONTAINS

  SUBROUTINE pre_init_sponge
  USE icosa
  USE disvert_mod
  IMPLICIT NONE
    INTEGER :: l

    tau_sponge = 1.e-5
    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
    
    IF (.NOT. (iflag_sponge == 1 .OR. iflag_sponge == 2 .OR. iflag_sponge == 3)) THEN
      PRINT*,'Bad selector for variable iflag_sponge : <',iflag_sponge,'> options are 0,1,2'
      STOP
    ELSE
      IF (.NOT. (mode_sponge == 1 .OR. mode_sponge == 2 .OR. mode_sponge == 3)) THEN
        PRINT*,'Bad selector for variable mode_sponge : <',mode_sponge,'> options are 0,1,2'
        STOP
      ENDIF
    ENDIF


    IF (iflag_sponge == 0) RETURN
    
    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.
        n_sponge=4
        begin_sponge=llm-3
    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.)
        DO l=llm,1,-1
          IF (lambda(l)/=0.) begin_sponge=l
        ENDDO
        n_sponge=llm-begin_sponge+1
    ENDIF
     
    IF (iflag_sponge>=1 .AND. (mode_sponge==2 .OR. mode_sponge==3)) THEN
     CALL pre_init_sponge_zonal_mean_u
    ENDIF
    
  END SUBROUTINE pre_init_sponge
 
 
  SUBROUTINE init_sponge
  USE icosa
  USE disvert_mod
  USE omp_para
  USE mpipara, ONLY: is_mpi_master
  IMPLICIT NONE
    INTEGER :: l
  
    IF (iflag_sponge == 0) RETURN

    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
    !$acc update device(rdamp(:), lambda(:)) async
 
   IF (iflag_sponge>=1 .AND. (mode_sponge == 2 .OR. mode_sponge == 3)) THEN
     CALL init_sponge_zonal_mean_u
   ENDIF
   
  END SUBROUTINE init_sponge

  
  SUBROUTINE sponge(f_mass, 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_mass(:)
    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(:,:), dtheta_rhodz)  async
        DO l=ll_begin,ll_end
         !$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)

            dtheta_rhodz(ij,l,:) = 0.0
          ENDDO
        ENDDO
       !$acc end parallel loop
      END DO
    ELSE IF (mode_sponge == 2 .OR. mode_sponge == 3) THEN
      CALL sponge_zonal_mean_u(mode_sponge, f_mass, f_ue,f_due,f_theta_rhodz,f_dtheta_rhodz) 
    ELSE
      PRINT*,'Arnaud Bad selector for variable mode_sponge : <',mode_sponge,             &
            '> options 2 and 3 not available for the moment!'
      STOP
    ENDIF

    DO ind=1,ndomain
      IF (.NOT. assigned_domain(ind)) CYCLE
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      due=f_due(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)*due(ij+u_right,l)
          due(ij+u_lup,l)          = -rdamp(l)*due(ij+u_lup,l)
          due(ij+u_ldown,l)        = -rdamp(l)*due(ij+u_ldown,l)
          dtheta_rhodz(ij,l,:)     = -rdamp(l)*dtheta_rhodz(ij,l,:)
        ENDDO
      ENDDO
     !$acc end parallel loop
    END DO

   CALL trace_end("sponge")

!$OMP BARRIER
  END SUBROUTINE sponge

  SUBROUTINE pre_init_sponge_zonal_mean_u()
  USE xios_mod
  USE grid_param, ONLY : iim_glo
  USE disvert_mod
  IMPLICIT NONE
  INTEGER :: nlon
  INTEGER :: nlat  

!$OMP MASTER
    nlon=360.*iim_glo/80.
    nlat=180.*iim_glo/80.

    CALL xios_set_axis_attr("lev_sponge", n_glo=n_sponge, value=presnivs(begin_sponge:))
    CALL xios_set_domain_attr("to_sponge", ni_glo=nlon, nj_glo=nlat)
    CALL xios_set_axis_attr("reduced_sponge", n_glo=nlat)
    CALL xios_set_field_attr("ulon_sponge_reduced", read_access=.TRUE.)
    CALL xios_set_field_attr("ulon_from_sponge", read_access=.TRUE.)
    CALL xios_set_field_attr("ulat_sponge_reduced", read_access=.TRUE.)
    CALL xios_set_field_attr("ulat_from_sponge", read_access=.TRUE.)
    CALL xios_set_field_attr("rhodz_sponge_reduced", read_access=.TRUE.)
    IF (mode_sponge==3) CALL xios_set_field_attr("theta_sponge_reduced", read_access=.TRUE.)
    IF (mode_sponge==3) CALL xios_set_field_attr("theta_from_sponge", read_access=.TRUE.)
!$OMP END MASTER

  END SUBROUTINE  pre_init_sponge_zonal_mean_u


  SUBROUTINE init_sponge_zonal_mean_u()
  USE xios_mod
  IMPLICIT NONE
    
!$OMP MASTER
    CALL xios_get_domain_attr("to_sponge", ni_glo=ni_glo, nj_glo=nj_glo,ni=ni, nj=nj, ibegin=ibegin,jbegin=jbegin)
!$OMP END MASTER
    CALL allocate_field(f_hflux, field_u, type_real, llm, name='hflux') 
    CALL allocate_field(f_ulon, field_t, type_real, llm, name='ulon') 
    CALL allocate_field(f_ulat, field_t, type_real, llm, name='ulat') 
    CALL allocate_field(f_ulon_sponge, field_t, type_real, n_sponge, name='ulon_sponge') 
    CALL allocate_field(f_ulat_sponge, field_t, type_real, n_sponge, name='ulat_sponge') 
    CALL allocate_field(f_rhodz_sponge, field_t, type_real, n_sponge, name='rhodz_sponge') 
    CALL allocate_field(f_theta_sponge, field_t, type_real, n_sponge, name='theta_sponge') 
    CALL allocate_field(f_dulon, field_t, type_real, llm,name='dulon_sponge') 
    CALL allocate_field(f_dulat, field_t, type_real, llm,name='dulat_sponge') 

  END SUBROUTINE init_sponge_zonal_mean_u

  
  SUBROUTINE sponge_zonal_mean_u(mode_sponge, f_rhodz, 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
  USE wind_mod
  USE xios_mod
  IMPLICIT NONE
    INTEGER,INTENT(IN)    :: mode_sponge
    TYPE(t_field),POINTER :: f_rhodz(:)
    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         :: rhodz(:,:)
    REAL(rstd),POINTER         :: hflux(:,:)
    REAL(rstd),POINTER         :: theta_rhodz(:,:,:)
    REAL(rstd),POINTER         :: dtheta_rhodz(:,:,:)!,dtheta_sponge(:,:)
    REAL(rstd),POINTER         :: ulon(:,:)
    REAL(rstd),POINTER         :: ulat(:,:)
    REAL(rstd),POINTER         :: ulon_sponge(:,:)
    REAL(rstd),POINTER         :: ulat_sponge(:,:)
    REAL(rstd),POINTER         :: theta_sponge(:,:)
    REAL(rstd),POINTER         :: rhodz_sponge(:,:)
    
    REAL(rstd) :: ulon_zonal_mean(nj_glo,n_sponge)
    REAL(rstd) :: ulat_zonal_mean(nj_glo,n_sponge)
    REAL(rstd) :: theta_zonal_mean(nj_glo,n_sponge)
    REAL(rstd) :: rhodz_zonal_mean(nj_glo,n_sponge)
    REAL(rstd) :: domain_zonal_mean(ni,nj,n_sponge)
   
    INTEGER :: i,j,ind,ij,l,ll 

    DO ind = 1 , ndomain
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      ue=f_ue(ind)
      rhodz=f_rhodz(ind)
      hflux = f_hflux(ind)
         
      DO l = ll_begin, ll_end
         !DIR$ IVDEP
         DO ij=ij_begin_ext,ij_end_ext         
            hflux(ij+u_right,l)=0.5*(rhodz(ij,l)+rhodz(ij+t_right,l))*ue(ij+u_right,l)
            hflux(ij+u_lup,l)=0.5*(rhodz(ij,l)+rhodz(ij+t_lup,l))*ue(ij+u_lup,l)
            hflux(ij+u_ldown,l)=0.5*(rhodz(ij,l)+rhodz(ij+t_ldown,l))*ue(ij+u_ldown,l)
         ENDDO
       ENDDO
    ENDDO

              
    CALL transfert_request(f_hflux,req_e1_vect)
    CALL un2ulonlat(f_hflux, f_ulon, f_ulat)

    DO ind = 1 , ndomain
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      
      ulon=f_ulon(ind)
      ulat=f_ulat(ind)
      ulon_sponge = f_ulon_sponge(ind)
      ulat_sponge = f_ulat_sponge(ind)
      rhodz_sponge = f_rhodz_sponge(ind)
      theta_sponge = f_theta_sponge(ind)
      theta_rhodz = f_theta_rhodz(ind)
              
      DO l = 1, llm
         IF (l<begin_sponge) CYCLE
         ll=l-begin_sponge+1
         !DIR$ IVDEP
         DO ij=ij_begin_ext,ij_end_ext         
            ulon_sponge(ij,ll)= ulon(ij,l) 
            ulat_sponge(ij,ll)= ulat(ij,l) 
            rhodz_sponge(ij,ll) = rhodz(ij,l)
            IF (mode_sponge==3) theta_sponge(ij,ll) = theta_rhodz(ij,l,1)
         ENDDO
       ENDDO
    ENDDO

!$OMP BARRIER
    
    CALL xios_write_field("ulon_sponge",f_ulon_sponge)
    CALL xios_write_field("ulat_sponge",f_ulat_sponge)
    CALL xios_write_field("rhodz_sponge",f_rhodz_sponge)
!$OMP MASTER
    CALL xios_recv_field("ulon_sponge_reduced",ulon_zonal_mean)
    CALL xios_recv_field("ulat_sponge_reduced",ulat_zonal_mean)
    CALL xios_recv_field("rhodz_sponge_reduced",rhodz_zonal_mean)
!$OMP END MASTER

    IF (mode_sponge==3) THEN
      CALL xios_write_field("theta_sponge",f_theta_sponge)
!$OMP MASTER
      CALL xios_recv_field("theta_sponge_reduced",theta_zonal_mean)
!$OMP END MASTER
    ENDIF
    
!$OMP MASTER
    DO j=1,nj
      DO i=1,ni
        domain_zonal_mean(i,j,:)=ulon_zonal_mean(jbegin+j,:)/rhodz_zonal_mean(jbegin+j,:)
      ENDDO
    ENDDO
    
    CALL xios_send_field("ulon_sponge_mean",domain_zonal_mean)
!$OMP END MASTER
    CALL xios_read_field("ulon_from_sponge",f_ulon_sponge)
  
!$OMP MASTER
    DO j=1,nj
      DO i=1,ni
        domain_zonal_mean(i,j,:)=ulat_zonal_mean(jbegin+j,:)/rhodz_zonal_mean(jbegin+j,:)
      ENDDO
    ENDDO
    
    CALL xios_send_field("ulat_sponge_mean",domain_zonal_mean)
!$OMP END MASTER
    CALL xios_read_field("ulat_from_sponge",f_ulat_sponge)
    
    IF (mode_sponge==3) THEN
!$OMP MASTER
      DO j=1,nj
        DO i=1,ni
          domain_zonal_mean(i,j,:)=theta_zonal_mean(jbegin+j,:)
        ENDDO
      ENDDO
      CALL xios_send_field("theta_sponge_mean",domain_zonal_mean)
!$OMP END MASTER
      CALL xios_read_field("theta_from_sponge",f_theta_sponge)
    ENDIF
    

    DO ind = 1 , ndomain
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      
      ulon=f_ulon(ind)
      ulat=f_ulat(ind)
      ulon_sponge = f_ulon_sponge(ind)
      ulat_sponge = f_ulat_sponge(ind)
      theta_sponge = f_theta_sponge(ind)
      dtheta_rhodz = f_dtheta_rhodz(ind)    
      theta_rhodz = f_theta_rhodz(ind)
      DO l = ll_begin, ll_end
        IF (l<begin_sponge) CYCLE
        ll=l-begin_sponge+1
        IF (mode_sponge==2) THEN
          DO ij=ij_begin_ext,ij_end_ext         
            ulon(ij,l) = ulon_sponge(ij,ll) 
            ulat(ij,l) = ulat_sponge(ij,ll) 
            dtheta_rhodz(ij,l,1)=theta_rhodz(ij,ll,1) 
          ENDDO
        ELSE  IF (mode_sponge==3) THEN
          DO ij=ij_begin_ext,ij_end_ext         
            ulon(ij,l) = ulon_sponge(ij,ll) 
            ulat(ij,l) = ulat_sponge(ij,ll) 
            dtheta_rhodz(ij,l,1) = theta_sponge(ij,ll)
          ENDDO
        ENDIF
      ENDDO
    ENDDO
    
    CALL transfert_request(f_ulon,req_i0)
    CALL transfert_request(f_ulat,req_i0)
    IF (mode_sponge==3) THEN
      CALL transfert_request(f_dtheta_rhodz,req_i0)
    ENDIF 
    CALL ulonlat2un(f_ulon,f_ulat,f_due)
!$OMP BARRIER
   
    DO ind = 1 , ndomain
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      due=f_due(ind)
      ue=f_ue(ind)
      dtheta_rhodz=f_dtheta_rhodz(ind)
      theta_rhodz=f_theta_rhodz(ind)
 
       DO l = ll_begin, ll_end
         IF (l<begin_sponge) THEN
           !DIR$ IVDEP
           DO ij=ij_begin,ij_end         
              due(ij+u_right,l) = 0.
              due(ij+u_lup,l)   = 0.
              due(ij+u_ldown,l) = 0.
              dtheta_rhodz(ij,l,:) = 0.
           ENDDO         
         ELSE
           !DIR$ IVDEP
           DO ij=ij_begin,ij_end         
             due(ij+u_right,l) = ue(ij+u_right,l) - due(ij+u_right,l)
             due(ij+u_lup,l)   = ue(ij+u_lup,l)   - due(ij+u_lup,l)
             due(ij+u_ldown,l) = ue(ij+u_ldown,l) - due(ij+u_ldown,l)
             dtheta_rhodz(ij,l,1) = theta_rhodz(ij,l,1) - dtheta_rhodz(ij,l,1)
          ENDDO
        ENDIF
      ENDDO
    ENDDO  
   
    CALL transfert_request(f_due,req_e1_vect)
!ym to check results
!    CALL un2ulonlat(f_due, f_dulon, f_dulat)
!    CALL xios_write_field("dulon_sponge",f_dulon)
!    CALL xios_write_field("dulat_sponge",f_dulat)
!    CALL xios_write_field("dtheta_sponge",f_dtheta_rhodz)
    
  END SUBROUTINE sponge_zonal_mean_u
    

END MODULE sponge_mod