MODULE time_scheme_base_mod
   USE prec, ONLY : rstd
   USE field_mod, ONLY: t_field
   IMPLICIT NONE
   PRIVATE
   SAVE

   TYPE(t_field), POINTER, PUBLIC :: &
      f_geopot(:), f_q(:), &
      f_rhodz(:), f_mass(:), f_massm1(:), f_massm2(:), f_dmass(:), &     ! current and previous time steps + tendency of mass,
      f_phis(:), f_ps(:), f_psm1(:), f_psm2(:), f_dps(:), &              ! column-integrated mass
      f_u(:), f_um1(:), f_um2(:), f_du(:), f_W(:), &                     ! 'horizontal' and vertical (NH) momentum
      f_theta_rhodz(:), f_theta_rhodzm1(:), f_theta_rhodzm2(:), f_dtheta_rhodz(:), & ! mass-weighted theta/entropy
      f_hflux(:), f_wflux(:), f_hfluxt(:), f_wfluxt(:), &                ! accumulated mass fluxes
      f_du_slow(:, :), f_du_fast(:, :), &                                ! slow/fast trend arrays
      f_dmass_slow(:, :), f_dps_slow(:, :), f_dtheta_rhodz_slow(:, :), & ! for HEVI time scheme
      f_dPhi_slow(:, :), f_dPhi_fast(:, :), &                            ! geopotential tendencies (NH)
      f_dW_slow(:, :), f_dW_fast(:, :)                                   ! vertical momentum tendencies (NH)

   INTEGER, PARAMETER, PUBLIC :: explicit = 1, hevi = 2, &
                                 euler = 1, rk4 = 2, mlf = 3, rk25 = 4, &
                                 ark11 = 5, ark23 = 6, ark33 = 7, ark25 = 8

   INTEGER, PARAMETER, PUBLIC :: max_nb_stage = 5

   INTEGER, PUBLIC :: nb_stage, nb_storage, matsuno_period, scheme, scheme_family
   !$OMP THREADPRIVATE(nb_stage, nb_storage, matsuno_period, scheme, scheme_family)

   PUBLIC :: compute_accumulate_fluxes, legacy_to_DEC, DEC_to_legacy

CONTAINS

   SUBROUTINE compute_accumulate_fluxes(hflux, wflux, hfluxt, wfluxt, tau, fluxt_zero)
      USE accumulate_fluxes_mod
      USE grid_param, ONLY: grid_type, grid_ico, grid_unst

      REAL(rstd), INTENT(IN)    :: hflux(:, :), wflux(:, :)    ! horizontal and vertical mass fluxes
      REAL(rstd), INTENT(INOUT) :: hfluxt(:, :), wfluxt(:, :)  ! time-integrated fluxes
      REAL(rstd), INTENT(IN) :: tau                            ! time interval
      LOGICAL, INTENT(INOUT) :: fluxt_zero                     ! .TRUE. to reset fluxes

      SELECT CASE (grid_type)
      CASE (grid_ico)
         CALL accumulate_fluxes_hex(hflux, wflux, hfluxt, wfluxt, tau, fluxt_zero)
      CASE (grid_unst)
         CALL accumulate_fluxes_unst(hflux, wflux, hfluxt, wfluxt, tau, fluxt_zero)
      END SELECT
   END SUBROUTINE compute_accumulate_fluxes

   SUBROUTINE legacy_to_DEC(f_ps, f_u)
      USE icosa
      USE disvert_mod
      USE omp_para
      USE trace
      TYPE(t_field), POINTER :: f_ps(:), f_u(:)
      REAL(rstd), POINTER :: ps(:), u(:, :)
      INTEGER :: ind, ij, l

      CALL trace_start("legacy_to_DEC")

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

         SELECT CASE (grid_type)

         CASE (grid_ico)

            IF (caldyn_eta == eta_mass .AND. is_omp_first_level) THEN ! update ps
               ps = f_ps(ind)
               !$acc parallel loop async default(present)
               !DIR$ SIMD
               DO ij = ij_begin, ij_end
                  ps(ij) = (ps(ij) - ptop)/g ! convert ps to column-integrated mass
               END DO
            END IF

            u = f_u(ind)
            !$acc parallel loop async default(present)
            DO l = ll_begin, ll_end
               !$acc loop
               !DIR$ SIMD
               DO ij = ij_begin, ij_end
                  u(ij + u_right, l) = u(ij + u_right, l)*de(ij + u_right)
                  u(ij + u_lup, l) = u(ij + u_lup, l)*de(ij + u_lup)
                  u(ij + u_ldown, l) = u(ij + u_ldown, l)*de(ij + u_ldown)
               END DO
            END DO

         CASE (grid_unst)

            IF (caldyn_eta == eta_mass .AND. is_omp_first_level) THEN ! update ps
               ps = f_ps(ind)
               ps(:) = (ps(:) - ptop)/g ! FIXME : OpenMP+unstructured
            END IF

            u = f_u(ind)

!          PRINT *, 'legacy_to_DEC', llm, edge_num, SHAPE(u), SHAPE(de) ! FIXME
            DO ij = 1, edge_num
               DO l = 1, llm
                  u(l, ij) = u(l, ij)*de(ij)
               END DO
            END DO
         CASE DEFAULT
            STOP 'Unsupported grid_type encountered in legacy_to_DEC'
         END SELECT
      END DO

      CALL trace_end("legacy_to_DEC")

   END SUBROUTINE Legacy_to_DEC

   SUBROUTINE DEC_to_legacy(f_ps, f_u)
      USE icosa
      USE disvert_mod
      USE omp_para
      USE trace
      TYPE(t_field), POINTER :: f_ps(:), f_u(:)
      REAL(rstd), POINTER :: ps(:), u(:, :)
      INTEGER :: ind, ij, l
      CALL trace_start("legacy_to_DEC")

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

         SELECT CASE (grid_type)

         CASE (grid_ico)
            IF (caldyn_eta == eta_mass .AND. is_omp_first_level) THEN
               ps = f_ps(ind)
               !$acc parallel loop async default(present)
               !DIR$ SIMD
               DO ij = ij_begin, ij_end
                  ps(ij) = ptop + ps(ij)*g ! convert column-integrated mass to ps
               END DO
            END IF

            u = f_u(ind)
            !$acc parallel loop async default(present)
            DO l = ll_begin, ll_end
               !$acc loop
               !DIR$ SIMD
               DO ij = ij_begin, ij_end
                  u(ij + u_right, l) = u(ij + u_right, l)/de(ij + u_right)
                  u(ij + u_lup, l) = u(ij + u_lup, l)/de(ij + u_lup)
                  u(ij + u_ldown, l) = u(ij + u_ldown, l)/de(ij + u_ldown)
               END DO
            END DO

         CASE (grid_unst)
            IF (caldyn_eta == eta_mass .AND. is_omp_first_level) THEN
               ps = f_ps(ind)
               ps(:) = ptop + ps(:)*g ! FIXME : OpenMP+unstructured
            END IF
            u = f_u(ind)
            !          PRINT *, 'DEC_to_legacy', llm, edge_num, SHAPE(u), SHAPE(de) ! FIXME
            DO ij = 1, edge_num
               DO l = 1, llm
                  u(l, ij) = u(l, ij)/de(ij)
               END DO
            END DO
         CASE DEFAULT
            STOP 'Unsupported grid_type encountered in legacy_to_DEC'
         END SELECT
      END DO

      CALL trace_end("DEC_to_legacy")
   END SUBROUTINE DEC_to_legacy

END MODULE time_scheme_base_mod
