MODULE hevi_scheme_mod
   USE prec, ONLY: rstd
   USE time_scheme_base_mod, ONLY: maxn => max_nb_stage
   IMPLICIT NONE
   PRIVATE
   SAVE

   REAL(rstd) :: wj(maxn)        ! flux accmulation weights
   REAL(rstd) :: bjl(maxn, maxn) ! update coefs (slow)
   REAL(rstd) :: cjl(maxn, maxn) ! update coefficients (fast)
   REAL(rstd) :: taujj(maxn)     ! diagonal coefficient on the fast Butcher tableau (implicit part of ARK schemes)
   INTEGER    :: jstore(maxn)    ! where to store the tendencies computed at stage j, for low-storage schemes

   REAL(rstd), PARAMETER, DIMENSION(3) :: zero = (/0., 0., 0./)

   PUBLIC :: set_coefs_ark11, set_coefs_ark23, set_coefs_ark33, set_coefs_ark25, hevi_scheme

CONTAINS

   SUBROUTINE set_coefs_ark23(dt)
      ! ARK2 scheme by Giraldo, Kelly, Constantinescu 2013
      ! See Weller et al., 2013 - ARK2 scheme Fig. 2
      REAL(rstd), INTENT(IN) :: dt
      REAL(rstd), PARAMETER :: delta = .5/SQRT(2.), gamma = 1.-2.*delta
!    REAL(rstd), PARAMETER :: alpha=(3.+SQRT(8.))/6. ! original value in GKC2013
      REAL(rstd), PARAMETER :: alpha = 0.7
      REAL(rstd), PARAMETER, DIMENSION(3) :: wj = (/delta, delta, gamma/)
      CALL set_coefs_hevi(3, dt, &
                          (/zero, (/2.*gamma, 0., 0./), (/1 - alpha, alpha, 0./), wj/), &
                          (/zero, (/gamma, gamma, 0./), wj, wj/))
   END SUBROUTINE set_coefs_ark23

   SUBROUTINE set_coefs_ark33(dt)
      ! Fully-explicit RK3 scheme disguised as ARK
      REAL(rstd), INTENT(IN) :: dt
      CALL set_coefs_rk(3, dt, (/zero, (/.5, 0., 0./), (/-1., 2., 0./), (/1./6., 2./3., 1./6./)/))
   END SUBROUTINE set_coefs_ark33

   SUBROUTINE set_coefs_ark11(dt)
      ! Euler scheme disguised as ARK for development purposes
      REAL(rstd) :: dt
      CALL set_coefs_rk(1, dt, (/zero, (/1., 0., 0./), (/0., 0., 0./), (/0., 0., 0./)/))
   END SUBROUTINE set_coefs_ark11

   SUBROUTINE set_coefs_rk(n, dt, ajl)
      INTEGER, INTENT(IN) :: n
      REAL(rstd), INTENT(IN) :: dt, ajl(n, n + 1)
      CALL set_coefs_hevi(n, dt, ajl, ajl)
   END SUBROUTINE set_coefs_rk

   SUBROUTINE set_coefs_hevi(n, dt, ajl_slow, ajl_fast)
      USE time_scheme_base_mod, ONLY: nb_stage, scheme_family, hevi
      ! from the slow and fast Butcher tableaus ajl_slow, ajl_fast
      ! we compute/extract
      !   * slow and fast update coefficients = difference between two lines of a Butcher tableau
      !   * implicit coefficients = diagonal of fast Butcher tableau
      !   * flux accumulation coefficients = last line of slow Butcher tableau
      INTEGER, INTENT(IN) :: n
      REAL(rstd), INTENT(IN) :: dt, ajl_slow(n, n + 1), ajl_fast(n, n + 1) ! fast/slow Butcher tableaus
      INTEGER :: j, k

      DO j = 1, n
         DO k = 1, n
            bjl(k, j) = dt*(ajl_slow(k, j + 1) - ajl_slow(k, j))
            cjl(k, j) = dt*(ajl_fast(k, j + 1) - ajl_fast(k, j))
         END DO
         taujj(j) = dt*ajl_fast(j, j)
         wj(j) = dt*ajl_slow(j, n + 1)
         jstore(j) = j
      END DO
   END SUBROUTINE set_coefs_hevi

   SUBROUTINE set_coefs_ark25(dt)
      ! Fully-explicit RK2.5 scheme disguised as ARK
      REAL(rstd), INTENT(IN) :: dt
      REAL(rstd), DIMENSION(5), PARAMETER :: coef_rk25 = (/.25, 1./6., 3./8., .5, 1./)
      CALL set_coefs_low_storage(5, dt, coef_rk25)
   END SUBROUTINE set_coefs_ark25

   SUBROUTINE set_coefs_low_storage(n, dt, aj)
      INTEGER, INTENT(IN) :: n
      REAL(rstd), INTENT(IN) :: dt, aj(n) ! coefficients
      ! explicit scheme of the form :
      !    x1 = x0
      !    x1 = x0 + a(1).f(x1)
      !    x1 = x0 + a(2).f(x1)
      !    ...
      !    x1 = x0 + a(n).f(x1)
      ! reformulation in low-storage 'update' form :
      !    k1 = f(x)
      !    x  = x + a(1)*k1
      !    k2 = f(x)
      !    x = x - a(1)*k1 + a(2)*k2
      !    k1 = f(x)
      !    x = x - a(2)*k2 + a(3)*k1
      INTEGER :: j
      bjl(:, :) = 0.
      taujj(:) = 0.
      wj(:) = 0.
      wj(n) = dt*aj(n) ! mass fluxes are accumumated only at the last step j=n
      ! j=1
      jstore(1) = 1
      bjl(1, 1) = dt*aj(1)
      ! j>1
      DO j = 2, n
         IF (MOD(j, 2) == 0) THEN
            jstore(j) = 2
            bjl(1, j) = -dt*aj(j - 1)
            bjl(2, j) = dt*aj(j)
         ELSE
            jstore(j) = 1
            bjl(2, j) = -dt*aj(j - 1)
            bjl(1, j) = dt*aj(j)
         END IF
      END DO
      cjl(:, :) = bjl(:, :)
   END SUBROUTINE set_coefs_low_storage

   SUBROUTINE HEVI_scheme(it, fluxt_zero)
      USE time_scheme_base_mod, ONLY: &
         DEC_to_legacy, legacy_to_DEC, compute_accumulate_fluxes, &
         f_phis, f_q, &
         f_ps, f_dps_slow, &
         f_mass, f_dmass_slow, &
         f_theta_rhodz, f_dtheta_rhodz_slow, &
         f_u, f_du_slow, f_du_fast, &
         f_W, f_dW_slow, f_dW_fast, &
         f_geopot, f_dPhi_slow, f_dPhi_fast, &
         f_hfluxt, f_wfluxt, f_hflux, f_wflux, &
         nb_stage, nb_storage
      USE caldyn_vars_mod, ONLY: hydrostatic
      USE domain_mod, ONLY: ndomain
      USE field_mod
      USE abort_mod, ONLY: abort_acc
      USE time_mod
      USE disvert_mod
      USE caldyn_hevi_mod
      USE omp_para
      USE checksum_mod
      LOGICAL :: fluxt_zero(ndomain) ! set to .TRUE. to start accumulating fluxes in time
      INTEGER :: it, j, jj, l, ind
      REAL(rstd), POINTER :: hflux(:, :), wflux(:, :), hfluxt(:, :), wfluxt(:, :)

      CALL legacy_to_DEC(f_ps, f_u)

      DO j = 1, nb_stage
         jj = jstore(j) ! store tendencies at index jj (jj=j except for low-storage schemes)
         CALL caldyn_hevi((j == 1) .AND. (MOD(it, itau_out) == 0), taujj(j), &
                          f_phis, f_ps, f_mass, f_theta_rhodz, f_u, f_q, &
                          f_W, f_geopot, f_hflux, f_wflux, &
                          f_dps_slow(:, jj), f_dmass_slow(:, jj), f_dtheta_rhodz_slow(:, jj), &
                          f_du_slow(:, jj), f_du_fast(:, jj), &
                          f_dPhi_slow(:, jj), f_dPhi_fast(:, jj), &
                          f_dW_slow(:, jj), f_dW_fast(:, jj))
         ! accumulate mass fluxes for transport scheme

         DO ind = 1, ndomain
            IF (.NOT. assigned_domain(ind)) CYCLE
            CALL swap_dimensions(ind)
            hflux = f_hflux(ind); hfluxt = f_hfluxt(ind)
            wflux = f_wflux(ind); wfluxt = f_wfluxt(ind)
            CALL compute_accumulate_fluxes(hflux, wflux, hfluxt, wfluxt, wj(j), fluxt_zero(ind))
         END DO

         ! update model state
         DO l = 1, MIN(j, nb_storage)
            IF (caldyn_eta == eta_mass) THEN
               CALL update_2D(bjl(l, j), f_ps, f_dps_slow(:, l))
            ELSE
               CALL update_3D(bjl(l, j), f_mass, f_dmass_slow(:, l))
            END IF
            CALL update_4D(bjl(l, j), f_theta_rhodz, f_dtheta_rhodz_slow(:, l))
            CALL update_3D(bjl(l, j), f_u, f_du_slow(:, l))
            CALL update_3D(cjl(l, j), f_u, f_du_fast(:, l))
            IF (.NOT. hydrostatic) THEN
               CALL abort_acc("HEVI_scheme/!hydrostatic")
               CALL update_3D(bjl(l, j), f_W, f_dW_slow(:, l))
               CALL update_3D(cjl(l, j), f_W, f_dW_fast(:, l))
               CALL update_3D(bjl(l, j), f_geopot, f_dPhi_slow(:, l))
               CALL update_3D(cjl(l, j), f_geopot, f_dPhi_fast(:, l))
            END IF
!$OMP BARRIER
         END DO
      END DO
      CALL DEC_to_legacy(f_ps, f_u)
   END SUBROUTINE HEVI_scheme

   SUBROUTINE update_4D(w, f_y, f_dy)
      USE dimensions
      USE field_mod
      USE grid_param, ONLY: nqdyn
      USE domain_mod, ONLY: ndomain, assigned_domain
      REAL(rstd) :: w
      TYPE(t_field) :: f_y(:), f_dy(:)
      REAL(rstd), POINTER :: y(:, :, :), dy(:, :, :)
      INTEGER :: ind, iq
      IF (w /= 0.) THEN
         DO ind = 1, ndomain
            IF (.NOT. assigned_domain(ind)) CYCLE
            CALL swap_dimensions(ind)
            dy = f_dy(ind); y = f_y(ind)
            DO iq = 1, nqdyn
               CALL compute_update_3D(w, y(:, :, iq), dy(:, :, iq))
            END DO
         END DO
      END IF
   END SUBROUTINE update_4D

   SUBROUTINE update_3D(w, f_y, f_dy)
      USE dimensions
      USE field_mod
      USE domain_mod, ONLY: ndomain, assigned_domain
      REAL(rstd) :: w
      TYPE(t_field) :: f_y(:), f_dy(:)
      REAL(rstd), POINTER :: y(:, :), dy(:, :)
      INTEGER :: ind
      IF (w /= 0.) THEN
         DO ind = 1, ndomain
            IF (.NOT. assigned_domain(ind)) CYCLE
            CALL swap_dimensions(ind)
            dy = f_dy(ind); y = f_y(ind)
            CALL compute_update_3D(w, y, dy)
         END DO
      END IF
   END SUBROUTINE update_3D

   SUBROUTINE compute_update_3D(w, y, dy)
      USE omp_para
      USE disvert_mod
      REAL(rstd) :: w
      REAL(rstd) :: y(:, :), dy(:, :)
      INTENT(INOUT) :: y
      INTENT(IN) :: dy
      INTEGER :: l, l_begin, l_end

      IF (grid_type == grid_ico) THEN
         l_begin = ll_begin
         l_end = ll_end
      ELSE ! unstructured : FIXME OpenMP
         l_begin = 1
         l_end = SIZE(y, 2)
      END IF

      !$acc kernels async default(present)
      DO l = l_begin, l_end
         y(:, l) = y(:, l) + w*dy(:, l)
      END DO
      !$acc end kernels
   END SUBROUTINE compute_update_3D

   SUBROUTINE update_2D(w, f_y, f_dy)
      USE dimensions
      USE omp_para
      USE field_mod
      USE domain_mod, ONLY: ndomain, assigned_domain
      REAL(rstd) :: w
      TYPE(t_field) :: f_y(:), f_dy(:)
      REAL(rstd), POINTER :: y(:), dy(:)
      INTEGER :: ind
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         dy = f_dy(ind); y = f_y(ind)
         IF (is_omp_level_master) CALL compute_update_2D(w, y, dy) ! FIXME OpenMP+unstructured
      END DO
   END SUBROUTINE update_2D

   SUBROUTINE compute_update_2D(w, y, dy)
      REAL(rstd) :: w, y(:), dy(:)
      INTENT(INOUT) :: y
      INTENT(IN) :: dy
      !$acc kernels async default(present)
      y(:) = y(:) + w*dy(:)
      !$acc end kernels
   END SUBROUTINE compute_update_2D

END MODULE hevi_scheme_mod
