MODULE timestep_unstructured_mod
   USE ISO_C_BINDING
   USE data_unstructured_mod
   USE caldyn_unstructured_mod
   USE transfer_unstructured_mod
   IMPLICIT NONE
   PRIVATE
   SAVE

CONTAINS

   SUBROUTINE scalar_laplacian(s, laps) BIND(C, name='dynamico_scalar_laplacian')
      REAL(C_DOUBLE), DIMENSION(llm, primal_num)   :: s, laps ! IN, OUT
      REAL(C_DOUBLE), DIMENSION(llm, edge_num)      :: grads  ! BUF
      !$OMP PARALLEL NUM_THREADS(nb_threads)
      CALL update_halo(transfer_primal, s)
      CALL compute_scalar_laplacian(s, grads, laps)
      !$OMP END PARALLEL
   END SUBROUTINE scalar_laplacian

   SUBROUTINE curl_laplacian(u, lapu) BIND(C, name='dynamico_curl_laplacian')
      REAL(C_DOUBLE), DIMENSION(llm, edge_num)   :: u, lapu ! IN, OUT
      REAL(C_DOUBLE), DIMENSION(llm, dual_num)   :: qv  ! BUF
      !$OMP PARALLEL NUM_THREADS(nb_threads)
      CALL update_halo(transfer_edge, u)
      CALL compute_curl_laplacian(u, qv, lapu)
      !$OMP END PARALLEL
   END SUBROUTINE curl_laplacian

   !
   SUBROUTINE caldyn_unstructured(tau, mass_col, rhodz, theta_rhodz, u, geopot, w, & ! IN : flow state
                                  theta, ps, pk, hflux, qv, &    ! OUT : diags (except surface geopot : IN)
                                  dmass_col, drhodz, dtheta_rhodz, du_fast, du_slow, &
                                  dPhi_fast, dPhi_slow, dW_fast, dW_slow) BIND(C, name='dynamico_caldyn_unstructured') ! OUT : tendencies
      REAL(C_DOUBLE), VALUE :: tau
      REAL(C_DOUBLE), DIMENSION(llm, primal_num)   :: rhodz, drhodz, pk, berni          ! IN, OUT, DIAG
      REAL(C_DOUBLE), DIMENSION(llm, primal_num, nqdyn)  :: theta_rhodz, dtheta_rhodz, theta  ! IN, OUT, DIAG
      REAL(C_DOUBLE), DIMENSION(llm + 1, primal_num) :: wflux, w, geopot, &               ! DIAG, INOUT
                                                        dPhi_fast, dPhi_slow, dW_fast, dW_slow       ! OUT
      REAL(C_DOUBLE), DIMENSION(llm, edge_num)      :: u, du_fast, du_slow, hflux, qu        ! INOUT,OUT,OUT,DIAG
      REAL(C_DOUBLE), DIMENSION(llm, dual_num)      :: qv                                ! DIAG
      REAL(C_DOUBLE), DIMENSION(primal_num)     :: ps, dmass_col, mass_col             ! OUT,OUT,IN (if eta_mass) or OUT,UNUSED,UNUSED (if eta_lag)
      ! buffers for fields that need to be shared among OpenMP threads
      ! vertical size is allocated as llm+1 even if only llm is needed
      REAL(C_DOUBLE), DIMENSION(llm + 1, primal_num) :: bufm1, bufm2, bufm3
      REAL(C_DOUBLE), DIMENSION(llm + 1, edge_num)     :: bufu1, bufu2, bufu3, bufu4

      REAL(C_DOUBLE)          :: time1, time2
      INTEGER :: ij

      time1 = OMP_GET_WTIME()

      IF (hydrostatic) THEN

         !$OMP PARALLEL NUM_THREADS(nb_threads)
         !$OMP DO SCHEDULE(STATIC)
         DO ij = 1, edge_num
            du_fast(:, ij) = 0.
            du_slow(:, ij) = 0.
         END DO
         !$OMP END DO
         CALL compute_theta(mass_col, rhodz, theta_rhodz, theta)
         CALL compute_geopot(rhodz, theta, ps, pk, geopot)

         CALL compute_caldyn_fast(tau, pk, berni, theta, geopot, du_fast, u)

         CALL compute_pvort_only(rhodz, u, qv, qu)
         CALL compute_caldyn_slow_hydro(rhodz, theta, u, berni, hflux, du_slow)
         CALL compute_coriolis(hflux, theta, qu, bufu1, drhodz, dtheta_rhodz, du_slow)
         IF (caldyn_eta == eta_mass) THEN
            CALL caldyn_vert(drhodz, rhodz, theta, u, dmass_col, wflux, dtheta_rhodz, du_slow, bufu1)
         END IF
         !$OMP END PARALLEL

      ELSE ! NH
         DO ij = 1, edge_num
            du_fast(:, ij) = 0.
            du_slow(:, ij) = 0.
         END DO
         DO ij = 1, primal_num
            wflux(1, ij) = 0.
            wflux(llm + 1, ij) = 0.
         END DO
         CALL compute_theta(mass_col, rhodz, theta_rhodz, theta)
         CALL compute_caldyn_solver(tau, rhodz, theta, bufm1, bufm2, bufm3, pk, geopot, W, dPhi_fast, dW_fast, du_fast)
         CALL compute_caldyn_fast(tau, pk, berni, theta, geopot, du_fast, u)
         CALL compute_pvort_only(rhodz, u, qv, qu)
         CALL compute_caldyn_slow_NH(u, rhodz, geopot, W, bufm1, bufm2, bufm3, bufu1, bufu2, bufu3, bufu4, &
                                     hflux, du_slow, dPhi_slow, dW_slow)
         CALL compute_coriolis(hflux, theta, qu, bufu1, drhodz, dtheta_rhodz, du_slow)
         IF (caldyn_eta == eta_mass) THEN
            CALL caldyn_vert(drhodz, rhodz, theta, u, dmass_col, wflux, dtheta_rhodz, du_slow, bufu1)
            CALL compute_caldyn_vert_NH(rhodz, geopot, W, wflux, bufm1, bufm2, bufm3, du_slow, dPhi_slow, dW_slow)
         END IF
      END IF

      time2 = OMP_GET_WTIME()
      IF (time2 > time1) elapsed = elapsed + time2 - time1

      CALL print_trace()
   END SUBROUTINE caldyn_unstructured
   !
   !----------------------------- Time stepping -------------------------------

   !
   SUBROUTINE ARK_step(nstep, mass_col, rhodz, theta_rhodz, u, geopot, w, & ! INOUT : flow state
                       theta, ps, pk, hflux, qv, &    ! OUT : diags (except surface geopot : IN)
                       dmass_col, drhodz, dtheta_rhodz, du_fast, du_slow, &
                       dPhi_fast, dPhi_slow, dW_fast, dW_slow) BIND(C, name='dynamico_ARK_step') ! OUT : tendencies
      INTEGER, VALUE :: nstep ! advance by nstep time steps
      REAL(C_DOUBLE), DIMENSION(primal_num)     :: mass_col, ps                     ! OUT,IN (if eta_mass) or OUT,UNUSED (if eta_lag)
      REAL(C_DOUBLE), DIMENSION(llm, primal_num)   :: rhodz, pk, berni                 ! IN, DIAG
      REAL(C_DOUBLE), DIMENSION(llm, primal_num, nqdyn)  :: theta_rhodz, theta               ! IN, DIAG
      REAL(C_DOUBLE), DIMENSION(llm, edge_num)      :: u, hflux, qu                       ! INOUT,DIAG*2
      REAL(C_DOUBLE), DIMENSION(llm + 1, primal_num) :: geopot, w, wflux                 ! IN, INOUT, DIAG
      REAL(C_DOUBLE), DIMENSION(llm, dual_num)      :: qv                               ! DIAG
      REAL(C_DOUBLE), DIMENSION(primal_num, max_nb_stage) :: dmass_col    ! OUT
      REAL(C_DOUBLE), DIMENSION(llm, primal_num, max_nb_stage) :: drhodz   ! OUT
      REAL(C_DOUBLE), DIMENSION(llm, primal_num, nqdyn, max_nb_stage) :: dtheta_rhodz ! OUT
      REAL(C_DOUBLE), DIMENSION(llm, edge_num, max_nb_stage) :: du_fast, du_slow ! OUT
      REAL(C_DOUBLE), DIMENSION(llm + 1, primal_num, max_nb_stage) :: &
         dPhi_fast, dPhi_slow, dW_fast, dW_slow      ! OUT
      ! buffers for fields that need to be shared among OpenMP threads
      ! vertical size is allocated as llm+1 even if only llm is needed
      REAL(C_DOUBLE), DIMENSION(llm + 1, primal_num) :: bufm1, bufm2, bufm3
      REAL(C_DOUBLE), DIMENSION(llm + 1, edge_num)     :: bufu1, bufu2, bufu3, bufu4
      REAL(C_DOUBLE)       :: time1, time2
      INTEGER :: step, stage, iq, ij

      time1 = OMP_GET_WTIME()

      !$OMP PARALLEL NUM_THREADS(nb_threads)

      !$OMP DO SCHEDULE(STATIC)
      DO ij = 1, primal_num
         wflux(1, ij) = 0.
         wflux(llm + 1, ij) = 0.
      END DO
      !$OMP END DO

      DO step = 1, nstep
         DO stage = 1, nb_stage

            CALL update_halo(transfer_primal, rhodz)
            DO iq = 1, nqdyn
               CALL update_halo(transfer_primal, theta_rhodz(:, :, iq))
            END DO

            IF (hydrostatic) THEN

               !$OMP DO SCHEDULE(STATIC)
               DO ij = 1, edge_num
                  du_fast(:, ij, stage) = 0.
                  du_slow(:, ij, stage) = 0.
               END DO
               !$OMP END DO

               CALL compute_theta(mass_col, rhodz, theta_rhodz, theta)
               CALL compute_geopot(rhodz, theta, ps, pk, geopot)

               CALL compute_caldyn_fast(tauj(stage), pk, berni, theta, geopot, du_fast(:, :, stage), u)
               CALL update_halo(transfer_edge, u)
               CALL compute_pvort_only(rhodz, u, qv, qu)
               CALL update_halo(transfer_edge, qu)

               CALL compute_caldyn_slow_hydro(rhodz, theta, u, berni, hflux, du_slow(:, :, stage))
               CALL compute_coriolis(hflux, theta, qu, bufu1, drhodz(:, :, stage), &
                                     dtheta_rhodz(:, :, :, stage), du_slow(:, :, stage))
               IF (caldyn_eta == eta_mass) THEN
                  CALL caldyn_vert(drhodz(:, :, stage), rhodz, theta, u, &
                                   dmass_col(:, stage), wflux, dtheta_rhodz(:, :, :, stage), du_slow(:, :, stage), bufu1)
               END IF

            ELSE ! NH
               CALL update_halo(transfer_primal, geopot)
               CALL update_halo(transfer_primal, W)

               CALL compute_theta(mass_col, rhodz, theta_rhodz, theta)
               CALL compute_caldyn_solver(tauj(stage), rhodz, theta, bufm1, bufm2, bufm3, pk, geopot, W, &
                                          dPhi_fast(:, :, stage), dW_fast(:, :, stage), du_fast(:, :, stage))

               CALL compute_caldyn_fast(tauj(stage), pk, berni, theta, geopot, du_fast(:, :, stage), u)
               CALL update_halo(transfer_edge, u)
               CALL compute_pvort_only(rhodz, u, qv, qu)
               CALL update_halo(transfer_edge, qu)

               CALL compute_caldyn_slow_NH(u, rhodz, geopot, W, bufm1, bufm2, bufm3, &
                                           bufu1, bufu2, bufu3, bufu4, hflux, &
                                           du_slow(:, :, stage), dPhi_slow(:, :, stage), dW_slow(:, :, stage))
               CALL compute_coriolis(hflux, theta, qu, bufu1, &
                                     drhodz(:, :, stage), dtheta_rhodz(:, :, :, stage), du_slow(:, :, stage))
               IF (caldyn_eta == eta_mass) THEN
                  CALL caldyn_vert(drhodz(:, :, stage), rhodz, theta, u, &
                                   dmass_col(:, stage), wflux, dtheta_rhodz(:, :, :, stage), du_slow(:, :, stage), bufu1)
                  CALL compute_caldyn_vert_NH(rhodz, geopot, W, wflux, bufm1, bufm2, bufm3, &
                                              du_slow(:, :, stage), dPhi_slow(:, :, stage), dW_slow(:, :, stage))
               END IF
            END IF ! NH

            ! FIXME : mass_col is computed from rhodz
            ! so that the DOFs are the same whatever caldyn_eta
            ! in DYNAMICO mass_col is prognosed rather than rhodz

            CALL update(stage, SIZE(rhodz), cslj, rhodz, drhodz)
            CALL update(stage, SIZE(theta_rhodz), cslj, theta_rhodz, dtheta_rhodz)
            CALL update(stage, SIZE(u), cslj, u, du_slow)
            CALL update(stage, SIZE(u), cflj, u, du_fast)

            IF (.NOT. hydrostatic) THEN
               CALL update(stage, SIZE(geopot), cslj, geopot, dPhi_slow)
               CALL update(stage, SIZE(geopot), cflj, geopot, dPhi_fast)
               CALL update(stage, SIZE(W), cslj, W, dW_slow)
               CALL update(stage, SIZE(W), cflj, W, dW_fast)
            END IF

         END DO
      END DO
      !$OMP END PARALLEL

      time2 = OMP_GET_WTIME()
      IF (time2 > time1) elapsed = elapsed + time2 - time1

      CALL print_trace()

   END SUBROUTINE ARK_step

   SUBROUTINE update(j, sz, clj, field, dfield)
      INTEGER :: j, sz ! stage in ARK scheme, field size
      REAL(C_DOUBLE), DIMENSION(max_nb_stage, max_nb_stage) :: clj ! modified Butcher tableau
      REAL(C_DOUBLE), DIMENSION(sz) :: field
      REAL(C_DOUBLE), DIMENSION(sz, max_nb_stage) :: dfield
      !
      INTEGER :: l, ij
      CALL enter_trace(id_update, 8*sz*(j + 1))
      !  PRINT *, clj(1:j,j)
      SELECT CASE (j)
      CASE (1)
         !$OMP DO SCHEDULE(static)
         DO ij = 1, sz
            field(ij) = field(ij) &
                        + clj(1, j)*dfield(ij, 1)
         END DO
      CASE (2)
         !$OMP DO SCHEDULE(static)
         DO ij = 1, sz
            field(ij) = field(ij) &
                        + clj(1, j)*dfield(ij, 1) &
                        + clj(2, j)*dfield(ij, 2)
         END DO
      CASE (3)
         !$OMP DO SCHEDULE(static)
         DO ij = 1, sz
            field(ij) = field(ij) &
                        + clj(1, j)*dfield(ij, 1) &
                        + clj(2, j)*dfield(ij, 2) &
                        + clj(3, j)*dfield(ij, 3)

         END DO
      CASE (4)
         !$OMP DO SCHEDULE(static)
         DO ij = 1, sz
            field(ij) = field(ij) &
                        + clj(1, j)*dfield(ij, 1) &
                        + clj(2, j)*dfield(ij, 2) &
                        + clj(3, j)*dfield(ij, 3) &
                        + clj(4, j)*dfield(ij, 4)
         END DO
      CASE (5)
         !$OMP DO SCHEDULE(static)
         DO ij = 1, sz
            field(ij) = field(ij) &
                        + clj(1, j)*dfield(ij, 1) &
                        + clj(2, j)*dfield(ij, 2) &
                        + clj(3, j)*dfield(ij, 3) &
                        + clj(4, j)*dfield(ij, 4) &
                        + clj(5, j)*dfield(ij, 5)
         END DO
      END SELECT
      CALL exit_trace()
   END SUBROUTINE update

   !---------------------------------------------- XIOS ----------------------------------------

END MODULE timestep_unstructured_mod
