MODULE compute_caldyn_vert_mod
  USE prec, ONLY : rstd
  USE caldyn_vars_mod
  USE grid_param
  USE disvert_mod
  USE omp_para
  USE trace
  IMPLICIT NONE
  PRIVATE
  SAVE

  PUBLIC :: compute_caldyn_vert_manual, &
       compute_caldyn_vert_hex
CONTAINS

#ifdef BEGIN_DYSL

KERNEL(caldyn_wflux)
  SEQUENCE_C0
    BODY('llm-1,1,-1')
      ! cumulate mass flux convergence from top to bottom
      convm(CELL) = convm(CELL) + convm(UP(CELL))
    END_BLOCK
    EPILOGUE(1)
      dmass_col(HIDX(CELL)) = convm(CELL)
    END_BLOCK
    BODY('2,llm')
      ! Compute vertical mass flux (l=1,llm+1 set to zero at init)
      wflux(CELL) = mass_bl(CELL) * dmass_col(HIDX(CELL)) - convm(CELL) 
    END_BLOCK
  END_BLOCK
  ! make sure wflux is up to date
  BARRIER
END_BLOCK

KERNEL(caldyn_dmass)
  FORALL_CELLS()
    ON_PRIMAL
      convm(CELL) = mass_dbk(CELL) * dmass_col(HIDX(CELL))
    END_BLOCK
  END_BLOCK
END_BLOCK

KERNEL(caldyn_vert)
  DO iq=1,nqdyn
    FORALL_CELLS('2', 'llm')
      ON_PRIMAL
        dtheta_rhodz(CELL,iq) = dtheta_rhodz(CELL,iq) + 0.5*(theta(CELL,iq)+theta(DOWN(CELL),iq))*wflux(CELL)  
      END_BLOCK
    END_BLOCK
    FORALL_CELLS('1', 'llm-1')
      ON_PRIMAL
        dtheta_rhodz(CELL,iq) = dtheta_rhodz(CELL,iq) - 0.5*(theta(CELL,iq)+theta(UP(CELL),iq))*wflux(UP(CELL))  
      END_BLOCK
    END_BLOCK
  END DO

  IF(caldyn_vert_variant == caldyn_vert_cons) THEN
    ! conservative vertical transport of momentum : (F/m)du/deta = 1/m (d/deta(Fu)-u.dF/deta)
    FORALL_CELLS('2','llm')
      ON_EDGES
        wwuu(EDGE) = .25*(wflux(CELL1)+wflux(CELL2))*(u(EDGE)+u(DOWN(EDGE))) ! Fu
      END_BLOCK
    END_BLOCK
    ! make sure wwuu is up to date
    BARRIER

    FORALL_CELLS()
      ON_EDGES
        dFu_deta = wwuu(UP(EDGE))-wwuu(EDGE)                                          ! d/deta (F*u)
	dF_deta  = .5*(wflux(UP(CELL1))+wflux(UP(CELL2))-(wflux(CELL1)+wflux(CELL2))) ! d/deta(F)
        du(EDGE) = du(EDGE) - (dFu_deta-u(EDGE)*dF_deta) / (.5*(rhodz(CELL1)+rhodz(CELL2)))  ! (F/m)du/deta
      END_BLOCK
    END_BLOCK
  ELSE  
    FORALL_CELLS('2','llm')
      ON_EDGES
        wwuu(EDGE) = .5*(wflux(CELL1)+wflux(CELL2))*(u(EDGE)-u(DOWN(EDGE)))
      END_BLOCK
    END_BLOCK

    ! make sure wwuu is up to date
    BARRIER

    FORALL_CELLS()
      ON_EDGES
        du(EDGE) = du(EDGE) - (wwuu(EDGE)+wwuu(UP(EDGE))) / (rhodz(CELL1)+rhodz(CELL2))
      END_BLOCK
    END_BLOCK
  END IF
END_BLOCK

#endif END_DYSL

  SUBROUTINE compute_caldyn_vert_hex(u,theta,rhodz,convm, wflux,wwuu, dps,dtheta_rhodz,du)
    REAL(rstd),INTENT(IN)  :: u(iim*3*jjm,llm)
    REAL(rstd),INTENT(IN)  :: theta(iim*jjm,llm,nqdyn)
    REAL(rstd),INTENT(IN)  :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(INOUT)  :: convm(iim*jjm,llm)  ! mass flux convergence
    REAL(rstd),INTENT(INOUT) :: wflux(iim*jjm,llm+1) ! vertical mass flux (kg/m2/s)
    REAL(rstd),INTENT(INOUT) :: wwuu(iim*3*jjm,llm+1)
    REAL(rstd),INTENT(INOUT) :: du(iim*3*jjm,llm)
    REAL(rstd),INTENT(INOUT) :: dtheta_rhodz(iim*jjm,llm,nqdyn)
    REAL(rstd),INTENT(OUT) :: dps(iim*jjm)

    ! temporary variable    
    INTEGER :: i,j,ij,l,iq
    REAL(rstd) :: p_ik, exner_ik, dF_deta, dFu_deta
    INTEGER    :: ij_omp_begin, ij_omp_end

    CALL trace_start("compute_caldyn_vert")

!$OMP BARRIER

    CALL distrib_level(ij_begin,ij_end, ij_omp_begin,ij_omp_end)

#define mass_bl(ij,l) bp(l)
#define dmass_col(ij) dps(ij)
#include "../kernels_hex/caldyn_wflux.k90"
#include "../kernels_hex/caldyn_vert.k90"
#undef mass_bl
#undef dmass_col

    CALL trace_end("compute_caldyn_vert")

  END SUBROUTINE compute_caldyn_vert_hex

  SUBROUTINE compute_caldyn_vert_manual(u,theta,rhodz,convm, wflux,wwuu, dps,dtheta_rhodz,du)
    REAL(rstd),INTENT(IN)  :: u(iim*3*jjm,llm)
    REAL(rstd),INTENT(IN)  :: theta(iim*jjm,llm,nqdyn)
    REAL(rstd),INTENT(IN)  :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(INOUT)  :: convm(iim*jjm,llm)  ! mass flux convergence
    REAL(rstd),INTENT(INOUT) :: wflux(iim*jjm,llm+1) ! vertical mass flux (kg/m2/s)
    REAL(rstd),INTENT(INOUT) :: wwuu(iim*3*jjm,llm+1)
    REAL(rstd),INTENT(INOUT) :: du(iim*3*jjm,llm)
    REAL(rstd),INTENT(INOUT) :: dtheta_rhodz(iim*jjm,llm,nqdyn)
    REAL(rstd),INTENT(OUT) :: dps(iim*jjm)

    ! temporary variable    
    INTEGER :: i,j,ij,l,iq
    REAL(rstd) :: p_ik, exner_ik, dF_deta, dFu_deta
    INTEGER    :: ij_omp_begin, ij_omp_end

    CALL trace_start("compute_caldyn_vert")

!$OMP BARRIER

    CALL distrib_level(ij_begin,ij_end, ij_omp_begin,ij_omp_end)

!!! cumulate mass flux convergence from top to bottom
    DO  l = llm-1, 1, -1
       !DIR$ SIMD
       DO ij=ij_omp_begin,ij_omp_end         
          convm(ij,l) = convm(ij,l) + convm(ij,l+1)
       ENDDO
    ENDDO
    !  ENDIF

    !$OMP BARRIER
    ! FLUSH on convm
    ! compute dmass_col
    IF (is_omp_first_level) THEN
       !DIR$ SIMD
       DO ij=ij_begin,ij_end         
          ! dps/dt = -int(div flux)dz
          dps(ij) = convm(ij,1)
       ENDDO
    ENDIF
!!! Compute vertical mass flux (l=1,llm+1 done by caldyn_BC)
    DO l=ll_beginp1,ll_end
       !    IF (caldyn_conserv==energy) CALL test_message(req_qu) 
       !DIR$ SIMD
       DO ij=ij_begin,ij_end         
          ! w = int(z,ztop,div(flux)dz) + B(eta)dps/dt
          ! => w>0 for upward transport
          wflux( ij, l ) = bp(l) * convm( ij, 1 ) - convm( ij, l ) 
       ENDDO
    ENDDO

    !--> flush wflux  
    !$OMP BARRIER

    DO iq=1,nqdyn
       DO l=ll_begin,ll_endm1
       !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             dtheta_rhodz(ij, l, iq) = dtheta_rhodz(ij, l, iq)  -   0.5 * &
                  ( wflux(ij,l+1) * (theta(ij,l,iq) + theta(ij,l+1,iq)))
          END DO
       END DO
       DO l=ll_beginp1,ll_end
          !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             dtheta_rhodz(ij, l, iq) = dtheta_rhodz(ij, l, iq)  +   0.5 * &
                  ( wflux(ij,l) * (theta(ij,l-1,iq) + theta(ij,l,iq) ) )
          END DO

       END DO
    END DO

    ! Compute vertical transport
    DO l=ll_beginp1,ll_end
       !DIR$ SIMD
       DO ij=ij_begin,ij_end         
          wwuu(ij+u_right,l) = 0.5*( wflux(ij,l) + wflux(ij+t_right,l)) * (u(ij+u_right,l) - u(ij+u_right,l-1))
          wwuu(ij+u_lup,l) = 0.5* ( wflux(ij,l) + wflux(ij+t_lup,l)) * (u(ij+u_lup,l) - u(ij+u_lup,l-1))
          wwuu(ij+u_ldown,l) = 0.5*( wflux(ij,l) + wflux(ij+t_ldown,l)) * (u(ij+u_ldown,l) - u(ij+u_ldown,l-1))
       ENDDO
    ENDDO

    !--> flush wwuu  
    !$OMP BARRIER

    ! Add vertical transport to du
    DO l=ll_begin,ll_end
       !DIR$ SIMD
       DO ij=ij_begin,ij_end         
          du(ij+u_right, l )   = du(ij+u_right,l)  - (wwuu(ij+u_right,l+1)+ wwuu(ij+u_right,l)) / (rhodz(ij,l)+rhodz(ij+t_right,l))
          du(ij+u_lup, l )     = du(ij+u_lup,l)    - (wwuu(ij+u_lup,l+1)  + wwuu(ij+u_lup,l))   / (rhodz(ij,l)+rhodz(ij+t_lup,l))
          du(ij+u_ldown, l )   = du(ij+u_ldown,l)  - (wwuu(ij+u_ldown,l+1)+ wwuu(ij+u_ldown,l)) / (rhodz(ij,l)+rhodz(ij+t_ldown,l))
       ENDDO
    ENDDO

    CALL trace_end("compute_caldyn_vert")

  END SUBROUTINE compute_caldyn_vert_manual

END MODULE compute_caldyn_vert_mod