MODULE caldyn_kernels_hevi_mod
  USE icosa
  USE trace
  USE omp_para
  USE disvert_mod
  USE transfert_mod
  USE caldyn_vars_mod
  IMPLICIT NONE
  PRIVATE

  REAL(rstd), PARAMETER :: pbot=1e5, rho_bot=1e6

  LOGICAL, SAVE :: debug_hevi_solver = .FALSE.

  PUBLIC :: compute_caldyn_fast,compute_NH_geopot

CONTAINS

  SUBROUTINE compute_NH_geopot(tau, phis, m_ik, m_il, theta, W_il, Phi_il)
    REAL(rstd),INTENT(IN)    :: tau ! solve Phi-tau*dPhi/dt = Phi_rhs
    REAL(rstd),INTENT(IN)    :: phis(iim*jjm)
    REAL(rstd),INTENT(IN)    :: m_ik(iim*jjm,llm)
    REAL(rstd),INTENT(IN)    :: m_il(iim*jjm,llm+1)
    REAL(rstd),INTENT(IN)    :: theta(iim*jjm,llm)
    REAL(rstd),INTENT(IN)    :: W_il(iim*jjm,llm+1)   ! vertical momentum
    REAL(rstd),INTENT(INOUT) :: Phi_il(iim*jjm,llm+1) ! geopotential

    REAL(rstd) :: Phi_star_il(iim*jjm,llm+1) 
    REAL(rstd) :: p_ik(iim*jjm,llm)      ! pressure
    REAL(rstd) :: R_il(iim*jjm,llm+1)    ! rhs of tridiag problem
    REAL(rstd) :: x_il(iim*jjm,llm+1)    ! solution of tridiag problem
    REAL(rstd) :: A_ik(iim*jjm,llm)      ! off-diagonal coefficients of tridiag problem
    REAL(rstd) :: B_il(iim*jjm,llm+1)    ! diagonal coefficients of tridiag problem
    REAL(rstd) :: C_ik(iim*jjm,llm)      ! Thomas algorithm
    REAL(rstd) :: D_il(iim*jjm,llm+1)    ! Thomas algorithm
    REAL(rstd) :: gamma, rho_ij, X_ij, Y_ij
    REAL(rstd) :: wil, tau2_g, g2, gm2, ml_g2, c2_mik, vreff

    INTEGER    :: iter, ij, l, ij_omp_begin_ext, ij_omp_end_ext

    CALL distrib_level(ij_begin_ext,ij_end_ext, ij_omp_begin_ext,ij_omp_end_ext)

    IF(dysl) THEN
#define PHI_BOT(ij) phis(ij)
#include "../kernels_hex/compute_NH_geopot.k90"
#undef PHI_BOT
    ELSE
!    FIXME : vertical OpenMP parallelism will not work
   
    tau2_g=tau*tau/g
    g2=g*g
    gm2 = g**-2
    gamma = 1./(1.-kappa)
    
    ! compute Phi_star
    DO l=1,llm+1
       !DIR$ SIMD
       DO ij=ij_begin_ext,ij_end_ext
          Phi_star_il(ij,l) = Phi_il(ij,l) + tau*g2*(W_il(ij,l)/m_il(ij,l)-tau)
       ENDDO
    ENDDO

    ! Newton-Raphson iteration : Phi_il contains current guess value
    DO iter=1,5 ! 2 iterations should be enough

       ! Compute pressure, A_ik
       DO l=1,llm
          !DIR$ SIMD
          DO ij=ij_begin_ext,ij_end_ext
             rho_ij = (g*m_ik(ij,l))/(Phi_il(ij,l+1)-Phi_il(ij,l))
             X_ij = (cpp/preff)*kappa*theta(ij,l)*rho_ij
             p_ik(ij,l) = preff*(X_ij**gamma)
             c2_mik = gamma*p_ik(ij,l)/(rho_ij*m_ik(ij,l)) ! c^2 = gamma*R*T = gamma*p/rho
             A_ik(ij,l) = c2_mik*(tau/g*rho_ij)**2
          ENDDO
       ENDDO

       ! Compute residual, B_il
       ! bottom interface l=1
       !DIR$ SIMD
       DO ij=ij_begin_ext,ij_end_ext
          ml_g2 = gm2*m_il(ij,1) 
          B_il(ij,1) = A_ik(ij,1) + ml_g2 + tau2_g*rho_bot
          R_il(ij,1) = ml_g2*( Phi_il(ij,1)-Phi_star_il(ij,1))  &
               + tau2_g*( p_ik(ij,1)-pbot+rho_bot*(Phi_il(ij,1)-phis(ij)) )
       ENDDO
       ! inner interfaces
       DO l=2,llm
          !DIR$ SIMD
          DO ij=ij_begin_ext,ij_end_ext
             ml_g2 = gm2*m_il(ij,l) 
             B_il(ij,l) = A_ik(ij,l)+A_ik(ij,l-1) + ml_g2
             R_il(ij,l) = ml_g2*( Phi_il(ij,l)-Phi_star_il(ij,l))  &
                     + tau2_g*(p_ik(ij,l)-p_ik(ij,l-1))
             ! consistency check : if Wil=0 and initial state is in hydrostatic balance
             ! then Phi_star_il(ij,l) = Phi_il(ij,l) - tau^2*g^2
             ! and residual = tau^2*(ml+(1/g)dl_pi)=0
          ENDDO
       ENDDO
       ! top interface l=llm+1
       !DIR$ SIMD
       DO ij=ij_begin_ext,ij_end_ext
          ml_g2 = gm2*m_il(ij,llm+1) 
          B_il(ij,llm+1) = A_ik(ij,llm) + ml_g2
          R_il(ij,llm+1) = ml_g2*( Phi_il(ij,llm+1)-Phi_star_il(ij,llm+1))  &
               + tau2_g*( ptop-p_ik(ij,llm) )
       ENDDO

       ! FIXME later
       ! the lines below modify the tridiag problem 
       ! for flat, rigid boundary conditions at top and bottom :
       ! zero out A(1), A(llm), R(1), R(llm+1)
       ! => x(l)=0 at l=1,llm+1
       DO ij=ij_begin_ext,ij_end_ext
          A_ik(ij,1) = 0.
          A_ik(ij,llm) = 0.
          R_il(ij,1) = 0.
          R_il(ij,llm+1) = 0.
       ENDDO

       IF(debug_hevi_solver) THEN ! print Linf(residual)
          PRINT *, '[hevi_solver] R,p', iter, MAXVAL(ABS(R_il)), MAXVAL(p_ik)
       END IF

       ! Solve -A(l-1)x(l-1) + B(l)x(l) - A(l)x(l+1) = R(l) using Thomas algorithm
       ! Forward sweep :
       ! C(0)=0, C(l) = -A(l) / (B(l)+A(l-1)C(l-1)), 
       ! D(0)=0, D(l) = (R(l)+A(l-1)D(l-1)) / (B(l)+A(l-1)C(l-1))
       ! bottom interface l=1
       !DIR$ SIMD
       DO ij=ij_begin_ext,ij_end_ext
          X_ij = 1./B_il(ij,1)
          C_ik(ij,1) = -A_ik(ij,1) * X_ij 
          D_il(ij,1) =  R_il(ij,1) * X_ij
       ENDDO
       ! inner interfaces/layers
       DO l=2,llm
          !DIR$ SIMD
          DO ij=ij_begin_ext,ij_end_ext
             X_ij = 1./(B_il(ij,l) + A_ik(ij,l-1)*C_ik(ij,l-1))
             C_ik(ij,l) = -A_ik(ij,l) * X_ij 
             D_il(ij,l) =  (R_il(ij,l)+A_ik(ij,l-1)*D_il(ij,l-1)) * X_ij
          ENDDO
       ENDDO
       ! top interface l=llm+1
       !DIR$ SIMD
       DO ij=ij_begin_ext,ij_end_ext
          X_ij = 1./(B_il(ij,llm+1) + A_ik(ij,llm)*C_ik(ij,llm))
          D_il(ij,llm+1) =  (R_il(ij,llm+1)+A_ik(ij,llm)*D_il(ij,llm)) * X_ij
       ENDDO
       
       ! Back substitution :
       ! x(i) = D(i)-C(i)x(i+1), x(N+1)=0
       ! + Newton-Raphson update
       x_il=0. ! FIXME
       ! top interface l=llm+1
       !DIR$ SIMD
       DO ij=ij_begin_ext,ij_end_ext
          x_il(ij,llm+1) = D_il(ij,llm+1)
          Phi_il(ij,llm+1) = Phi_il(ij,llm+1) - x_il(ij,llm+1)
       ENDDO
       ! lower interfaces
       DO l=llm,1,-1
          !DIR$ SIMD
          DO ij=ij_begin_ext,ij_end_ext
             x_il(ij,l) = D_il(ij,l) - C_ik(ij,l)*x_il(ij,l+1)
             Phi_il(ij,l) = Phi_il(ij,l) - x_il(ij,l)
          ENDDO
       ENDDO

       IF(debug_hevi_solver) THEN
          PRINT *, '[hevi_solver] A,B', iter, MAXVAL(ABS(A_ik)),MAXVAL(ABS(B_il))
          PRINT *, '[hevi_solver] C,D', iter, MAXVAL(ABS(C_ik)),MAXVAL(ABS(D_il))
          DO l=1,llm+1
             WRITE(*,'(A,I2.1,I3.2,E9.2)') '[hevi_solver] x', iter,l, MAXVAL(ABS(x_il(:,l)))
          END DO
       END IF

    END DO ! Newton-Raphson

    END IF ! dysl
    
  END SUBROUTINE compute_NH_geopot

  SUBROUTINE compute_caldyn_fast(tau,u,rhodz,theta,pk,geopot,du)
    REAL(rstd),INTENT(IN)    :: tau                ! "solve" u-tau*du/dt = rhs
    REAL(rstd),INTENT(INOUT) :: u(iim*3*jjm,llm)   ! OUT if tau>0
    REAL(rstd),INTENT(IN)    :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(IN)    :: theta(iim*jjm,llm,nqdyn)
    REAL(rstd),INTENT(INOUT) :: pk(iim*jjm,llm)
    REAL(rstd),INTENT(INOUT) :: geopot(iim*jjm,llm+1)
    REAL(rstd),INTENT(INOUT)   :: du(iim*3*jjm,llm)
    REAL(rstd) :: berni(iim*jjm,llm)  ! Bernoulli function
    REAL(rstd) :: berniv(iim*jjm,llm)  ! moist Bernoulli function

    INTEGER :: i,j,ij,l
    REAL(rstd) :: cp_ik, qv, temp, chi, nu, due, due_right, due_lup, due_ldown

    CALL trace_start("compute_caldyn_fast")

    IF(dysl_caldyn_fast) THEN
#include "../kernels_hex/caldyn_fast.k90"
    ELSE

    ! Compute Bernoulli term
    IF(boussinesq) THEN
       DO l=ll_begin,ll_end
          !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             berni(ij,l) = pk(ij,l)
             ! from now on pk contains the vertically-averaged geopotential
             pk(ij,l) = .5*(geopot(ij,l)+geopot(ij,l+1))
          END DO
       END DO
    ELSE ! compressible

       DO l=ll_begin,ll_end
          SELECT CASE(caldyn_thermo)
          CASE(thermo_theta) ! vdp = theta.dpi => B = Phi
             !DIR$ SIMD
             DO ij=ij_begin,ij_end         
                berni(ij,l) = .5*(geopot(ij,l)+geopot(ij,l+1))
             END DO
          CASE(thermo_entropy) ! vdp = dG + sdT => B = Phi + G, G=h-Ts=T*(cpp-s) 
             !DIR$ SIMD
             DO ij=ij_begin,ij_end         
                berni(ij,l) = .5*(geopot(ij,l)+geopot(ij,l+1)) &
                     + pk(ij,l)*(cpp-theta(ij,l,1)) ! pk=temperature, theta=entropy
             END DO
          CASE(thermo_moist) 
             !DIR$ SIMD
             DO ij=ij_begin,ij_end         
                ! du/dt = grad(Bd)+rv.grad(Bv)+s.grad(T)
                ! Bd = Phi + gibbs_d
                ! Bv = Phi + gibbs_v
                ! pk=temperature, theta=entropy
                qv = theta(ij,l,2)
                temp = pk(ij,l)
                chi = log(temp/Treff)
                nu = (chi*(cpp+qv*cppv)-theta(ij,l,1))/(Rd+qv*Rv) ! log(p/preff)
                berni(ij,l) = .5*(geopot(ij,l)+geopot(ij,l+1)) &
                     + temp*(cpp*(1.-chi)+Rd*nu)
                berniv(ij,l) = .5*(geopot(ij,l)+geopot(ij,l+1)) &
                     + temp*(cppv*(1.-chi)+Rv*nu)
            END DO
          END SELECT
       END DO

    END IF ! Boussinesq/compressible

!!! u:=u+tau*du, du = -grad(B)-theta.grad(pi)
    DO l=ll_begin,ll_end
       IF(caldyn_thermo == thermo_moist) THEN
          !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             due_right =  berni(ij+t_right,l)-berni(ij,l)       &
                  + 0.5*(theta(ij,l,1)+theta(ij+t_right,l,1))   &
                       *(pk(ij+t_right,l)-pk(ij,l))             &
                  + 0.5*(theta(ij,l,2)+theta(ij+t_right,l,2))   &
                       *(berniv(ij+t_right,l)-berniv(ij,l))
             
             due_lup = berni(ij+t_lup,l)-berni(ij,l)            &
                  + 0.5*(theta(ij,l,1)+theta(ij+t_lup,l,1))     &
                       *(pk(ij+t_lup,l)-pk(ij,l))               &
                  + 0.5*(theta(ij,l,2)+theta(ij+t_lup,l,2))     &
                       *(berniv(ij+t_lup,l)-berniv(ij,l))
             
             due_ldown = berni(ij+t_ldown,l)-berni(ij,l)        &
                  + 0.5*(theta(ij,l,1)+theta(ij+t_ldown,l,1))   &
                       *(pk(ij+t_ldown,l)-pk(ij,l))             &
                  + 0.5*(theta(ij,l,2)+theta(ij+t_ldown,l,2))   &
                       *(berniv(ij+t_ldown,l)-berniv(ij,l))
             
             du(ij+u_right,l) = du(ij+u_right,l) - ne_right*due_right
             du(ij+u_lup,l)   = du(ij+u_lup,l)   - ne_lup*due_lup
             du(ij+u_ldown,l) = du(ij+u_ldown,l) - ne_ldown*due_ldown
             u(ij+u_right,l) = u(ij+u_right,l) + tau*du(ij+u_right,l)
             u(ij+u_lup,l)   = u(ij+u_lup,l)   + tau*du(ij+u_lup,l)
             u(ij+u_ldown,l) = u(ij+u_ldown,l) + tau*du(ij+u_ldown,l)
          END DO
       ELSE
          !DIR$ SIMD
          DO ij=ij_begin,ij_end         
             due_right =  0.5*(theta(ij,l,1)+theta(ij+t_right,l,1))  &
                  *(pk(ij+t_right,l)-pk(ij,l))        &
                  +  berni(ij+t_right,l)-berni(ij,l)
             due_lup = 0.5*(theta(ij,l,1)+theta(ij+t_lup,l,1))    &
                  *(pk(ij+t_lup,l)-pk(ij,l))          &
                  +  berni(ij+t_lup,l)-berni(ij,l)
             due_ldown = 0.5*(theta(ij,l,1)+theta(ij+t_ldown,l,1)) &
                  *(pk(ij+t_ldown,l)-pk(ij,l))      &
                  +   berni(ij+t_ldown,l)-berni(ij,l)
             du(ij+u_right,l) = du(ij+u_right,l) - ne_right*due_right
             du(ij+u_lup,l)   = du(ij+u_lup,l)   - ne_lup*due_lup
             du(ij+u_ldown,l) = du(ij+u_ldown,l) - ne_ldown*due_ldown
             u(ij+u_right,l) = u(ij+u_right,l) + tau*du(ij+u_right,l)
             u(ij+u_lup,l)   = u(ij+u_lup,l)   + tau*du(ij+u_lup,l)
             u(ij+u_ldown,l) = u(ij+u_ldown,l) + tau*du(ij+u_ldown,l)
          END DO
       END IF
    END DO

    END IF ! dysl
    CALL trace_end("compute_caldyn_fast")

  END SUBROUTINE compute_caldyn_fast

END MODULE caldyn_kernels_hevi_mod