MODULE sponge_zonal_mod
   USE icosa

   TYPE(t_field), POINTER, SAVE :: f_ue_mean(:)
   TYPE(t_field), POINTER, SAVE :: f_rhodz_mean(:)
   TYPE(t_field), POINTER, SAVE :: f_theta_rhodz_mean(:)
   INTEGER, SAVE               :: nb_it_mean ! number of iteration used to build mean
   !$OMP THREADPRIVATE(nb_it_mean)
   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_ue_sponge_zonal(:)
   TYPE(t_field), POINTER, SAVE :: f_theta_sponge_zonal(:)

   TYPE(t_field), POINTER, SAVE :: f_dulon(:)
   TYPE(t_field), POINTER, SAVE :: f_dulat(:)

   LOGICAL, SAVE               :: is_read_sponge
   !$OMP THREADPRIVATE(is_read_sponge)
   INTEGER, SAVE  :: n_sponge
   !$OMP THREADPRIVATE(n_sponge)
   INTEGER, SAVE  :: begin_sponge
   !$OMP THREADPRIVATE(begin_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)

CONTAINS

   SUBROUTINE pre_init_sponge_zonal_mean(mode_sponge)
      USE xios_mod
      USE grid_param, ONLY: iim_glo
      USE disvert_mod
      USE restart_mod
      USE transfert_mod
      IMPLICIT NONE
      INTEGER, INTENT(IN) :: mode_sponge
      INTEGER :: nlon
      INTEGER :: nlat

      LOGICAL :: ok_sponge_zonal

      ok_sponge_zonal = restart_var_present("sponge_nb_it")
      is_read_sponge = .FALSE.

!$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.)
      ! for restart
      CALL xios_set_fieldgroup_attr("sponge_restart", enabled=.TRUE.)
      CALL xios_set_field_attr("sponge_rhodz_mean_restart", enabled=.TRUE.)
      IF (mode_sponge == 3) CALL xios_set_field_attr("sponge_theta_mean_restart", enabled=.TRUE.)
      IF (mode_sponge == 3) CALL xios_set_field_attr("sponge_theta_zonal_restart", enabled=.TRUE.)

      IF (is_write_start) THEN
         CALL xios_set_fieldgroup_attr("sponge_start0", enabled=.TRUE.)
         CALL xios_set_field_attr("sponge_rhodz_mean_start0", enabled=.TRUE.)
         IF (mode_sponge == 3) CALL xios_set_field_attr("sponge_theta_mean_start0", enabled=.TRUE.)
         IF (mode_sponge == 3) CALL xios_set_field_attr("sponge_theta_zonal_start0", enabled=.TRUE.)
      END IF

      IF (is_read_start .AND. .NOT. is_read_interp) THEN
         IF (ok_sponge_zonal) THEN
            CALL xios_set_fieldgroup_attr("sponge_start", enabled=.TRUE.)
            CALL xios_set_field_attr("sponge_rhodz_mean_start", enabled=.TRUE.)
            IF (mode_sponge == 3) CALL xios_set_field_attr("sponge_theta_mean_start", enabled=.TRUE.)
            IF (mode_sponge == 3) CALL xios_set_field_attr("sponge_theta_zonal_start", enabled=.TRUE.)
            is_read_sponge = .TRUE.
         END IF
      END IF
!$OMP END MASTER
      CALL bcast(is_read_sponge)

   END SUBROUTINE pre_init_sponge_zonal_mean

   SUBROUTINE init_sponge_zonal_mean()
      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_ue_mean, field_u, type_real, llm, name='sponge_ue_mean')
      CALL allocate_field(f_rhodz_mean, field_t, type_real, llm, name='sponge_rhodz_mean')
      CALL allocate_field(f_theta_rhodz_mean, field_t, type_real, llm, name='sponge_theta_mean')
      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='sponge_rhodz')
      CALL allocate_field(f_theta_sponge, field_t, type_real, n_sponge, name='sponge_theta')

      CALL allocate_field(f_ue_sponge_zonal, field_u, type_real, llm, name='sponge_ue_zonal')
      CALL allocate_field(f_theta_sponge_zonal, field_t, type_real, llm, name='sponge_theta_zonal')

      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

  SUBROUTINE sponge_zonal_mean(it, mode_sponge, itau_sponge, itau_sponge_zonal_update, 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_to_lonlat_mod
      USE xios_mod
      USE checksum_mod
      IMPLICIT NONE
      INTEGER, INTENT(IN)   :: it
      INTEGER, INTENT(IN)   :: mode_sponge
      INTEGER, INTENT(IN)   :: itau_sponge
      INTEGER, INTENT(IN)   :: itau_sponge_zonal_update
      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(:)

      INTEGER :: it_sponge

      it_sponge = (it - 1)/itau_sponge
      if (it_sponge == 0) CALL sponge_reset_mean(mode_sponge)
      CALL sponge_update_mean(mode_sponge, f_rhodz, f_ue, f_theta_rhodz)
!    IF (is_master) PRINT *, "calling sponge at it=",it," sponge iteration",it_sponge
!    IF (is_master) PRINT *, "nb_it_mean=",nb_it_mean
!    CALL checksum(f_ue_mean)
!    CALL checksum(f_rhodz_mean)
!    CALL checksum(f_theta_rhodz_mean)

      IF (MOD(it_sponge, itau_sponge_zonal_update) == 0) THEN
         CALL sponge_update_zonal_mean(mode_sponge)
         CALL sponge_reset_mean(mode_sponge)
      END IF
!    CALL checksum(f_ue_sponge_zonal)
!    CALL checksum(f_theta_sponge_zonal)
      CALL sponge_apply_zonal_mean(mode_sponge, f_ue, f_due, f_theta_rhodz, f_dtheta_rhodz)

   END SUBROUTINE sponge_zonal_mean

   SUBROUTINE sponge_update_zonal_mean(mode_sponge)
      USE icosa
      USE theta2theta_rhodz_mod
      USE exner_mod
      USE geopotential_mod
      USE trace
      USE time_mod
      USE omp_para
      USE wind_from_lonlat_mod
      USE wind_to_lonlat_mod
      USE xios_mod
      IMPLICIT NONE
      INTEGER, INTENT(IN)    :: mode_sponge

      REAL(rstd), POINTER         :: ue(:, :)
      REAL(rstd), POINTER         :: rhodz(:, :)
      REAL(rstd), POINTER         :: hflux(:, :)
      REAL(rstd), POINTER         :: theta_rhodz(:, :)
      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), POINTER         :: theta_sponge_zonal(:, :)

      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

      CALL transfert_request(f_rhodz_mean, req_i1)
      CALL transfert_request(f_ue_mean, req_e0_vect)

      DO ind = 1, ndomain
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ue = f_ue_mean(ind)
         rhodz = f_rhodz_mean(ind)
         hflux = f_hflux(ind)

         DO l = ll_begin, ll_end
            !DIR$ IVDEP
            DO ij = ij_begin, ij_end
               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)
            END DO
         END DO
      END DO

      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_mean(ind)
         rhodz = f_rhodz_mean(ind)

         DO l = 1, llm
            IF (l < begin_sponge) CYCLE
            ll = l - begin_sponge + 1
            !DIR$ IVDEP
            DO ij = ij_begin, ij_end
               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)
            END DO
         END DO
      END DO

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

!$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, :)
         END DO
      END DO

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

      CALL xios_send_field("ulat_sponge_mean", domain_zonal_mean)
!$OMP END MASTER
      CALL xios_read_field("ulat_from_sponge", f_ulat_sponge)

      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)
         DO l = ll_begin, ll_end
            IF (l < begin_sponge) THEN
               DO ij = ij_begin, ij_end
                  ulon(ij, l) = 0.
                  ulat(ij, l) = 0.
               END DO
            ELSE
               ll = l - begin_sponge + 1
               DO ij = ij_begin, ij_end
                  ulon(ij, l) = ulon_sponge(ij, ll)
                  ulat(ij, l) = ulat_sponge(ij, ll)
               END DO
            END IF
         END DO
      END DO
      CALL transfert_request(f_ulon, req_i0)
      CALL transfert_request(f_ulon, req_i1)
      CALL transfert_request(f_ulat, req_i0)
      CALL transfert_request(f_ulat, req_i1)
      CALL ulonlat2un(f_ulon, f_ulat, f_ue_sponge_zonal)
      CALL transfert_request(f_ue_sponge_zonal, req_e0_vect)

      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, :)
            END DO
         END DO
         CALL xios_send_field("theta_sponge_mean", domain_zonal_mean)
!$OMP END MASTER
         CALL xios_read_field("theta_from_sponge", f_theta_sponge)

         DO ind = 1, ndomain
            CALL swap_dimensions(ind)
            CALL swap_geometry(ind)
            theta_sponge_zonal = f_theta_sponge_zonal(ind)
            theta_sponge = f_theta_sponge(ind)
            DO l = ll_begin, ll_end
               IF (l < begin_sponge) THEN
                  DO ij = ij_begin, ij_end
                     theta_sponge_zonal(ij, l) = 0
                  END DO
               ELSE
                  ll = l - begin_sponge + 1
                  DO ij = ij_begin, ij_end
                     theta_sponge_zonal(ij, l) = theta_sponge(ij, ll)
                  END DO
               END IF
            END DO
         END DO
         CALL transfert_request(f_theta_sponge_zonal, req_i0)
      END IF

   END SUBROUTINE sponge_update_zonal_mean

   SUBROUTINE sponge_apply_zonal_mean(mode_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
      USE wind_to_lonlat_mod
      USE xios_mod
      IMPLICIT NONE
      INTEGER, INTENT(IN)    :: mode_sponge
      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(:, :)
      REAL(rstd), POINTER         :: ue(:, :)
      REAL(rstd), POINTER         :: theta_rhodz(:, :, :)
      REAL(rstd), POINTER         :: dtheta_rhodz(:, :, :)
      REAL(rstd), POINTER         :: ulon(:, :)
      REAL(rstd), POINTER         :: ulat(:, :)
      REAL(rstd), POINTER         :: ue_sponge_zonal(:, :)
      REAL(rstd), POINTER         :: theta_sponge_zonal(:, :)

      INTEGER :: ind, ij, l, ll

      DO ind = 1, ndomain
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)

         ue_sponge_zonal = f_ue_sponge_zonal(ind)
         due = f_due(ind)
         ue = f_ue(ind)
         theta_sponge_zonal = f_theta_sponge_zonal(ind)
         dtheta_rhodz = f_dtheta_rhodz(ind)
         theta_rhodz = f_theta_rhodz(ind)
         DO l = ll_begin, ll_end
            IF (l < begin_sponge) CYCLE
            IF (mode_sponge == 2) THEN
               DO ij = ij_begin, ij_end
                  due(ij + u_right, l) = ue_sponge_zonal(ij + u_right, l)
                  due(ij + u_lup, l) = ue_sponge_zonal(ij + u_lup, l)
                  due(ij + u_ldown, l) = ue_sponge_zonal(ij + u_ldown, l)
                  dtheta_rhodz(ij, l, 1) = theta_rhodz(ij, l, 1)
               END DO
            ELSE IF (mode_sponge == 3) THEN
               DO ij = ij_begin, ij_end
                  due(ij + u_right, l) = ue_sponge_zonal(ij + u_right, l)
                  due(ij + u_lup, l) = ue_sponge_zonal(ij + u_lup, l)
                  due(ij + u_ldown, l) = ue_sponge_zonal(ij + u_ldown, l)
                  dtheta_rhodz(ij, l, 1) = theta_sponge_zonal(ij, l)
               END DO
            END IF
         END DO
      END DO

!$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.
               END DO
            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)
               END DO
            END IF
         END DO
      END DO

!ym to check results
!    CALL transfert_request(f_due,req_e1_vect)
!    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_apply_zonal_mean

   SUBROUTINE sponge_reset_mean(mode_sponge)
      USE icosa
      USE theta2theta_rhodz_mod
      USE exner_mod
      USE geopotential_mod
      USE trace
      USE time_mod
      USE omp_para
      USE wind_to_lonlat_mod
      USE xios_mod
      IMPLICIT NONE
      INTEGER, INTENT(IN)    :: mode_sponge

      REAL(rstd), POINTER         :: ue_mean(:, :)
      REAL(rstd), POINTER         :: rhodz_mean(:, :)
      REAL(rstd), POINTER         :: theta_rhodz_mean(:, :)

      INTEGER :: ind, ij, l, ll

      DO ind = 1, ndomain
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ue_mean = f_ue_mean(ind)
         rhodz_mean = f_rhodz_mean(ind)
         theta_rhodz_mean = f_theta_rhodz_mean(ind)

         DO l = ll_begin, ll_end
            !DIR$ IVDEP
            DO ij = ij_begin_ext, ij_end_ext
               ue_mean(ij + u_right, l) = 0
               ue_mean(ij + u_lup, l) = 0
               ue_mean(ij + u_ldown, l) = 0
               rhodz_mean(ij, l) = 0
            END DO
         END DO

         IF (mode_sponge == 3) THEN
            DO l = ll_begin, ll_end
               !DIR$ IVDEP
               DO ij = ij_begin_ext, ij_end_ext
                  theta_rhodz_mean(ij, l) = 0
               END DO
            END DO
         END IF

      END DO
      nb_it_mean = 0
   END SUBROUTINE sponge_reset_mean

   SUBROUTINE sponge_update_mean(mode_sponge, f_rhodz, f_ue, f_theta_rhodz)
      USE icosa
      USE theta2theta_rhodz_mod
      USE exner_mod
      USE geopotential_mod
      USE trace
      USE time_mod
      USE omp_para
      USE wind_to_lonlat_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_theta_rhodz(:)

      REAL(rstd), POINTER         :: ue(:, :)
      REAL(rstd), POINTER         :: rhodz(:, :)
      REAL(rstd), POINTER         :: theta_rhodz(:, :, :)
      REAL(rstd), POINTER         :: ue_mean(:, :)
      REAL(rstd), POINTER         :: rhodz_mean(:, :)
      REAL(rstd), POINTER         :: theta_rhodz_mean(:, :)

      INTEGER :: ind, ij, l

      CALL transfert_request(f_ue, req_e0_vect)

      DO ind = 1, ndomain
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ue = f_ue(ind)
         rhodz = f_rhodz(ind)
         theta_rhodz = f_theta_rhodz(ind)
         ue_mean = f_ue_mean(ind)
         rhodz_mean = f_rhodz_mean(ind)
         theta_rhodz_mean = f_theta_rhodz_mean(ind)

         DO l = ll_begin, ll_end
            !DIR$ IVDEP
            DO ij = ij_begin, ij_end
               ue_mean(ij + u_right, l) = (ue_mean(ij + u_right, l)*nb_it_mean + ue(ij + u_right, l))/(nb_it_mean + 1)
               ue_mean(ij + u_lup, l) = (ue_mean(ij + u_lup, l)*nb_it_mean + ue(ij + u_lup, l))/(nb_it_mean + 1)
               ue_mean(ij + u_ldown, l) = (ue_mean(ij + u_ldown, l)*nb_it_mean + ue(ij + u_ldown, l))/(nb_it_mean + 1)
               rhodz_mean(ij, l) = (rhodz_mean(ij, l)*nb_it_mean + rhodz(ij, l))/(nb_it_mean + 1)
            END DO
         END DO

         IF (mode_sponge == 3) THEN
            DO l = ll_begin, ll_end
               !DIR$ IVDEP
               DO ij = ij_begin, ij_end
                  theta_rhodz_mean(ij, l) = (theta_rhodz_mean(ij, l)*nb_it_mean + theta_rhodz(ij, l, 1))/(nb_it_mean + 1)
               END DO
            END DO
         END IF

      END DO
      nb_it_mean = nb_it_mean + 1
   END SUBROUTINE sponge_update_mean

END MODULE sponge_zonal_mod

