MODULE advect_tracer_mod
   USE icosa
   USE advect_mod
   IMPLICIT NONE
   PRIVATE

   TYPE(t_field), SAVE, POINTER :: f_normal(:)
   TYPE(t_field), SAVE, POINTER :: f_tangent(:)
   TYPE(t_field), SAVE, POINTER :: f_gradq3d(:)
   TYPE(t_field), SAVE, POINTER :: f_cc(:)  ! starting point of backward-trajectory (Miura approach)
   TYPE(t_field), SAVE, POINTER :: f_sqrt_leng(:)

   TYPE(t_message), SAVE :: req_u, req_cc, req_wfluxt, req_q, req_rhodz, req_gradq3d, req_qsat

   ! temporary shared variable for vlz
   TYPE(t_field), SAVE, POINTER :: f_dzqw(:)   ! vertical finite difference of q
   TYPE(t_field), SAVE, POINTER :: f_adzqw(:)  ! abs(dzqw)
   TYPE(t_field), SAVE, POINTER :: f_dzq(:)    ! limited slope of q
   TYPE(t_field), SAVE, POINTER :: f_wq(:)     ! time-integrated vertical flux of q
   TYPE(t_field), SAVE, POINTER :: f_hqflux(:) ! time-integrated horizontal flux of q
   TYPE(t_field), SAVE, POINTER :: f_massq(:)  ! mass of q
   TYPE(t_field), SAVE, POINTER :: f_pmid(:)      ! pressure at mid level
   TYPE(t_field), SAVE, POINTER :: f_temperature(:)      ! temperature
   TYPE(t_field), SAVE, POINTER :: f_qsat(:)      ! saturating vapor

   PUBLIC init_advect_tracer, advect_tracer

CONTAINS

   SUBROUTINE init_advect_tracer
      USE omp_para
      REAL(rstd), POINTER :: tangent(:, :)
      REAL(rstd), POINTER :: normal(:, :)
      REAL(rstd), POINTER :: sqrt_leng(:)
      INTEGER :: ind

      CALL allocate_field(f_normal, field_u, type_real, 3, name='normal', ondevice=.TRUE.)
      CALL allocate_field(f_tangent, field_u, type_real, 3, name='tangent', ondevice=.TRUE.)
      CALL allocate_field(f_gradq3d, field_t, type_real, llm, 3, name='gradq3d', ondevice=.TRUE.)
      CALL allocate_field(f_cc, field_u, type_real, llm, 3, name='cc', ondevice=.TRUE.)
      CALL allocate_field(f_sqrt_leng, field_t, type_real, name='sqrt_leng', ondevice=.TRUE.)
      CALL allocate_field(f_dzqw, field_t, type_real, llm, name='dzqw', ondevice=.TRUE.)
      CALL allocate_field(f_adzqw, field_t, type_real, llm, name='adzqw', ondevice=.TRUE.)
      CALL allocate_field(f_dzq, field_t, type_real, llm, name='dzq', ondevice=.TRUE.)
      CALL allocate_field(f_wq, field_t, type_real, llm + 1, nqtot + 1, name='wq', ondevice=.TRUE.)
      CALL allocate_field(f_massq, field_t, type_real, llm, nqtot + 1, name='massq', ondevice=.TRUE.)
      CALL allocate_field(f_hqflux, field_u, type_real, llm, nqtot + 1, name='hqflux', ondevice=.TRUE.)
      CALL allocate_field(f_pmid, field_t, type_real, llm+1, name='pmid', ondevice=.TRUE.)
      CALL allocate_field(f_temperature, field_t, type_real, llm, name='temperature', ondevice=.TRUE.)
      CALL allocate_field(f_qsat, field_t, type_real, llm, name='qsat', ondevice=.TRUE.)

      IF (grid_type == grid_ico) THEN
         DO ind = 1, ndomain
            IF (.NOT. assigned_domain(ind)) CYCLE
            CALL swap_dimensions(ind)
            CALL swap_geometry(ind)
            normal = f_normal(ind)
            tangent = f_tangent(ind)
            sqrt_leng = f_sqrt_leng(ind)
            IF (is_omp_level_master) CALL init_advect(normal, tangent, sqrt_leng)
         END DO
      END IF

      CALL update_device_field(f_tangent)
      CALL update_device_field(f_normal)
      CALL update_device_field(f_sqrt_leng)

   END SUBROUTINE init_advect_tracer

   SUBROUTINE advect_tracer(f_hfluxt, f_wfluxt, f_u, f_q, f_rhodz, &
                            frac, f_masst, f_qmasst, f_massfluxt, f_qfluxt, f_theta_rhodz, f_ps)
      USE omp_para
      USE trace
      USE write_field_mod
      USE tracer_icosa_mod
      USE compute_transport_mod, ONLY: compute_backward_traj
      USE pression_mod, ONLY: pression_mid
      USE theta2theta_rhodz_mod, ONLY: theta_rhodz2temperature
      USE qsat_mod, ONLY: qsat

      USE abort_mod
      TYPE(t_field), POINTER :: f_hfluxt(:)   ! time-integrated horizontal mass flux
      TYPE(t_field), POINTER :: f_wfluxt(:)   ! time-integrated vertical mass flux
      TYPE(t_field), POINTER :: f_u(:)        ! velocity (for back-trajectories)
      TYPE(t_field), POINTER :: f_q(:)        ! tracer
      TYPE(t_field), POINTER :: f_rhodz(:)    ! mass field at beginning of macro time step
      REAL(rstd), INTENT(in):: frac          ! ratio itau_adv/itau_out or 0. if not diagflux_on
      TYPE(t_field), POINTER :: f_masst(:)    ! time-integrated mass
      TYPE(t_field), POINTER :: f_qmasst(:)   ! time-integrated tracer mass
      TYPE(t_field), POINTER :: f_massfluxt(:)! time-integrated horizontal mass flux
      TYPE(t_field), POINTER :: f_qfluxt(:)   ! time-integrated horizontal tracer flux
      TYPE(t_field), POINTER :: f_theta_rhodz(:)   ! potential temperature
      TYPE(t_field), POINTER :: f_ps(:)       ! surface pressure

      REAL(rstd), POINTER :: q(:, :, :), normal(:, :), tangent(:, :), sqrt_leng(:), gradq3d(:, :, :), cc(:, :, :)
      REAL(rstd), POINTER :: hfluxt(:, :), wfluxt(:, :), masst(:, :), qmasst(:, :, :), &
                             massfluxt(:, :), qfluxt(:, :, :), hqflux(:, :, :)
      REAL(rstd), POINTER :: rhodz(:, :), u(:, :)
      ! temporary shared variable for vlz
      REAL(rstd), POINTER ::  dzqw(:, :)         ! vertical finite difference of q
      REAL(rstd), POINTER ::  adzqw(:, :)        ! abs(dzqw)
      REAL(rstd), POINTER ::  dzq(:, :)          ! limited slope of q
      REAL(rstd), POINTER ::  wq(:, :, :)           ! time-integrated flux of q
      REAL(rstd), POINTER ::  massq(:, :, :)        ! time-integrated flux of q

      INTEGER :: ind, k, nq_last
      LOGICAL, SAVE :: first = .TRUE.
      LOGICAl       :: ok_qsat=.FALSE.
!$OMP THREADPRIVATE(first)

      IF (first) THEN
         first = .FALSE.
         CALL init_message(f_u, req_e1_vect, req_u, 'req_u')
         CALL init_message(f_cc, req_e1_scal, req_cc, 'req_cc')
         CALL init_message(f_wfluxt, req_i1, req_wfluxt, 'req_wfluxt')
         CALL init_message(f_q, req_i1, req_q, 'req_q')
         CALL init_message(f_rhodz, req_i1, req_rhodz, 'req_rhodz')
         CALL init_message(f_gradq3d, req_i1, req_gradq3d, 'req_gradq3d')
         CALL init_message(f_qsat, req_i1, req_qsat, 'req_qsat')
      END IF

!!$OMP BARRIER

      IF (nqtot < 1) RETURN
      nq_last = -1

      DO k = 1, nqtot
         IF ((tracers(k)%advection_scheme == advect_vanleer .OR. tracers(k)%advection_scheme == advect_vanleer_qsat) &
         &    .AND. .NOT. tracers(k)%has_parent) nq_last = k
         IF (tracers(k)%advection_scheme == advect_vanleer_qsat) ok_qsat=.TRUE.
      END DO

      IF (nq_last < 0) RETURN
      

      CALL trace_start("advect_tracer")
      
      IF (ok_qsat) THEN
        ! qsat not computed onto GPU, to do later...
        CALL update_host_field(f_ps)
        CALL update_host_field(f_theta_rhodz)
        CALL pression_mid(f_ps, f_pmid)
        CALL theta_rhodz2temperature(f_ps, f_theta_rhodz, f_temperature)
        CALL qsat(f_temperature, f_pmid, f_qsat)
        CALL update_device_field(f_qsat)
      ENDIF  
      
      CALL send_message(f_u, req_u)
      CALL send_message(f_wfluxt, req_wfluxt)
      CALL send_message(f_q, req_q)
      CALL send_message(f_rhodz, req_rhodz)
      IF (ok_qsat) CALL send_message(f_qsat, req_qsat)

      CALL wait_message(req_u)
      CALL wait_message(req_wfluxt)
      CALL wait_message(req_q)
      CALL wait_message(req_rhodz)
      IF (ok_qsat)  CALL wait_message(req_qsat)

      ! copy mass and mass flux for tracer 0 (air)
      ! also compute back-trajectories (same for all tracers)
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         wq = f_wq(ind)
         massq = f_massq(ind)
         rhodz = f_rhodz(ind)
         wfluxt = f_wfluxt(ind)
         hfluxt = f_hfluxt(ind)
         hqflux = f_hqflux(ind)
         IF (grid_type == grid_ico) THEN
            !$acc kernels default(present) async
            wq(:, ll_begin:ll_endp1, 1) = wfluxt(:, ll_begin:ll_endp1)
            massq(:, ll_begin:ll_end, 1) = rhodz(:, ll_begin:ll_end)
            hqflux(:, ll_begin:ll_end, 1) = hfluxt(:, ll_begin:ll_end)
            !$acc end kernels
         ELSE
            wq(:, :, 1) = wfluxt(:, :)
            massq(:, :, 1) = rhodz(:, :)
            hqflux(:, :, 1) = hfluxt(:, :)
         END IF

         normal = f_normal(ind)
         tangent = f_tangent(ind)
         cc = f_cc(ind)
         u = f_u(ind)
         CALL compute_backward_traj(tangent, normal, u, 0.5*dt*itau_adv, cc, &
                                    xyz_e, de, wee, le) ! metrics terms
      END DO

      CALL send_message(f_cc, req_cc)
      CALL wait_message(req_cc)

      ! 1/2 vertical transport
      DO k = 1, nqtot
         IF (.NOT. tracers(k)%has_parent) &
            CALL recursive_vlz(k == nq_last, 0.5, k, 0, 1, f_q, f_wq, f_massq, f_dzqw, f_adzqw, f_dzq)
      END DO

      ! horizontal transport

      DO k = 1, nqtot
         IF (.NOT. tracers(k)%has_parent) THEN
            IF (tracers(k)%advection_scheme == advect_vanleer) THEN
              CALL recursive_advect_horiz(k == nq_last, frac, k, 0, f_hqflux, f_cc, f_gradq3d, f_massq, f_q, &
                                          f_qfluxt, f_sqrt_leng, f_qmasst, f_masst, f_massfluxt)
            ELSE IF (tracers(k)%advection_scheme == advect_vanleer_qsat) THEN
              CALL recursive_advect_horiz(k == nq_last, frac, k, 0, f_hqflux, f_cc, f_gradq3d, f_massq, f_q, &
                                          f_qfluxt, f_sqrt_leng, f_qmasst, f_masst, f_massfluxt, f_qsat)
            ENDIF
         ENDIF   
      END DO

      ! 1/2 vertical transport
!!$OMP BARRIER

      DO k = 1, nqtot
         IF (.NOT. tracers(k)%has_parent) &
            CALL recursive_vlz(k == nq_last, 0.5, k, 0, 0, f_q, f_wq, f_massq, f_dzqw, f_adzqw, f_dzq)

         IF (k == nq_last) THEN
            DO ind = 1, ndomain
               IF (.NOT. assigned_domain(ind)) CYCLE
               CALL swap_dimensions(ind)
               CALL swap_geometry(ind)
               massq = f_massq(ind)
               rhodz = f_rhodz(ind)

               ! update mass
               IF (grid_type == grid_ico) THEN
                  !$acc kernels default(present) async
                  rhodz(:, ll_begin:ll_end) = massq(:, ll_begin:ll_end, 1)
                  !$acc end kernels
               ELSE
                  rhodz(:, :) = massq(:, :, 1)
               END IF
            END DO
         END IF
      END DO

      CALL trace_end("advect_tracer")

!!$OMP BARRIER

   END SUBROUTINE advect_tracer

   RECURSIVE SUBROUTINE recursive_vlz(update_mass, fac, nq, parent, halo, f_q, f_wq, f_massq, f_dzqw, f_adzqw, f_dzq)
      USE tracer_icosa_mod, ONLY: tracers
      USE compute_transport_mod, ONLY: compute_vlz
      IMPLICIT NONE

      LOGICAL, INTENT(IN)      :: update_mass
      REAL(rstd), INTENT(IN)   :: fac
      INTEGER, INTENT(IN)      :: nq
      INTEGER, INTENT(IN)      :: parent
      INTEGER, INTENT(IN)      :: halo

      TYPE(t_field), POINTER :: f_q(:)
      TYPE(t_field), POINTER :: f_wq(:)
      TYPE(t_field), POINTER :: f_massq(:)
      TYPE(t_field), POINTER :: f_dzqw(:)
      TYPE(t_field), POINTER :: f_adzqw(:)    ! mass field at beginning of macro time step
      TYPE(t_field), POINTER :: f_dzq(:)    ! time-integrated mass
      REAL(rstd), POINTER :: q(:, :, :)
! temporary shared variable for vlz
      REAL(rstd), POINTER ::  dzqw(:, :)         ! vertical finite difference of q
      REAL(rstd), POINTER ::  adzqw(:, :)        ! abs(dzqw)
      REAL(rstd), POINTER ::  dzq(:, :)          ! limited slope of q
      REAL(rstd), POINTER ::  wq(:, :, :)           ! time-integrated flux of q
      REAL(rstd), POINTER ::  massq(:, :, :)        ! time-integrated flux of q

      INTEGER :: child, ind, nc

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         q = f_q(ind)
         dzqw = f_dzqw(ind)
         adzqw = f_adzqw(ind)
         dzq = f_dzq(ind)
         wq = f_wq(ind)
         massq = f_massq(ind)
         CALL compute_vlz(update_mass, fac, halo, &
                          massq(:, :, parent + 1), wq(:, :, parent + 1), &
                          massq(:, :, nq + 1), wq(:, :, nq + 1), &
                          q(:, :, nq), dzqw, adzqw, dzq)
      END DO

      DO nc = 1, tracers(nq)%nb_children
         child = tracers(nq)%children(nc)
         CALL recursive_vlz(.FALSE., fac, child, nq, halo, f_q, f_wq, f_massq, f_dzqw, f_adzqw, f_dzq)
      END DO

   END SUBROUTINE recursive_vlz

   RECURSIVE SUBROUTINE recursive_advect_horiz(update_mass, frac, nq, parent, f_hqflux, f_cc, f_gradq3d, f_massq, f_q, &
                                               f_qfluxt, f_sqrt_leng, f_qmasst, f_masst, f_massfluxt, f_qsat)
      USE tracer_icosa_mod, ONLY: tracers
      USE compute_transport_mod, ONLY: compute_gradq3d, compute_advect_horiz
      USE omp_para
      USE abort_mod
      IMPLICIT NONE

      LOGICAL, INTENT(IN)   :: update_mass
      REAL(rstd), INTENT(IN):: frac
      INTEGER, INTENT(IN)  :: nq
      INTEGER, INTENT(IN)  :: parent
      TYPE(t_field), POINTER :: f_hqflux(:)
      TYPE(t_field), POINTER :: f_cc(:)
      TYPE(t_field), POINTER :: f_gradq3d(:)
      TYPE(t_field), POINTER :: f_massq(:)
      TYPE(t_field), POINTER :: f_q(:)
      TYPE(t_field), POINTER :: f_qfluxt(:)
      TYPE(t_field), POINTER :: f_sqrt_leng(:)
      TYPE(t_field), POINTER :: f_qmasst(:)
      TYPE(t_field), POINTER :: f_masst(:)
      TYPE(t_field), POINTER :: f_massfluxt(:)
      TYPE(t_field), POINTER, OPTIONAL :: f_qsat(:)

      REAL(rstd), POINTER ::  q(:, :, :)
      REAL(rstd), POINTER ::  gradq3d(:, :, :)
      REAL(rstd), POINTER ::  sqrt_leng(:)
      REAL(rstd), POINTER ::  hqflux(:, :, :)
      REAL(rstd), POINTER ::  cc(:, :, :)
      REAL(rstd), POINTER ::  massq(:, :, :)
      REAL(rstd), POINTER ::  qfluxt(:, :, :)
      REAL(rstd), POINTER ::  qmasst(:, :, :)
      REAL(rstd), POINTER ::  masst(:, :)
      REAL(rstd), POINTER ::  massfluxt(:, :)
      REAL(rstd), POINTER ::  qsat(:, :)
      INTEGER :: ind, nc, child
      ! horizontal transport - split in two to place transfer of gradq3d

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         q = f_q(ind)
         gradq3d = f_gradq3d(ind)
         sqrt_leng = f_sqrt_leng(ind)
         CALL compute_gradq3d(q(:, :, nq), sqrt_leng, gradq3d, xyz_i, xyz_v)
      END DO

      CALL send_message(f_gradq3d, req_gradq3d)
      CALL wait_message(req_gradq3d)

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         cc = f_cc(ind)
         q = f_q(ind)
         massq = f_massq(ind)
         hqflux = f_hqflux(ind)
         qfluxt = f_qfluxt(ind)
         gradq3d = f_gradq3d(ind)
         IF (PRESENT(f_qsat)) qsat = f_qsat(ind)

         IF (frac > 0.) THEN ! accumulate mass, mass flux and tracer mass
            CALL abort_acc("frac>0.")
            qmasst = f_qmasst(ind)
            !$acc kernels default(present) async
            qmasst(:, ll_begin:ll_end, nq) = qmasst(:, ll_begin:ll_end, nq) + &
                                             frac*massq(:, ll_begin:ll_end, parent)*q(:, ll_begin:ll_end, nq)
            !$acc end kernels
            IF (update_mass) THEN
               masst = f_masst(ind)
               massfluxt = f_massfluxt(ind)
               !$acc kernels default(present) async
               masst(:, ll_begin:ll_end) = masst(:, ll_begin:ll_end) + frac*massq(:, ll_begin:ll_end, parent)
               massfluxt(:, ll_begin:ll_end) = massfluxt(:, ll_begin:ll_end) + hqflux(:, ll_begin:ll_end, parent)
               !$acc end kernels
            END IF
         END IF

         IF (PRESENT(f_qsat)) THEN
           CALL compute_advect_horiz(update_mass, frac > 0., &
                                     massq(:, :, parent + 1), hqflux(:, :, parent + 1), & ! parent mass and mass flux
                                     massq(:, :, nq + 1), hqflux(:, :, nq + 1), & ! tracer mass and mass flux
                                     q(:, :, nq), gradq3d(:, :, :), & ! tracer concentration and gradient
                                     qfluxt(:, :, nq), cc, qsat)
         ELSE
           CALL compute_advect_horiz(update_mass, frac > 0., &
                                     massq(:, :, parent + 1), hqflux(:, :, parent + 1), & ! parent mass and mass flux
                                     massq(:, :, nq + 1), hqflux(:, :, nq + 1), & ! tracer mass and mass flux
                                     q(:, :, nq), gradq3d(:, :, :), & ! tracer concentration and gradient
                                     qfluxt(:, :, nq), cc)
         ENDIF            
      END DO

      DO nc = 1, tracers(nq)%nb_children
         child = tracers(nq)%children(nc)

         CALL recursive_advect_horiz(.FALSE., frac, child, nq, f_hqflux, f_cc, f_gradq3d, f_massq, f_q, &
                                     f_qfluxt, f_sqrt_leng, f_qmasst, f_masst, f_massfluxt)
      END DO

   END SUBROUTINE recursive_advect_horiz

END MODULE advect_tracer_mod
