MODULE advect_tracer_mod
  USE icosa
  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

  REAL(rstd), PARAMETER :: pente_max=2.0 ! for vlz

! 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 flux of q

  PUBLIC init_advect_tracer, advect_tracer

CONTAINS

  SUBROUTINE init_advect_tracer
    USE advect_mod
    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')
    CALL allocate_field(f_tangent,field_u,type_real,3, name='tangent')
    CALL allocate_field(f_gradq3d,field_t,type_real,llm,3, name='gradq3d')
    CALL allocate_field(f_cc,field_u,type_real,llm,3, name='cc')
    CALL allocate_field(f_sqrt_leng,field_t,type_real, name='sqrt_leng')
    CALL allocate_field(f_dzqw, field_t, type_real, llm, name='dzqw')
    CALL allocate_field(f_adzqw, field_t, type_real, llm, name='adzqw')
    CALL allocate_field(f_dzq, field_t, type_real, llm, name='dzq')
    CALL allocate_field(f_wq, field_t, type_real, llm+1, name='wq')
    
    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 SUBROUTINE init_advect_tracer

  SUBROUTINE advect_tracer(f_hfluxt, f_wfluxt,f_u, f_q,f_rhodz)
    USE advect_mod
    USE mpipara
    USE trace
    USE write_field
    IMPLICIT NONE
    
    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),POINTER :: q(:,:,:), normal(:,:), tangent(:,:), sqrt_leng(:), gradq3d(:,:,:), cc(:,:,:)
    REAL(rstd),POINTER :: hfluxt(:,:), wfluxt(:,:)
    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
    
     INTEGER :: ind,k
    LOGICAL,SAVE :: first=.TRUE.
!$OMP THREADPRIVATE(first)

    IF (first) THEN
      first=.FALSE.
      CALL init_message(f_u,req_e1_vect,req_u)
      CALL init_message(f_cc,req_e1_scal,req_cc)
      CALL init_message(f_wfluxt,req_i1,req_wfluxt)
      CALL init_message(f_q,req_i1,req_q)
      CALL init_message(f_rhodz,req_i1,req_rhodz)
      CALL init_message(f_gradq3d,req_i1,req_gradq3d)
    ENDIF
    
!!$OMP BARRIER

    CALL trace_start("advect_tracer") 

    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)

    CALL wait_message(req_u)
    CALL wait_message(req_wfluxt)
    CALL wait_message(req_q)
    CALL wait_message(req_rhodz)
    
    ! 1/2 vertical transport + back-trajectories
    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)
       cc      = f_cc(ind)
       u       = f_u(ind)
       q       = f_q(ind)
       rhodz   = f_rhodz(ind)
       wfluxt  = f_wfluxt(ind) 
       dzqw    = f_dzqw(ind)
       adzqw   = f_adzqw(ind)
       dzq     = f_dzq(ind)
       wq      = f_wq(ind) 

       DO k = 1, nqtot
          CALL vlz(k==nqtot,0.5, wfluxt,rhodz,q(:,:,k),1,dzqw, adzqw, dzq, wq)
       END DO

       CALL compute_backward_traj(tangent,normal,u,0.5*dt*itau_adv, cc) 

    END DO

    CALL send_message(f_cc,req_cc)


    ! horizontal transport - split in two to place transfer of gradq3d
    DO k = 1, nqtot
       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(:,:,k),sqrt_leng,gradq3d,xyz_i,xyz_v)

       END DO

       CALL send_message(f_gradq3d,req_gradq3d)
       CALL wait_message(req_cc)
       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)
          rhodz   = f_rhodz(ind)
          hfluxt  = f_hfluxt(ind) 
          gradq3d = f_gradq3d(ind)
          CALL compute_advect_horiz(k==nqtot,hfluxt,cc,gradq3d, rhodz,q(:,:,k))
       END DO
    END DO 
    
    ! 1/2 vertical transport
!!$OMP BARRIER

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       q       = f_q(ind)
       rhodz   = f_rhodz(ind)
       wfluxt  = f_wfluxt(ind) 
       dzqw    = f_dzqw(ind)
       adzqw   = f_adzqw(ind)
       dzq     = f_dzq(ind)
       wq      = f_wq(ind) 

       DO k = 1,nqtot
          CALL vlz(k==nqtot, 0.5,wfluxt,rhodz, q(:,:,k),0, dzqw, adzqw, dzq, wq)
       END DO

    END DO

    CALL trace_end("advect_tracer")

!!$OMP BARRIER

  END SUBROUTINE advect_tracer

  SUBROUTINE vlz(update_mass, fac,wfluxt,mass, q, halo, dzqw, adzqw, dzq, wq)
    !
    !     Auteurs:   P.Le Van, F.Hourdin, F.Forget, T. Dubos
    !
    !    ********************************************************************
    !     Update tracers using vertical mass flux only
    !     Van Leer scheme with minmod limiter
    !     wfluxt >0 for upward transport
    !    ********************************************************************
    USE trace
    USE omp_para
    IMPLICIT NONE
    LOGICAL, INTENT(IN)       :: update_mass
    REAL(rstd), INTENT(IN)    :: fac, wfluxt(iim*jjm,llm+1) ! vertical mass flux
    REAL(rstd), INTENT(INOUT) :: mass(iim*jjm,llm)
    REAL(rstd), INTENT(INOUT) :: q(iim*jjm,llm)
    INTEGER, INTENT(IN) :: halo

! temporary shared variable
    REAL(rstd),INTENT(INOUT) :: dzqw(iim*jjm,llm),        & ! vertical finite difference of q 
                                adzqw(iim*jjm,llm),       & ! abs(dzqw)
                                dzq(iim*jjm,llm),         & ! limited slope of q
                                wq(iim*jjm,llm+1)           ! time-integrated flux of q


    REAL(rstd) :: dzqmax, newmass, sigw, qq, w
    INTEGER :: i,ij,l,j,ijb,ije

    CALL trace_start("vlz")
     
     ijb=((jj_begin-halo)-1)*iim+ii_begin-halo
     ije = ((jj_end+halo)-1)*iim+ii_end+halo

    ! finite difference of q

     DO l=ll_beginp1,ll_end
!$SIMD
       DO ij=ijb,ije
         dzqw(ij,l)=q(ij,l)-q(ij,l-1)
         adzqw(ij,l)=abs(dzqw(ij,l))
       ENDDO
    ENDDO

!--> flush dzqw, adzqw
!$OMP BARRIER

    ! minmod-limited slope of q
    ! dzq = slope*dz, i.e. the reconstructed q varies by dzq inside level l

     DO l=ll_beginp1,ll_endm1
!$SIMD
       DO ij=ijb,ije 
         IF(dzqw(ij,l)*dzqw(ij,l+1).gt.0.) THEN
             dzq(ij,l) = 0.5*( dzqw(ij,l)+dzqw(ij,l+1) )
             dzqmax    = pente_max * min( adzqw(ij,l),adzqw(ij,l+1) )
             dzq(ij,l) = sign( min(abs(dzq(ij,l)),dzqmax) , dzq(ij,l) )  ! NB : sign(a,b)=a*sign(b)
          ELSE
             dzq(ij,l)=0.
          ENDIF
       ENDDO
    ENDDO


    ! 0 slope in top and bottom layers
    IF (is_omp_first_level) THEN
      DO ij=ijb,ije
           dzq(ij,1)=0.
      ENDDO
    ENDIF
      
    IF (is_omp_last_level) THEN
      DO ij=ijb,ije
          dzq(ij,llm)=0.
      ENDDO
    ENDIF

!---> flush dzq
!$OMP BARRIER  

    ! sigw = fraction of mass that leaves level l/l+1
    ! then amount of q leaving level l/l+1 = wq = w * qq
     DO l=ll_beginp1,ll_end
!$SIMD
       DO ij=ijb,ije
             w = fac*wfluxt(ij,l)
             IF(w>0.) THEN  ! upward transport, upwind side is at level l 
                sigw       = w/mass(ij,l-1)
                qq         = q(ij,l-1)+0.5*(1.-sigw)*dzq(ij,l-1) ! qq = q if sigw=1 , qq = q+dzq/2 if sigw=0
             ELSE           ! downward transport, upwind side is at level l+1
                sigw       = w/mass(ij,l)
                qq         = q(ij,l)-0.5*(1.+sigw)*dzq(ij,l) ! qq = q if sigw=-1 , qq = q-dzq/2 if sigw=0                
             ENDIF
             wq(ij,l) = w*qq
       ENDDO
    END DO
    ! wq = 0 at top and bottom
    IF (is_omp_first_level) THEN
       DO ij=ijb,ije
            wq(ij,1)=0.
      END DO
    ENDIF
    
    IF (is_omp_last_level) THEN
      DO ij=ijb,ije
            wq(ij,llm+1)=0.
      END DO
    ENDIF

! --> flush wq
!$OMP BARRIER


    ! update q, mass is updated only after all q's have been updated
    DO l=ll_begin,ll_end
!$SIMD
       DO ij=ijb,ije
             newmass = mass(ij,l) + fac*(wfluxt(ij,l)-wfluxt(ij,l+1))
             q(ij,l) = ( q(ij,l)*mass(ij,l) + wq(ij,l)-wq(ij,l+1) ) / newmass
             IF(update_mass) mass(ij,l)=newmass
       ENDDO
    END DO

    CALL trace_end("vlz")

  END SUBROUTINE vlz

END MODULE advect_tracer_mod