MODULE etat0_dcmip4_mod
   USE icosa, ONLY: rstd, pi, nqtot, getin
   IMPLICIT NONE
   PRIVATE
   SAVE

   REAL(rstd), PARAMETER :: eta0 = 0.252
   REAL(rstd), PARAMETER :: etat = 0.2
   REAL(rstd), PARAMETER :: ps0 = 1e5
   REAL(rstd), PARAMETER :: u0 = 35
   REAL(rstd), PARAMETER :: T0 = 288
   REAL(rstd), PARAMETER :: DeltaT = 4.8e5
   REAL(rstd), PARAMETER :: Rd = 287
   REAL(rstd), PARAMETER :: Gamma = 0.005
   REAL(rstd), PARAMETER :: lonc = Pi/9, latc = 2*Pi/9, latw = 2*Pi/9
   REAL(rstd), PARAMETER :: pw = 34000
   REAL(rstd), PARAMETER :: q0 = 0.021

   REAL(rstd) :: up0 = 1.              ! default value, can be modified in run.def
   !$OMP THREADPRIVATE(up0)

   INTEGER :: testcase
   !$OMP THREADPRIVATE(testcase)

   PUBLIC getin_etat0, compute_etat0

CONTAINS

   SUBROUTINE getin_etat0
      USE mpipara, ONLY: is_mpi_root
      IF (nqtot < 2) THEN
         IF (is_mpi_root) THEN
            PRINT *, "nqtot must be at least 2 for test case DCMIP4"
         END IF
         STOP
      END IF
      testcase = 1
      CALL getin("dcmip4_testcase", testcase)
      CALL getin("dcmip4_up0", up0)
   END SUBROUTINE getin_etat0

   SUBROUTINE compute_etat0(ngrid, lon, lat, phis, ps, volume, ulon, ulat, q)
      USE icosa
      USE disvert_mod
      USE omp_para
      INTEGER, INTENT(IN) :: ngrid
      REAL(rstd), INTENT(IN) :: lon(ngrid)
      REAL(rstd), INTENT(IN) :: lat(ngrid)
      REAL(rstd), INTENT(OUT) :: phis(ngrid)
      REAL(rstd), INTENT(OUT) :: ps(ngrid)
      REAL(rstd), INTENT(OUT) :: volume(ngrid, llm) ! specific volume
      REAL(rstd), INTENT(OUT) :: ulon(ngrid, llm)
      REAL(rstd), INTENT(OUT) :: ulat(ngrid, llm)
      REAL(rstd), INTENT(OUT) :: q(ngrid, llm, nqtot)

      INTEGER :: l, ij
      REAL(rstd) :: etal, etavl, etas, etavs, sinlat, coslat, &
                    Y, Tave, T, phis_ave, vort, utot, &
                    dthetaodeta_ave, dthetaodeta, dthetaodlat, duodeta, K, r

      etas = ap(1)/preff + bp(1)
      etavs = (etas - eta0)*Pi/2
      phis_ave = T0*g/Gamma*(1 - etas**(Rd*Gamma/g))

      DO ij = 1, ngrid
         sinlat = SIN(lat(ij))
         coslat = COS(lat(ij))
         phis(ij) = phis_ave + u0*cos(etavs)**1.5*((-2*sinlat**6*(coslat**2 + 1./3) + 10./63)*u0*cos(etavs)**1.5 &
                                                   + (8./5*coslat**3*(sinlat**2 + 2./3) - Pi/4)*radius*Omega)
         ps(ij) = ps0
      END DO

      DO l = ll_begin, ll_end
         etal = 0.5*(ap(l)/preff + bp(l) + ap(l + 1)/preff + bp(l + 1))
         etavl = (etal - eta0)*Pi/2

         Tave = T0*etal**(Rd*Gamma/g)
         dthetaodeta_ave = T0*(Rd*Gamma/g - kappa)*etal**(Rd*Gamma/g - kappa - 1)
         IF (etat > etal) THEN
            Tave = Tave + DeltaT*(etat - etal)**5
            dthetaodeta_ave = dthetaodeta_ave - DeltaT*(5*(etat - etal)**4*etal**(-kappa) &
                                                        + kappa*(etat - etal)**5*etal**(-kappa - 1))
         END IF

         DO ij = 1, ngrid
            sinlat = SIN(lat(ij))
            coslat = COS(lat(ij))

            K = sin(latc)*sinlat + cos(latc)*coslat*cos(lon(ij) - lonc)
            r = radius*acos(K)
            utot = u0*cos(etavl)**1.5*sin(2*lat(ij))**2 + up0*exp(-(r/(0.1*radius))**2)
            ulon(ij, l) = utot
            ulat(ij, l) = 0.
            Y = ((-2*sinlat**6*(coslat**2 + 1./3) + 10./63)*2*u0*cos(etavl)**1.5 &
                 + (8./5*coslat**3*(sinlat**2 + 2./3) - Pi/4)*radius*Omega)
            T = Tave + 0.75*(etal*Pi*u0/Rd)*sin(etavl)*cos(etavl)**0.5*Y
            volume(ij, l) = Rd*T/(preff*etal) ! p=preff*etal, V=RT/p

            IF (testcase == 1) THEN
               q(ij, l, 1) = T*etal**(-kappa)
               IF (nqtot > 2) q(ij, l, 3) = 1.
               dthetaodeta = dthetaodeta_ave + 3./4.*Pi*u0/Rd*(1 - kappa)*etal**(-kappa)*sin(etavl)*cos(etavl)**0.5*Y &
                             + 3/8.*Pi**2*u0/Rd*etal**(1 - kappa)*cos(etavl)**1.5*Y &
                             - 3./16.*Pi**2*u0/Rd*etal**(1 - kappa)*sin(etavl)**2*cos(etavl)**(-0.5)*Y &
                             - 9./8.*Pi**2*u0/Rd*etal**(1 - kappa)*sin(etavl)**2*cos(etavl) &
                             *(-2*sinlat**6*(coslat**2 + 1./3.) + 10./63.)
               dthetaodlat = 3./4.*Pi*u0/Rd*etal**(1 - kappa)*sin(etavl)*cos(etavl)**0.5 &
                             *(2*u0*cos(etavl)**1.5*(-12*coslat*sinlat**5*(coslat**2 + 1./3.) + 4*coslat*sinlat**7) &
                               + radius*omega*(-24./5.*sinlat*coslat**2*(sinlat**2 + 2./3.) + 16./5.*coslat**4*sinlat))

               duodeta = -u0*sin(2*lat(ij))**2*3./4.*Pi*cos(etavl)**0.5*sin(etavl)

               vort = -4*u0/radius*cos(etavl)**1.5*sinlat*coslat*(2.-5.*sinlat**2) &
                      + up0/radius*exp(-(r/(0.1*radius))**2)*(tan(lat(ij)) - 2*(radius/(0.1*radius))**2*acos(K)*(sin(latc)*coslat &
                                                                           - cos(latc)*sinlat*cos(lon(ij) - lonc))/(sqrt(1 - K**2)))
               q(ij, l, 2) = ABS(g/preff*(-1./radius*duodeta*dthetaodlat - (2*sinlat*omega + vort)*dthetaodeta))
               IF (nqtot > 3) q(ij, l, 4) = cos(lon(ij))*coslat
            ELSE IF (testcase == 2) THEN
               q(ij, l, 1) = q0*exp(-(lat(ij)/latw)**4)*exp(-((etal - 1)*preff/pw)**2)
            END IF
         END DO
      END DO

   END SUBROUTINE compute_etat0

END MODULE etat0_dcmip4_mod
