MODULE explicit_scheme_mod
   USE prec, ONLY: rstd
   IMPLICIT NONE
   PRIVATE

   REAL(rstd), DIMENSION(4), PARAMETER :: coef_rk4 = (/.25, 1./3., .5, 1./)
   REAL(rstd), DIMENSION(5), PARAMETER :: coef_rk25 = (/.25, 1./6., 3./8., .5, 1./)

   PUBLIC :: explicit_scheme
CONTAINS

   SUBROUTINE explicit_scheme(it, fluxt_zero)
      USE time_scheme_base_mod, ONLY : &
         DEC_to_legacy, legacy_to_DEC, compute_accumulate_fluxes, &
         f_geopot, f_phis, f_q, &
         f_ps, f_psm1, f_psm2, f_dps, &
         f_mass, f_massm1, f_massm2, f_dmass, &
         f_u, f_du, f_um1, f_um2, &
         f_theta_rhodz, f_theta_rhodzm1, f_theta_rhodzm2, f_dtheta_rhodz, &
         f_hfluxt, f_wfluxt, f_hflux, f_wflux, &
         scheme, nb_stage, euler, rk4, rk25, mlf

      USE euler_scheme_mod, ONLY: euler_scheme
      USE domain_mod
      USE field_mod
      USE time_mod
      USE omp_para
      USE caldyn_mod
      USE trace
      USE dimensions
      USE geometry
      REAL(rstd), POINTER :: ps(:), psm1(:), psm2(:), dps(:)
      REAL(rstd), POINTER :: u(:, :), um1(:, :), um2(:, :), du(:, :)
      REAL(rstd), POINTER :: mass(:, :), massm1(:, :), dmass(:, :)
      REAL(rstd), POINTER :: theta_rhodz(:, :, :), theta_rhodzm1(:, :, :), theta_rhodzm2(:, :, :), dtheta_rhodz(:, :, :)
      REAL(rstd), POINTER :: hflux(:, :), wflux(:, :), hfluxt(:, :), wfluxt(:, :)
      LOGICAL :: fluxt_zero(ndomain) ! set to .TRUE. to start accumulating fluxes in time
      INTEGER :: it, stage

      CALL legacy_to_DEC(f_ps, f_u)
      DO stage = 1, nb_stage
         !       CALL checksum(f_ps)
         !       CALL checksum(f_theta_rhodz)
         !       CALL checksum(f_mass)
         CALL caldyn((stage == 1) .AND. (MOD(it, itau_out) == 0), &
                     f_phis, f_ps, f_mass, f_theta_rhodz, f_u, f_q, &
                     f_geopot, f_hflux, f_wflux, f_dps, f_dmass, f_dtheta_rhodz, f_du)
         !       CALL checksum(f_dps)
         !       CALL checksum(f_dtheta_rhodz)
         !       CALL checksum(f_dmass)
         SELECT CASE (scheme)
         CASE (euler)
            CALL euler_scheme(.TRUE., fluxt_zero)
         CASE (rk4)
            CALL rk_scheme(stage, coef_rk4)
         CASE (rk25)
            CALL rk_scheme(stage, coef_rk25)
         CASE (mlf)
            CALL leapfrog_matsuno_scheme(stage)
         CASE DEFAULT
            PRINT *, 'Internal error in explicit_scheme : illegal value of scheme = ', scheme
            STOP
         END SELECT
      END DO
      CALL DEC_to_legacy(f_ps, f_u)

   CONTAINS

      SUBROUTINE RK_scheme(stage, coef)
         USE disvert_mod
         INTEGER :: ind, stage
         REAL(rstd), INTENT(IN) :: coef(:)
         REAL(rstd) :: tau
         INTEGER :: ij, l

         CALL trace_start("RK_scheme")

         tau = dt*coef(stage)

         ! if mass coordinate, deal with ps first on one core
         IF (caldyn_eta == eta_mass) THEN
            IF (is_omp_first_level) THEN

               DO ind = 1, ndomain
                  IF (.NOT. assigned_domain(ind)) CYCLE
                  CALL swap_dimensions(ind)
                  CALL swap_geometry(ind)
                  ps = f_ps(ind); psm1 = f_psm1(ind); dps = f_dps(ind)

                  IF (stage == 1) THEN ! first stage : save model state in XXm1
                     !DIR$ SIMD
                     DO ij = ij_begin, ij_end
                        psm1(ij) = ps(ij)
                     END DO
                  END IF

                  ! updates are of the form x1 := x0 + tau*f(x1)
                  !DIR$ SIMD
                  DO ij = ij_begin, ij_end
                     ps(ij) = psm1(ij) + tau*dps(ij)
                  END DO
               END DO
            END IF
            !         CALL send_message(f_ps,req_ps)
            !ym no overlap for now
            !         CALL wait_message(req_ps)

         ELSE ! Lagrangian coordinate, deal with mass
            DO ind = 1, ndomain
               IF (.NOT. assigned_domain(ind)) CYCLE
               CALL swap_dimensions(ind)
               CALL swap_geometry(ind)
               mass = f_mass(ind); dmass = f_dmass(ind); massm1 = f_massm1(ind)

               IF (stage == 1) THEN ! first stage : save model state in XXm1
                  DO l = ll_begin, ll_end
                     !DIR$ SIMD
                     DO ij = ij_begin, ij_end
                        massm1(ij, l) = mass(ij, l)
                     END DO
                  END DO
               END IF

               ! updates are of the form x1 := x0 + tau*f(x1)
               DO l = ll_begin, ll_end
                  !DIR$ SIMD
                  DO ij = ij_begin, ij_end
                     mass(ij, l) = massm1(ij, l) + tau*dmass(ij, l)
                  END DO
               END DO
            END DO
            !         CALL send_message(f_mass,req_mass)
            !ym no overlap for now
            !         CALL wait_message(req_mass)

         END IF

         ! now deal with other prognostic variables
         DO ind = 1, ndomain
            IF (.NOT. assigned_domain(ind)) CYCLE
            CALL swap_dimensions(ind)
            CALL swap_geometry(ind)
            u = f_u(ind); du = f_du(ind); um1 = f_um1(ind)
            theta_rhodz = f_theta_rhodz(ind)
            theta_rhodzm1 = f_theta_rhodzm1(ind)
            dtheta_rhodz = f_dtheta_rhodz(ind)

            IF (stage == 1) THEN ! first stage : save model state in XXm1
               DO l = ll_begin, ll_end
                  !DIR$ SIMD
                  DO ij = ij_begin, ij_end
                     um1(ij + u_right, l) = u(ij + u_right, l)
                     um1(ij + u_lup, l) = u(ij + u_lup, l)
                     um1(ij + u_ldown, l) = u(ij + u_ldown, l)
                     theta_rhodzm1(ij, l, 1) = theta_rhodz(ij, l, 1)
                  END DO
               END DO
            END IF

            DO l = ll_begin, ll_end
               !DIR$ SIMD
               DO ij = ij_begin, ij_end
                  u(ij + u_right, l) = um1(ij + u_right, l) + tau*du(ij + u_right, l)
                  u(ij + u_lup, l) = um1(ij + u_lup, l) + tau*du(ij + u_lup, l)
                  u(ij + u_ldown, l) = um1(ij + u_ldown, l) + tau*du(ij + u_ldown, l)
                  theta_rhodz(ij, l, 1) = theta_rhodzm1(ij, l, 1) + tau*dtheta_rhodz(ij, l, 1)
               END DO
            END DO

            IF (stage == nb_stage) THEN ! accumulate mass fluxes at last stage
               hflux = f_hflux(ind); hfluxt = f_hfluxt(ind)
               wflux = f_wflux(ind); wfluxt = f_wfluxt(ind)
               CALL compute_accumulate_fluxes(hflux, wflux, hfluxt, wfluxt, dt, fluxt_zero(ind))
            END IF
         END DO

         CALL trace_end("RK_scheme")

      END SUBROUTINE RK_scheme

      SUBROUTINE leapfrog_matsuno_scheme(stage)
         IMPLICIT NONE
         INTEGER :: ind, stage
         REAL :: tau

         CALL trace_start("leapfrog_matsuno_scheme")

         tau = dt/nb_stage
         DO ind = 1, ndomain
            IF (.NOT. assigned_domain(ind)) CYCLE
            CALL swap_dimensions(ind)
            CALL swap_geometry(ind)

            ps = f_ps(ind); u = f_u(ind); theta_rhodz = f_theta_rhodz(ind)
            psm1 = f_psm1(ind); um1 = f_um1(ind); theta_rhodzm1 = f_theta_rhodzm1(ind)
            psm2 = f_psm2(ind); um2 = f_um2(ind); theta_rhodzm2 = f_theta_rhodzm2(ind)
            dps = f_dps(ind); du = f_du(ind); dtheta_rhodz = f_dtheta_rhodz(ind)

            IF (stage == 1) THEN
               psm1(:) = ps(:); um1(:, :) = u(:, :); theta_rhodzm1(:, :, :) = theta_rhodz(:, :, :)
               ps(:) = psm1(:) + tau*dps(:)
               u(:, :) = um1(:, :) + tau*du(:, :)
               theta_rhodz(:, :, :) = theta_rhodzm1(:, :, :) + tau*dtheta_rhodz(:, :, :)

            ELSE IF (stage == 2) THEN

               ps(:) = psm1(:) + tau*dps(:)
               u(:, :) = um1(:, :) + tau*du(:, :)
               theta_rhodz(:, :, :) = theta_rhodzm1(:, :, :) + tau*dtheta_rhodz(:, :, :)

               psm2(:) = psm1(:); theta_rhodzm2(:, :, :) = theta_rhodzm1(:, :, :); um2(:, :) = um1(:, :)
               psm1(:) = ps(:); theta_rhodzm1(:, :, :) = theta_rhodz(:, :, :); um1(:, :) = u(:, :)

            ELSE

               ps(:) = psm2(:) + 2*tau*dps(:)
               u(:, :) = um2(:, :) + 2*tau*du(:, :)
               theta_rhodz(:, :, :) = theta_rhodzm2(:, :, :) + 2*tau*dtheta_rhodz(:, :, :)

               psm2(:) = psm1(:); theta_rhodzm2(:, :, :) = theta_rhodzm1(:, :, :); um2(:, :) = um1(:, :)
               psm1(:) = ps(:); theta_rhodzm1(:, :, :) = theta_rhodz(:, :, :); um1(:, :) = u(:, :)

            END IF

         END DO
         CALL trace_end("leapfrog_matsuno_scheme")

      END SUBROUTINE leapfrog_matsuno_scheme

   END SUBROUTINE Explicit_scheme

END MODULE explicit_scheme_mod
