!======================================================================
!
! !IROUTINE: WETBULBCALC -- Calculate wet bulb temperature (C)
!
! !DESCRIPTION:
!
!   Calculates wet bulb temperature in C, given pressure in
!      temperature in K and mixing ratio in kg/kg.
!
! !INPUT:
!    nx     - index for x dimension
!    ny     - index for y dimension
!    nz     - index for z dimension
!    prs    - pressure (mb)
!    tmk    - temperature (K)
!    qvp    - water vapor mixing ratio (kg/kg)
!
! !OUTPUT:
!    twb    - Wet bulb temperature (C)
!
! !ASSUMPTIONS:
!
! !REVISION HISTORY:
!     2009-March  - Mark T. Stoelinga - from RIP4.5
!     2010-August - J. Schramm
!     2014-March - A. Jaye - modified to run with NCL and ARW wrf output
!
! !INTERFACE:
! ------------------------------------------------------------------

! NCLFORTSTART
SUBROUTINE WETBULBCALC(prs, tmk, qvp, twb, nx, ny, nz, psafile, errstat, errmsg)
    USE wrf_constants, ONLY : ALGERR, GAMMA, GAMMAMD, TLCLC1, TLCLC2, TLCLC3, &
                          EPS, TLCLC4, THTECON1, THTECON2, THTECON3

    IMPLICIT NONE

    !f2py threadsafe
    !f2py intent(in,out) :: twb

    INTEGER, INTENT(IN) :: nx, ny, nz
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: prs
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: tmk
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: qvp
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(OUT) :: twb
    CHARACTER(LEN=*), INTENT(IN) :: psafile
    INTEGER, INTENT(INOUT) :: errstat
    CHARACTER(LEN=*), INTENT(INOUT) :: errmsg

!NCLEND

    INTEGER :: i, j, k
    INTEGER :: jt, ip
    REAL(KIND=8) :: q, t, p, e, tlcl, eth
    REAL(KIND=8) :: fracip, fracip2, fracjt, fracjt2
    REAL(KIND=8), DIMENSION(150) :: PSADITHTE, PSADIPRS
    REAL(KIND=8), DIMENSION(150,150) :: PSADITMK
    REAL(KIND=8) :: tonpsadiabat

    INTEGER :: l1, h1, mid1, rang1, l2, h2, mid2, rang2
    INTEGER :: errcnt1, errcnt2, bad_i, bad_j, bad_k
    REAL(KIND=8) :: bad_p, bad_eth
    !INTEGER :: ip, ipch, jt, jtch

    errcnt1 = 0
    errcnt2 = 0
    bad_i = -1
    bad_j = -1
    bad_k = -1

    !  Before looping, set lookup table for getting temperature on
    !  a pseudoadiabat.

    CALL DLOOKUP_TABLE(PSADITHTE, PSADIPRS, PSADITMK, psafile, errstat, errmsg)

    IF (errstat .NE. 0) THEN
        RETURN
    END IF

    !$OMP PARALLEL DO COLLAPSE(3) PRIVATE (i, j, k, jt, ip, q, t, p, e, tlcl, &
    !$OMP eth, fracip, fracip2, fracjt, fracjt2, l1, h1, mid1, rang1, l2, h2, &
    !$OMP mid2, rang2, tonpsadiabat) REDUCTION(+:errcnt1, errcnt2) &
    !$OMP SCHEDULE(runtime)
    DO k=1,nz
        DO j=1,ny
            DO i=1,nx
                q = MAX(qvp(i,j,k), 1.D-15)
                t = tmk(i,j,k)
                p = prs(i,j,k)/100.
                e = q*p/(EPS + q)
                tlcl = TLCLC1/(LOG(t**TLCLC2/e) - TLCLC3) + TLCLC4
                eth = t*(1000./p)**(GAMMA*(1. + GAMMAMD*q))*&
                    EXP((THTECON1/tlcl - THTECON2)*q*(1. + THTECON3*q))

                ! Now we need to find the temperature (in K) on a moist adiabat
                ! (specified by eth in K) given pressure in hPa.  It uses a
                ! lookup table, with data that was generated by the Bolton (1980)
                ! formula for theta_e.

                ! First check if pressure is less than min pressure in lookup table.
                ! If it is, assume parcel is so dry that the given theta-e value can
                ! be interpretted as theta, and get temperature from the simple dry
                ! theta formula.

                IF (p .LE. PSADIPRS(150)) THEN
                    tonpsadiabat = eth*(p/1000.)**GAMMA
                ELSE
                    ! Otherwise, look for the given thte/prs point in the lookup table.
                    jt = -1
                    l1 = 1
                    h1 = 149
                    rang1 = h1 - l1
                    mid1 = (h1 + l1) / 2
                    DO WHILE(rang1 .GT. 1)
                        IF (eth .GE. psadithte(mid1)) THEN
                             l1 = mid1
                        ELSE
                             h1 = mid1
                        END IF
                        rang1 = h1 - l1
                        mid1 = (h1 + l1) / 2
                    END DO
                    jt = l1

                    ip = -1
                    l2 = 1
                    h2 = 149
                    rang2 = h2 - l2
                    mid2 = (h2 + l2) / 2
                    DO WHILE(rang2 .GT. 1)
                        IF (p .LE. psadiprs(mid2)) THEN
                            l2 = mid2
                        ELSE
                            h2 = mid2
                        END IF
                        rang2 = h2 - l2
                        mid2 = (h2 + l2) / 2
                    END DO
                    ip = l2

                    IF (jt .EQ. -1 .OR. ip .EQ. -1) THEN
                        errcnt1 = errcnt1 + 1
                        !$OMP CRITICAL
                        ! Only do this the first time
                        IF (bad_i .EQ. -1) THEN
                            bad_i = i
                            bad_j = j
                            bad_k = k
                            bad_p = p
                            bad_eth = eth
                        END IF
                        !$OMP END CRITICAL
                        CYCLE
                    ENDIF

                    fracjt = (eth - PSADITHTE(jt))/(PSADITHTE(jt+1) - PSADITHTE(jt))
                    fracjt2 = 1. - fracjt
                    fracip = (PSADIPRS(ip) - p)/(PSADIPRS(ip) - PSADIPRS(ip+1))
                    fracip2 = 1. - fracip


                    IF (PSADITMK(ip,jt) .GT. 1e9 .OR. PSADITMK(ip+1,jt) .GT. 1e9 .OR. &
                        PSADITMK(ip,jt+1) .GT. 1e9 .OR. PSADITMK(ip+1,jt+1) .GT. 1e9) THEN
                        errcnt2 = errcnt2 + 1
                        !$OMP CRITICAL
                        ! Only do this the first time
                        IF (bad_i .EQ. -1) THEN
                            bad_i = i
                            bad_j = j
                            bad_k = k
                            bad_p = p
                            bad_eth = eth
                        END IF
                        !$OMP END CRITICAL
                        CYCLE
                    ENDIF

                    tonpsadiabat = fracip2*fracjt2*PSADITMK(ip,jt) + &
                                   fracip*fracjt2*PSADITMK(ip+1,jt) + &
                                   fracip2*fracjt*PSADITMK(ip,jt+1) + &
                                   fracip*fracjt*PSADITMK(ip+1,jt+1)
                ENDIF

                twb(i,j,k) = tonpsadiabat

            END DO
        END DO
    END DO

    !$OMP END PARALLEL DO

    IF (errcnt1 > 0) THEN
        errstat = ALGERR
        WRITE(errmsg, *) "Outside of lookup table bounds. i=", bad_i, ",j=", bad_j, ",k=", bad_k, ",p=", bad_p, ",eth=", bad_eth
        RETURN
    END IF

    IF (errcnt2 > 0) THEN
        errstat = ALGERR
        WRITE(errmsg, *) "PRS and THTE unreasonable. i=", bad_i, ",j=", bad_j, ",k=", bad_k, ",p=", bad_p, ",eth=", bad_eth
        RETURN
    END IF

    RETURN

END SUBROUTINE WETBULBCALC


! !IROUTINE: omgcalc -- Calculate omega (dp/dt)
!
! !DESCRIPTION:
!
!   Calculate approximate omega, based on vertical velocity w (dz/dt).
!   It is approximate because it cannot take into account the vertical
!   motion of pressure surfaces.
!
! !INPUT:
!    mx - index for x dimension
!    my - index for y dimension
!    mx -  index for vertical dimension
!    qvp - water vapor mixing ratio (kg/kg)
!    tmk - temperature (K)
!    www - vertical velocity (m/s)
!    prs -  pressure (Pa)
!
! !OUTPUT:
!    omg - omega (Pa/sec)
!
! !ASSUMPTIONS:
!
! !REVISION HISTORY:
!     2009-March  - Mark T. Stoelinga - from RIP4.5
!     2010-August - J. Schramm
!     2014-March - A. Jaye - modified to run with NCL and ARW wrf output
!
! ------------------------------------------------------------------

!NCLFORTSTART
SUBROUTINE OMGCALC(qvp, tmk, www, prs, omg, mx, my, mz)
    USE wrf_constants, ONLY : G, RD, EPS

    IMPLICIT NONE

    !f2py threadsafe
    !f2py intent(in,out) :: omg

    INTEGER,INTENT(IN) :: mx, my, mz
    REAL(KIND=8), INTENT(IN), DIMENSION(mx,my,mz) :: qvp
    REAL(KIND=8), INTENT(IN), DIMENSION(mx,my,mz) :: tmk
    REAL(KIND=8), INTENT(IN), DIMENSION(mx,my,mz) :: www
    REAL(KIND=8), INTENT(IN), DIMENSION(mx,my,mz) :: prs
    REAL(KIND=8), INTENT(OUT), DIMENSION(mx,my,mz) :: omg

!NCLEND

    ! Local variables
    INTEGER :: i, j, k
    !REAL(KIND=8), PARAMETER :: GRAV=9.81, RGAS=287.04, EPS=0.622

    !$OMP PARALLEL DO COLLAPSE(3) SCHEDULE(runtime)
    DO k=1,mz
        DO j=1,my
            DO i=1,mx
                omg(i,j,k) = -G*prs(i,j,k)/&
                              (RD*((tmk(i,j,k)*(EPS + qvp(i,j,k)))/&
                              (EPS*(1. + qvp(i,j,k)))))*www(i,j,k)
            END DO
       END DO
    END DO
    !$OMP END PARALLEL DO

    RETURN

END SUBROUTINE OMGCALC

!======================================================================
!
! !IROUTINE: VIRTUAL_TEMP -- Calculate virtual temperature (K)
!
! !DESCRIPTION:
!
!   Calculates virtual temperature in K, given temperature
!      in K and mixing ratio in kg/kg.
!
! !INPUT:
!    NX     - index for x dimension
!    NY     - index for y dimension
!    NZ     - index for z dimension
!    RATMIX - water vapor mixing ratio (kg/kg)
!    TEMP   - temperature (K)
!
! !OUTPUT:
!    TV     - Virtual temperature (K)
!
! !ASSUMPTIONS:
!
! !REVISION HISTORY:
!     2009-March  - Mark T. Stoelinga - from RIP4.5
!     2010-August - J. Schramm
!     2014-March - A. Jaye - modified to run with NCL and ARW wrf output
!
! ------------------------------------------------------------------
!NCLFORTSTART
SUBROUTINE VIRTUAL_TEMP(temp, ratmix, tv, nx, ny, nz)
    USE wrf_constants, ONLY : EPS

    IMPLICIT NONE

    !f2py threadsafe
    !f2py intent(in,out) :: tv

    INTEGER, INTENT(IN) :: nx, ny, nz
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: temp
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: ratmix
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(OUT) :: tv

!NCLEND

    INTEGER :: i,j,k
    !REAL(KIND=8),PARAMETER :: EPS = 0.622D0

    !$OMP PARALLEL DO COLLAPSE(3) SCHEDULE(runtime)
    DO k=1,nz
        DO j=1,ny
            DO i=1,nx
                tv(i,j,k) = temp(i,j,k)*(EPS + ratmix(i,j,k))/(EPS*(1.D0 + ratmix(i,j,k)))
            END DO
        END DO
    END DO
    !$OMP END PARALLEL DO

    RETURN

END SUBROUTINE VIRTUAL_TEMP

