MODULE caldyn_kernels_base_mod
  USE icosa
  USE transfert_mod
  IMPLICIT NONE
  PRIVATE

  LOGICAL, PARAMETER, PUBLIC :: DEC = .TRUE.

  INTEGER, PARAMETER,PUBLIC :: energy=1, enstrophy=2
  TYPE(t_field),POINTER,PUBLIC :: f_out_u(:), f_qu(:), f_qv(:)
  REAL(rstd),SAVE,POINTER :: out_u(:,:), p(:,:), qu(:,:)
  !$OMP THREADPRIVATE(out_u, p, qu)

  ! temporary shared variables for caldyn
  TYPE(t_field),POINTER,PUBLIC :: f_pk(:),f_wwuu(:),f_planetvel(:)

  INTEGER, PUBLIC :: caldyn_conserv
  !$OMP THREADPRIVATE(caldyn_conserv) 

  TYPE(t_message),PUBLIC :: req_ps, req_mass, req_theta_rhodz, req_u, req_qu, req_geopot, req_w

  PUBLIC :: compute_geopot, compute_caldyn_vert

CONTAINS

  !**************************** Geopotential *****************************

  SUBROUTINE compute_geopot(ps,rhodz,theta, pk,geopot) 
    USE icosa
    USE disvert_mod
    USE exner_mod
    USE trace
    USE omp_para
    IMPLICIT NONE
    REAL(rstd),INTENT(INOUT) :: ps(iim*jjm)
    REAL(rstd),INTENT(IN)    :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(IN)    :: theta(iim*jjm,llm)    ! potential temperature
    REAL(rstd),INTENT(OUT)   :: pk(iim*jjm,llm)       ! Exner function (compressible) /Lagrange multiplier (Boussinesq)
    REAL(rstd),INTENT(INOUT) :: geopot(iim*jjm,llm+1) ! geopotential

    INTEGER :: i,j,ij,l
    REAL(rstd) :: p_ik, exner_ik
    INTEGER    :: ij_omp_begin_ext, ij_omp_end_ext


    CALL trace_start("compute_geopot")

    CALL distrib_level(ij_end_ext-ij_begin_ext+1,ij_omp_begin_ext,ij_omp_end_ext)
    ij_omp_begin_ext=ij_omp_begin_ext+ij_begin_ext-1
    ij_omp_end_ext=ij_omp_end_ext+ij_begin_ext-1

    IF(caldyn_eta==eta_mass .AND. .NOT. DEC) THEN
!!! Compute exner function and geopotential
       DO l = 1,llm
          !DIR$ SIMD
          DO ij=ij_omp_begin_ext,ij_omp_end_ext
             p_ik = ptop + mass_ak(l) + mass_bk(l)*ps(ij) ! FIXME : leave ps for the moment ; change ps to Ms later
             exner_ik = cpp * (p_ik/preff) ** kappa
             pk(ij,l) = exner_ik
             ! specific volume v = kappa*theta*pi/p = dphi/g/rhodz
             geopot(ij,l+1) = geopot(ij,l) + (g*kappa)*rhodz(ij,l)*theta(ij,l)*exner_ik/p_ik
          ENDDO
       ENDDO
    ELSE 
       ! We are using DEC or a Lagrangian vertical coordinate
       ! Pressure is computed first top-down (temporarily stored in pk)
       ! Then Exner pressure and geopotential are computed bottom-up
       ! Works also when caldyn_eta=eta_mass          

       IF(boussinesq) THEN ! compute geopotential and pk=Lagrange multiplier
          ! specific volume 1 = dphi/g/rhodz
          !         IF (is_omp_level_master) THEN ! no openMP on vertical due to dependency
          DO l = 1,llm
             !DIR$ SIMD
             DO ij=ij_omp_begin_ext,ij_omp_end_ext         
                geopot(ij,l+1) = geopot(ij,l) + g*rhodz(ij,l)
             ENDDO
          ENDDO
          ! use hydrostatic balance with theta*rhodz to find pk (Lagrange multiplier=pressure) 
          ! uppermost layer
          !DIR$ SIMD
          DO ij=ij_begin_ext,ij_end_ext         
             pk(ij,llm) = ptop + (.5*g)*theta(ij,llm)*rhodz(ij,llm)
          END DO
          ! other layers
          DO l = llm-1, 1, -1
             !          !$OMP DO SCHEDULE(STATIC) 
             !DIR$ SIMD
             DO ij=ij_begin_ext,ij_end_ext         
                pk(ij,l) = pk(ij,l+1) + (.5*g)*(theta(ij,l)*rhodz(ij,l)+theta(ij,l+1)*rhodz(ij,l+1))
             END DO
          END DO
          ! now pk contains the Lagrange multiplier (pressure)
       ELSE ! non-Boussinesq, compute geopotential and Exner pressure
          ! uppermost layer

          !DIR$ SIMD
          DO ij=ij_omp_begin_ext,ij_omp_end_ext
             pk(ij,llm) = ptop + (.5*g)*rhodz(ij,llm)
          END DO
          ! other layers
          DO l = llm-1, 1, -1
             !DIR$ SIMD
             DO ij=ij_omp_begin_ext,ij_omp_end_ext         
                pk(ij,l) = pk(ij,l+1) + (.5*g)*(rhodz(ij,l)+rhodz(ij,l+1))
             END DO
          END DO
          ! surface pressure (for diagnostics)
          IF(caldyn_eta==eta_lag) THEN
             DO ij=ij_omp_begin_ext,ij_omp_end_ext         
                ps(ij) = pk(ij,1) + (.5*g)*rhodz(ij,1)
             END DO
          END IF
          ! specific volume v = kappa*theta*pi/p = dphi/g/rhodz
          DO l = 1,llm
             !DIR$ SIMD
             DO ij=ij_omp_begin_ext,ij_omp_end_ext         
                p_ik = pk(ij,l)
                exner_ik = cpp * (p_ik/preff) ** kappa
                pk(ij,l) = exner_ik
                geopot(ij,l+1) = geopot(ij,l) + (g*kappa)*rhodz(ij,l)*theta(ij,l)*exner_ik/p_ik 
             ENDDO
          ENDDO
       END IF

    END IF

    !ym flush geopot
    !$OMP BARRIER

    CALL trace_end("compute_geopot")

  END SUBROUTINE compute_geopot

  SUBROUTINE compute_caldyn_vert(u,theta,rhodz,convm, wflux,wwuu, dps,dtheta_rhodz,du)
    USE icosa
    USE disvert_mod
    USE exner_mod
    USE trace
    USE omp_para
    IMPLICIT NONE
    REAL(rstd),INTENT(IN)  :: u(iim*3*jjm,llm)
    REAL(rstd),INTENT(IN)  :: theta(iim*jjm,llm)
    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)
    REAL(rstd),INTENT(OUT) :: dps(iim*jjm)

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

    CALL trace_start("compute_caldyn_vert")

    CALL distrib_level(ij_end-ij_begin+1,ij_omp_begin,ij_omp_end)
    ij_omp_begin=ij_omp_begin+ij_begin-1
    ij_omp_end=ij_omp_end+ij_begin-1

    !    REAL(rstd) :: wwuu(iim*3*jjm,llm+1) ! tmp var, don't know why but gain 30% on the whole code in opemp
    ! need to be understood

    !    wwuu=wwuu_out
    CALL trace_start("compute_caldyn_vert")

    !$OMP BARRIER   
!!! cumulate mass flux convergence from top to bottom
    !  IF (is_omp_level_master) THEN
    DO  l = llm-1, 1, -1
       !    IF (caldyn_conserv==energy) CALL test_message(req_qu) 

!!$OMP DO SCHEDULE(STATIC) 
       !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 dps
    IF (is_omp_first_level) THEN
       !DIR$ SIMD
       DO ij=ij_begin,ij_end         
          ! dps/dt = -int(div flux)dz
          IF(DEC) THEN
             dps(ij) = convm(ij,1)
          ELSE
             dps(ij) = convm(ij,1) * g 
          END IF
       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 l=ll_begin,ll_endm1
       !DIR$ SIMD
       DO ij=ij_begin,ij_end         
          dtheta_rhodz(ij, l ) = dtheta_rhodz(ij, l )  -   0.5 * (  wflux(ij,l+1) * (theta(ij,l) + theta(ij,l+1)))
       ENDDO
    ENDDO

    DO l=ll_beginp1,ll_end
       !DIR$ SIMD
       DO ij=ij_begin,ij_end         
          dtheta_rhodz(ij, l ) = dtheta_rhodz(ij, l )  +   0.5 * ( wflux(ij,l  ) * (theta(ij,l-1) + theta(ij,l) ) )
       ENDDO
    ENDDO


    ! 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

    !  DO l=ll_beginp1,ll_end
    !!DIR$ SIMD
    !      DO ij=ij_begin,ij_end         
    !        wwuu_out(ij+u_right,l) = wwuu(ij+u_right,l)
    !        wwuu_out(ij+u_lup,l)   = wwuu(ij+u_lup,l)
    !        wwuu_out(ij+u_ldown,l) = wwuu(ij+u_ldown,l)
    !     ENDDO
    !   ENDDO

    CALL trace_end("compute_caldyn_vert")

  END SUBROUTINE compute_caldyn_vert

END MODULE caldyn_kernels_base_mod