MODULE etat0_collocated_mod
   USE icosa
   USE grid_param
   USE omp_para, ONLY: is_omp_level_master
   USE caldyn_vars_mod, ONLY: hydrostatic
   IMPLICIT NONE
   PRIVATE

   CHARACTER(len=255), SAVE :: etat0_type
!$OMP THREADPRIVATE(etat0_type)

   PUBLIC :: etat0_type, etat0_collocated

! Important notes for OpenMP
! When etat0 is called, vertical OpenMP parallelism is deactivated.
! Therefore only the omp_level_master thread must work, i.e. :
!   !$OMP BARRIER
!    DO ind=1,ndomain
!      IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
!      ...
!    END DO
!   !$OMP BARRIER
! There MUST be NO OMP BARRIER inside the DO-LOOP or any routine it calls.

CONTAINS

   SUBROUTINE etat0_collocated(f_phis, f_ps, f_mass, f_theta_rhodz, f_u, f_geopot, f_W, f_q)
      USE theta2theta_rhodz_mod
      TYPE(t_field), POINTER :: f_ps(:)
      TYPE(t_field), POINTER :: f_mass(:)
      TYPE(t_field), POINTER :: f_phis(:)
      TYPE(t_field), POINTER :: f_theta_rhodz(:)
      TYPE(t_field), POINTER :: f_u(:)
      TYPE(t_field), POINTER :: f_geopot(:)
      TYPE(t_field), POINTER :: f_W(:)
      TYPE(t_field), POINTER :: f_q(:)

      TYPE(t_field), POINTER, SAVE :: f_temp(:) ! SAVE => all threads see the same f_temp
      REAL(rstd), POINTER :: ps(:)
      REAL(rstd), POINTER :: mass(:, :)
      REAL(rstd), POINTER :: phis(:)
      REAL(rstd), POINTER :: theta_rhodz(:, :, :)
      REAL(rstd), POINTER :: temp(:, :)
      REAL(rstd), POINTER :: u(:, :)
      REAL(rstd), POINTER :: geopot(:, :)
      REAL(rstd), POINTER :: W(:, :)
      REAL(rstd), POINTER :: q(:, :, :)
      INTEGER :: ind

      CALL allocate_field(f_temp, field_t, type_real, llm, name='temp')

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ps = f_ps(ind)
         mass = f_mass(ind)
         phis = f_phis(ind)
         theta_rhodz = f_theta_rhodz(ind)
         temp = f_temp(ind)
         u = f_u(ind)
         geopot = f_geopot(ind)
         w = f_w(ind)
         q = f_q(ind)

         SELECT CASE (grid_type)
         CASE (grid_ico)
            CALL compute_etat0_collocated_hex(phis, ps, mass, theta_rhodz, u, geopot, W, q)
         CASE (grid_unst)
            CALL compute_etat0_collocated_unst(phis, ps, mass, theta_rhodz(:, :, 1), u, geopot, W, q)
         CASE DEFAULT
            STOP 'Unexpected value of grid_type encountered in etat0_collocated.'
         END SELECT

      END DO

      CALL deallocate_field(f_temp)

   END SUBROUTINE etat0_collocated

   SUBROUTINE compute_etat0_collocated_hex(phis, ps, mass, theta_rhodz, u, geopot, W, q)
      REAL(rstd), INTENT(INOUT) :: ps(iim*jjm)
      REAL(rstd), INTENT(INOUT) :: mass(iim*jjm, llm)
      REAL(rstd), INTENT(OUT) :: theta_rhodz(iim*jjm, llm, nqdyn)
      REAL(rstd), INTENT(OUT) :: phis(iim*jjm)
      REAL(rstd), INTENT(OUT) :: u(3*iim*jjm, llm)
      REAL(rstd), INTENT(OUT) :: W(iim*jjm, llm + 1)
      REAL(rstd), INTENT(OUT) :: geopot(iim*jjm, llm + 1)
      REAL(rstd), INTENT(OUT) :: q(iim*jjm, llm, nqtot)

      REAL(rstd) :: ps_e(3*iim*jjm)
      REAL(rstd) :: phis_e(3*iim*jjm)
      REAL(rstd) :: mass_e(3*iim*jjm, llm)
      REAL(rstd) :: theta_rhodz_e(3*iim*jjm, llm, nqdyn)
      REAL(rstd) :: geopot_e(3*iim*jjm, llm + 1)
      REAL(rstd) :: q_e(3*iim*jjm, llm, nqtot)

      w(:, :) = 0
      CALL compute_etat0_collocated(iim*jjm, lon_i, lat_i, phis, ps, mass, theta_rhodz, geopot, q)
      CALL compute_etat0_collocated(3*iim*jjm, lon_e, lat_e, phis_e, ps_e, mass_e, theta_rhodz_e, geopot_e, q_e, ep_e, u)

   END SUBROUTINE compute_etat0_collocated_hex

   SUBROUTINE compute_etat0_collocated_unst(phis, ps, mass, theta_rhodz, u, geopot, W, q)
      REAL(rstd), INTENT(INOUT) :: ps(primal_num)
      REAL(rstd), INTENT(INOUT) :: mass(llm, primal_num)
      REAL(rstd), INTENT(OUT) :: theta_rhodz(llm, primal_num, nqdyn)
      REAL(rstd), INTENT(OUT) :: phis(primal_num)
      REAL(rstd), INTENT(OUT) :: u(llm, edge_num)
      REAL(rstd), INTENT(OUT) :: W(llm + 1, primal_num)
      REAL(rstd), INTENT(OUT) :: geopot(llm + 1, primal_num)
      REAL(rstd), INTENT(OUT) :: q(llm, primal_num, nqtot)
      ! The 2D/3D arrays below have the shape expected by compute_etat0_collocated
      REAL(rstd) :: ps_e(edge_num)
      REAL(rstd) :: phis_e(edge_num)
      REAL(rstd) :: u_e(edge_num, llm)
      REAL(rstd) :: ep(edge_num, 3)
      REAL(rstd) :: mass_i(primal_num, llm), theta_rhodz_i(primal_num, llm, nqdyn), mass_e(edge_num, llm)
      REAL(rstd) :: geopot_i(edge_num, llm + 1), geopot_e(edge_num, llm + 1)
      REAL(rstd) :: q_i(primal_num, llm, nqtot), q_e(edge_num, llm, nqtot)
      INTEGER :: iq

      w(:, :) = 0
      ep = TRANSPOSE(ep_e)
      CALL compute_etat0_collocated(primal_num, lon_i, lat_i, phis, ps, mass_i, theta_rhodz_i, geopot_i, q_i)
      CALL compute_etat0_collocated(edge_num, lon_e, lat_e, phis_e, ps_e, mass_e, mass_e, geopot_e, q_e, ep, u_e)
      mass = TRANSPOSE(mass_i)
      geopot = TRANSPOSE(geopot_i)
      u = TRANSPOSE(u_e)
      DO iq = 1, nqtot
         q(:, :, iq) = TRANSPOSE(q_i(:, :, iq))
      END DO
      DO iq = 1, nqdyn
         theta_rhodz(:, :, iq) = TRANSPOSE(theta_rhodz_i(:, :, iq))
      END DO
   END SUBROUTINE compute_etat0_collocated_unst

   SUBROUTINE compute_etat0_collocated(ngrid, lon, lat, phis, ps, mass, theta_rhodz, geopot, q, ep, uperp)
      USE etat0_isothermal_mod, ONLY: compute_isothermal => compute_etat0
      USE etat0_jablonowsky06_mod, ONLY: compute_jablonowsky06 => compute_etat0
      USE etat0_dcmip1_mod, ONLY: compute_dcmip1 => compute_etat0
      USE etat0_dcmip2_mod, ONLY: compute_dcmip2 => compute_etat0
      USE etat0_dcmip3_mod, ONLY: compute_dcmip3 => compute_etat0
      USE etat0_dcmip4_mod, ONLY: compute_dcmip4 => compute_etat0
      USE etat0_dcmip5_mod, ONLY: compute_dcmip5 => compute_etat0
      USE etat0_bubble_mod, ONLY: compute_bubble => compute_etat0
      USE etat0_williamson_mod, ONLY: compute_w91_6 => compute_etat0
      USE etat0_temperature_mod, ONLY: compute_temperature => compute_etat0
      USE etat0_dcmip2016_baroclinic_wave_mod, ONLY: compute_dcmip2016_baroclinic_wave => compute_etat0
      USE etat0_dcmip2016_cyclone_mod, ONLY: compute_dcmip2016_cyclone => compute_etat0
      USE etat0_dcmip2016_supercell_mod, ONLY: compute_dcmip2016_supercell => compute_etat0
      USE disvert_mod, ONLY: ptop, ap, bp, mass_ak, mass_bk, mass_dak, mass_dbk
      USE thermodynamics_mod, ONLY: temperature_from_pVq, theta_from_pTq, volume_from_pCq
      USE earth_const, ONLY: thermo

      INTEGER :: ngrid
      REAL(rstd), INTENT(IN)    :: lon(ngrid), lat(ngrid)
      REAL(rstd), INTENT(INOUT) :: ps(ngrid)
      REAL(rstd), INTENT(INOUT) :: mass(ngrid, llm)
      REAL(rstd), INTENT(OUT)   :: theta_rhodz(ngrid, llm, nqdyn)
      REAL(rstd), INTENT(OUT)   :: phis(ngrid)
      REAL(rstd), INTENT(OUT)   :: geopot(ngrid, llm + 1)
      REAL(rstd), INTENT(OUT)   :: q(ngrid, llm, nqtot)
      REAL(rstd), INTENT(OUT), OPTIONAL   :: ep(ngrid, 3), uperp(ngrid, llm)

      REAL(rstd)   :: ulon(ngrid, llm)  ! zonal wind
      REAL(rstd)   :: ulat(ngrid, llm)  ! meridional wind
      REAL(rstd)   :: temp(ngrid, llm)  ! temperature
      REAL(rstd)   :: vol(ngrid, llm)   ! specific volume
      REAL(rstd)   :: pres(ngrid, llm)  ! pressure

      LOGICAL :: autoinit_mass, autoinit_NH, init_volume
      INTEGER :: ij, l, iq
      REAL(rstd) :: clat, slat, clon, slon, u3d(3), p_ik, mass_ik, v_ik
      REAL(rstd) :: q_ik, r_ik, chi, entropy, theta, log_p_preff

      autoinit_mass = .FALSE. ! Shallow-water test cases initialize mass and theta_rhodz => dot not change them later
      SELECT CASE (TRIM(etat0_type))
      CASE ('williamson91.6')
         CALL compute_w91_6(ngrid, lon, lat, phis, mass(:, 1), theta_rhodz(:, 1, 1), ulon(:, 1), ulat(:, 1))
      CASE DEFAULT
         ! The next test cases are 3D. They initialize surface pressure, from which
         ! the 3D mass field is deduced
         autoinit_mass = .TRUE.
      END SELECT

      autoinit_NH = .FALSE. ! These test cases compute geopotential, do not overvwrite it
      SELECT CASE (TRIM(etat0_type))
      CASE ('bubble')
         CALL compute_bubble(ngrid, lon, lat, phis, ps, temp, ulon, ulat, geopot, q)
      CASE DEFAULT
         ! The next test cases do not initialize geopotential.
         ! If hydrostatic, this is not needed : gepootential is diagnostic.
         ! For NH, geopotential and vertical momentum are initialized
         ! to hydrostatic geopotential and zero
         autoinit_NH = .NOT. hydrostatic
      END SELECT

      init_volume = .FALSE. ! These test cases initialize temperature
      SELECT CASE (TRIM(etat0_type))
      CASE ('isothermal')
         CALL compute_isothermal(ngrid, phis, ps, temp, ulon, ulat, q)
      CASE ('temperature_profile')
         CALL compute_temperature(ngrid, phis, ps, temp, ulon, ulat, q)
      CASE ('jablonowsky06')
         CALL compute_jablonowsky06(ngrid, lon, lat, phis, ps, temp, ulon, ulat)
      CASE ('dcmip1')
         CALL compute_dcmip1(ngrid, lon, lat, phis, ps, temp, ulon, ulat, q)
      CASE ('dcmip2_mountain', 'dcmip2_schaer_noshear', 'dcmip2_schaer_shear')
         CALL compute_dcmip2(ngrid, lon, lat, phis, ps, temp, ulon, ulat)
      CASE ('dcmip3')
         CALL compute_dcmip3(ngrid, lon, lat, phis, ps, temp, ulon, ulat, geopot, q)
         autoinit_NH = .FALSE. ! compute_dcmip3 initializes geopot
      CASE ('dcmip5')
         CALL compute_dcmip5(ngrid, lon, lat, phis, ps, temp, ulon, ulat, q)
      CASE DEFAULT
         ! The last test cases initialize specific volume, not temperature
         ! This makes them thermodynamics-agnostic
         init_volume = .TRUE.
      END SELECT

      SELECT CASE (TRIM(etat0_type))
      CASE ('dcmip4')
         CALL compute_dcmip4(ngrid, lon, lat, phis, ps, vol, ulon, ulat, q)
      CASE ('dcmip2016_baroclinic_wave')
         CALL compute_dcmip2016_baroclinic_wave(ngrid, lon, lat, phis, ps, vol, ulon, ulat, q)
      CASE ('dcmip2016_cyclone')
         CALL compute_dcmip2016_cyclone(ngrid, lon, lat, phis, ps, vol, ulon, ulat, q)
      CASE ('dcmip2016_supercell')
         CALL compute_dcmip2016_supercell(ngrid, lon, lat, phis, ps, vol, ulon, ulat, q)
      END SELECT

      IF (PRESENT(uperp)) THEN
         DO l = 1, llm
            DO ij = 1, ngrid
               clat = COS(lat(ij))
               slat = SIN(lat(ij))
               clon = COS(lon(ij))
               slon = SIN(lon(ij))
               u3d(1) = -clon*slat*ulat(ij, l) - slon*ulon(ij, l)
               u3d(2) = -slon*slat*ulat(ij, l) + clon*ulon(ij, l)
               u3d(3) = clat*ulat(ij, l)
               uperp(ij, l) = SUM(u3d*ep(ij, :))
            END DO
         END DO
      END IF

      IF (autoinit_mass) THEN

         DO l = 1, llm
            DO ij = 1, ngrid
               pres(ij, l) = ptop + mass_ak(l) + mass_bk(l)*ps(ij)
            END DO
         END DO

         IF (init_volume) CALL temperature_from_pVq(thermo, llm*ngrid, nqtot, temp, pres, vol, q)
         CALL theta_from_pTq(thermo, llm*ngrid, nqtot, temp, pres, temp, q) ! temp => theta

         DO l = 1, llm
            DO ij = 1, ngrid
               mass_ik = (mass_dak(l) + mass_dbk(l)*ps(ij))/g
               mass(ij, l) = mass_ik
               theta_rhodz(ij, l, 1) = mass_ik*temp(ij, l)
               DO iq = 2, nqdyn
                  theta_rhodz(ij, l, iq) = mass_ik*q(ij, l, iq - 1)
               END DO
            END DO
         END DO

         IF (autoinit_NH) THEN
            CALL volume_from_pCq(thermo, llm*ngrid, nqtot, vol, pres, temp, q)
            geopot(:, 1) = phis(:) ! surface geopotential
            DO l = 1, llm
               DO ij = 1, ngrid
                  geopot(ij, l + 1) = geopot(ij, l) + mass(ij, l)*vol(ij, l)*g
               END DO
            END DO
         END IF

      END IF

   END SUBROUTINE compute_etat0_collocated

END MODULE etat0_collocated_mod
