MODULE advect_tracer_mod
  USE icosa
  IMPLICIT NONE
  PRIVATE

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

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

  PUBLIC init_advect_tracer, advect_tracer

CONTAINS

  SUBROUTINE init_advect_tracer
    USE advect_mod
    REAL(rstd),POINTER :: tangent(:,:)
    REAL(rstd),POINTER :: normal(:,:)
    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')

    DO ind=1,ndomain
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       normal=f_normal(ind)
       tangent=f_tangent(ind)
       CALL init_advect(normal,tangent)
    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(:,:), gradq3d(:,:,:), cc(:,:,:)
    REAL(rstd),POINTER :: hfluxt(:,:), wfluxt(:,:)
    REAL(rstd),POINTER :: rhodz(:,:), u(:,:) 
    INTEGER :: ind,k

    CALL trace_start("advect_tracer") 

    CALL transfert_request(f_u,req_e1_vect)
!    CALL transfert_request(f_hfluxt,req_e1)  ! BUG : This (unnecessary) transfer makes the computation go wrong
    CALL transfert_request(f_wfluxt,req_i1)
    CALL transfert_request(f_q,req_i1)
    CALL transfert_request(f_rhodz,req_i1)
	
    IF (is_mpi_root) PRINT *, 'Advection scheme'

!    DO ind=1,ndomain
!       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)
!       hfluxt  = f_hfluxt(ind) 
!       wfluxt  = f_wfluxt(ind) 
!       gradq3d = f_gradq3d(ind)
!
!       ! 1/2 vertical transport
!       DO k = 1, nqtot
!          CALL vlz(k==nqtot,0.5, wfluxt,rhodz,q(:,:,k))
!       END DO
!
!       ! horizontal transport
!       CALL compute_backward_traj(tangent,normal,u,0.5*dt*itau_adv, cc) 
!       DO k = 1,nqtot
!          CALL compute_gradq3d(q(:,:,k),gradq3d)
!          CALL compute_advect_horiz(k==nqtot,hfluxt,cc,gradq3d, rhodz,q(:,:,k))
!       END DO
!
!       ! 1/2 vertical transport
!       DO k = 1,nqtot
!          CALL vlz(k==nqtot, 0.5,wfluxt,rhodz, q(:,:,k))
!       END DO
!    END DO

    ! 1/2 vertical transport + back-trajectories
    DO ind=1,ndomain
       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) 
       DO k = 1, nqtot
          CALL vlz(k==nqtot,0.5, wfluxt,rhodz,q(:,:,k))
       END DO
       CALL compute_backward_traj(tangent,normal,u,0.5*dt*itau_adv, cc) 
    END DO

    CALL transfert_request(f_q,req_i1)      ! necessary ?
    CALL transfert_request(f_rhodz,req_i1)  ! necessary ?

    ! horizontal transport - split in two to place transfer of gradq3d

    DO k = 1, nqtot
       DO ind=1,ndomain
          CALL swap_dimensions(ind)
          CALL swap_geometry(ind)
          q       = f_q(ind)
          gradq3d = f_gradq3d(ind)
          CALL compute_gradq3d(q(:,:,k),gradq3d)
       END DO

       CALL transfert_request(f_gradq3d,req_i1)

       DO ind=1,ndomain
          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 
    
    CALL transfert_request(f_q,req_i1)      ! necessary ?
    CALL transfert_request(f_rhodz,req_i1)  ! necessary ?
    
    ! 1/2 vertical transport
    DO ind=1,ndomain
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       q       = f_q(ind)
       rhodz   = f_rhodz(ind)
       wfluxt  = f_wfluxt(ind) 
       DO k = 1,nqtot
          CALL vlz(k==nqtot, 0.5,wfluxt,rhodz, q(:,:,k))
       END DO
    END DO

    CALL trace_end("advect_tracer")

  END SUBROUTINE advect_tracer

  SUBROUTINE vlz(update_mass, fac,wfluxt,mass, q)
    !
    !     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
    !    ********************************************************************
    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)

    REAL(rstd) :: dq(iim*jjm,llm), & ! increase of q
         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

    ! finite difference of q
    DO l=2,llm
       DO j=jj_begin-1,jj_end+1
          DO i=ii_begin-1,ii_end+1
             ij=(j-1)*iim+i
             dzqw(ij,l)=q(ij,l)-q(ij,l-1)
             adzqw(ij,l)=abs(dzqw(ij,l))
          ENDDO
       ENDDO
    ENDDO

    ! minmod-limited slope of q
    ! dzq = slope*dz, i.e. the reconstructed q varies by dzq inside level l
    DO l=2,llm-1
       DO j=jj_begin-1,jj_end+1
          DO i=ii_begin-1,ii_end+1
             ij=(j-1)*iim+i
             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
    ENDDO

    ! 0 slope in top and bottom layers
    DO j=jj_begin-1,jj_end+1
       DO i=ii_begin-1,ii_end+1
          ij=(j-1)*iim+i
          dzq(ij,1)=0.
          dzq(ij,llm)=0.
       ENDDO
    ENDDO

    ! sigw = fraction of mass that leaves level l/l+1
    ! then amount of q leaving level l/l+1 = wq = w * qq
    DO l = 1,llm-1
       DO j=jj_begin-1,jj_end+1
          DO i=ii_begin-1,ii_end+1
             ij=(j-1)*iim+i
             w = fac*wfluxt(ij,l+1)
             IF(w>0.) THEN  ! upward transport, upwind side is at level l 
                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
             ELSE           ! downward transport, upwind side is at level l+1
                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                
             ENDIF
             wq(ij,l+1) = w*qq
          ENDDO
       ENDDO
    END DO
    ! wq = 0 at top and bottom
    DO j=jj_begin-1,jj_end+1
       DO i=ii_begin-1,ii_end+1
          ij=(j-1)*iim+i
          wq(ij,llm+1)=0.
          wq(ij,1)=0.
       ENDDO
    END DO

    ! update q, mass is updated only after all q's have been updated
    DO l=1,llm
       DO j=jj_begin-1,jj_end+1
          DO i=ii_begin-1,ii_end+1
             ij=(j-1)*iim+i
             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
       ENDDO
    END DO

  END SUBROUTINE vlz

END MODULE advect_tracer_mod