MODULE etat0_academic_mod
   USE icosa
   IMPLICIT NONE

   PRIVATE

   PUBLIC :: etat0

CONTAINS

   SUBROUTINE test_etat0_academic
      USE icosa
      USE kinetic_mod
      TYPE(t_field), POINTER, SAVE :: f_ps(:)
      TYPE(t_field), POINTER, SAVE :: f_phis(:)
      TYPE(t_field), POINTER, SAVE :: f_theta_rhodz(:)
      TYPE(t_field), POINTER, SAVE :: f_u(:)
      TYPE(t_field), POINTER, SAVE :: f_q(:)
      TYPE(t_field), POINTER, SAVE :: f_Ki(:)
      TYPE(t_field), POINTER, SAVE :: f_temp(:)

      CALL allocate_field(f_ps, field_t, type_real)
      CALL allocate_field(f_phis, field_t, type_real)
      CALL allocate_field(f_theta_rhodz, field_t, type_real, llm)
      CALL allocate_field(f_u, field_u, type_real, llm)
      CALL allocate_field(f_Ki, field_t, type_real, llm)
      CALL allocate_field(f_temp, field_t, type_real)

      CALL etat0(f_ps, f_phis, f_theta_rhodz, f_u, f_q)

      CALL kinetic(f_u, f_Ki)

      CALL writefield('ps', f_ps)
      CALL writefield('theta', f_theta_rhodz)

      CALL writefield('Ki', f_Ki)

   END SUBROUTINE test_etat0_academic

   SUBROUTINE etat0(f_ps, f_phis, f_theta_rhodz, f_u, f_q)
      USE icosa
      TYPE(t_field), POINTER :: f_ps(:)
      TYPE(t_field), POINTER :: f_phis(:)
      TYPE(t_field), POINTER :: f_theta_rhodz(:)
      TYPE(t_field), POINTER :: f_u(:)
      TYPE(t_field), POINTER :: f_q(:)

      REAL(rstd), POINTER :: ps(:)
      REAL(rstd), POINTER :: phis(:)
      REAL(rstd), POINTER :: theta_rhodz(:, :)
      REAL(rstd), POINTER :: u(:, :)
      INTEGER :: ind

      PRINT *, 'etat0_academic needs an upgrade for 4D theta'
      PRINT *, 'STOP in etat0_academic.f90/etat0'
      STOP

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ps = f_ps(ind)
         phis = f_phis(ind)
         theta_rhodz = f_theta_rhodz(ind)
         u = f_u(ind)
         CALL compute_etat0_academic(ps, phis, theta_rhodz, u)
      END DO

   END SUBROUTINE etat0

   SUBROUTINE compute_etat0_academic(ps, phis, theta_rhodz, u)
      USE icosa
      USE disvert_mod
      USE pression_mod, ONLY: compute_pression
      USE exner_mod
      USE geopotential_mod
      USE theta2theta_rhodz_mod
      REAL(rstd), INTENT(OUT) :: ps(iim*jjm)
      REAL(rstd), INTENT(OUT) :: phis(iim*jjm)
      REAL(rstd), INTENT(OUT) :: theta_rhodz(iim*jjm, llm)
      REAL(rstd), INTENT(OUT) :: u(3*iim*jjm, llm)

      INTEGER :: i, j, l, ij
      REAL(rstd) :: r
      REAL(rstd) :: theta(iim*jjm, llm)
      REAL(rstd) :: zsig
      INTEGER    :: lsup
      REAL(rstd) :: ddsin
      REAL(rstd) :: thetarappel
      REAL(rstd) :: lat
      REAL(rstd) :: p(iim*jjm, llm + 1)
      REAL(rstd) :: phi(iim*jjm, llm)
      REAL(rstd) :: x
      REAL(rstd) :: fact(3*iim*jjm)
      REAL(rstd) :: ut(3*iim*jjm, llm)

      DO l = 1, llm
         zsig = ap(l)/preff + bp(l)
         IF (zsig .gt. 0.3) THEN
            lsup = l
            thetarappel = 1./8.*(-log(zsig) - .5)
            DO j = jj_begin - 1, jj_end + 1
               DO i = ii_begin - 1, ii_end + 1
                  ij = (j - 1)*iim + i
                  ddsin = sin(lat_i(ij)) - sin(pi/20.)
                  theta(ij, l) = 300.*(1 + 1./18.*(1.-3.*ddsin*ddsin) + thetarappel)
               END DO
            END DO
         ELSE
!   Choix isotherme au-dessus de 300 mbar
            DO j = jj_begin - 1, jj_end + 1
               DO i = ii_begin - 1, ii_end + 1
                  ij = (j - 1)*iim + i
                  theta(ij, l) = theta(ij, lsup)*(0.3/zsig)**kappa
               END DO
            END DO
         END IF
      END DO

      DO j = jj_begin - 1, jj_end + 1
         DO i = ii_begin - 1, ii_end + 1
            ij = (j - 1)*iim + i
            ps(ij) = 1.e5
            phis(ij) = 0
         END DO
      END DO

!$OMP BARRIER
      CALL compute_pression(ps, p, 1)
!$OMP BARRIER
      CALL compute_geopotential(phis, ps, theta, phi, 1)

      DO j = jj_begin - 1, jj_end + 1
         DO i = ii_begin - 1, ii_end + 1
            ij = (j - 1)*iim + i

            lat = lat_e(ij + u_right)
            IF (abs(sin(lat)) < 1.e-4) lat = 1e-4
            x = cos(lat); x = x*x; x = x*x; x = x*x
            fact(ij + u_right) = (1.-x)/(2.*omega*sin(lat))

            lat = lat_e(ij + u_lup)
            IF (abs(sin(lat)) < 1.e-4) lat = 1e-4
            x = cos(lat); x = x*x; x = x*x; x = x*x
            fact(ij + u_lup) = (1.-x)/(2.*omega*sin(lat))

            lat = lat_e(ij + u_ldown)
            IF (abs(sin(lat)) < 1.e-4) lat = 1e-4
            x = cos(lat); x = x*x; x = x*x; x = x*x
            fact(ij + u_ldown) = (1.-x)/(2.*omega*sin(lat))
         END DO
      END DO

      ! gradient ==> tangential component
      DO l = 1, llm
         DO j = jj_begin - 1, jj_end + 1
            DO i = ii_begin - 1, ii_end + 1
               ij = (j - 1)*iim + i
               ut(ij + u_right, l) = 1/de(ij + u_right)*(ne(ij, right)*phi(ij, l) + ne(ij + t_right, left)*phi(ij + t_right, l))
               ut(ij + u_lup, l) = 1/de(ij + u_lup)*(ne(ij, lup)*phi(ij, l) + ne(ij + t_lup, rdown)*phi(ij + t_lup, l))
               ut(ij + u_ldown, l) = 1/de(ij + u_ldown)*(ne(ij, ldown)*phi(ij, l) + ne(ij + t_ldown, rup)*phi(ij + t_ldown, l))
            END DO
         END DO
      END DO

! attention au signe ??
! reconstruction of perpendicular component
      DO l = 1, llm
         DO j = jj_begin, jj_end
            DO i = ii_begin, ii_end
               ij = (j - 1)*iim + i
               u(ij + u_right, l) = -fact(ij + u_right)/de(ij + u_right)* &
                                    (wee(ij + u_right, 1)*le(ij + u_rup)*ut(ij + u_rup, l) + &
                                     wee(ij + u_right, 2)*le(ij + u_lup)*ut(ij + u_lup, l) + &
                                     wee(ij + u_right, 3)*le(ij + u_left)*ut(ij + u_left, l) + &
                                     wee(ij + u_right, 4)*le(ij + u_ldown)*ut(ij + u_ldown, l) + &
                                     wee(ij + u_right, 5)*le(ij + u_rdown)*ut(ij + u_rdown, l) + &
                                     wee(ij + u_right, 6)*le(ij + t_right + u_ldown)*ut(ij + t_right + u_ldown, l) + &
                                     wee(ij + u_right, 7)*le(ij + t_right + u_rdown)*ut(ij + t_right + u_rdown, l) + &
                                     wee(ij + u_right, 8)*le(ij + t_right + u_right)*ut(ij + t_right + u_right, l) + &
                                     wee(ij + u_right, 9)*le(ij + t_right + u_rup)*ut(ij + t_right + u_rup, l) + &
                                     wee(ij + u_right, 10)*le(ij + t_right + u_lup)*ut(ij + t_right + u_lup, l))

               u(ij + u_lup, l) = -fact(ij + u_lup)/de(ij + u_lup)* &
                                  (wee(ij + u_lup, 1)*le(ij + u_left)*ut(ij + u_left, l) + &
                                   wee(ij + u_lup, 2)*le(ij + u_ldown)*ut(ij + u_ldown, l) + &
                                   wee(ij + u_lup, 3)*le(ij + u_rdown)*ut(ij + u_rdown, l) + &
                                   wee(ij + u_lup, 4)*le(ij + u_right)*ut(ij + u_right, l) + &
                                   wee(ij + u_lup, 5)*le(ij + u_rup)*ut(ij + u_rup, l) + &
                                   wee(ij + u_lup, 6)*le(ij + t_lup + u_right)*ut(ij + t_lup + u_right, l) + &
                                   wee(ij + u_lup, 7)*le(ij + t_lup + u_rup)*ut(ij + t_lup + u_rup, l) + &
                                   wee(ij + u_lup, 8)*le(ij + t_lup + u_lup)*ut(ij + t_lup + u_lup, l) + &
                                   wee(ij + u_lup, 9)*le(ij + t_lup + u_left)*ut(ij + t_lup + u_left, l) + &
                                   wee(ij + u_lup, 10)*le(ij + t_lup + u_ldown)*ut(ij + t_lup + u_ldown, l))

               u(ij + u_ldown, l) = -fact(ij + u_ldown)/de(ij + u_ldown)* &
                                    (wee(ij + u_ldown, 1)*le(ij + u_rdown)*ut(ij + u_rdown, l) + &
                                     wee(ij + u_ldown, 2)*le(ij + u_right)*ut(ij + u_right, l) + &
                                     wee(ij + u_ldown, 3)*le(ij + u_rup)*ut(ij + u_rup, l) + &
                                     wee(ij + u_ldown, 4)*le(ij + u_lup)*ut(ij + u_lup, l) + &
                                     wee(ij + u_ldown, 5)*le(ij + u_left)*ut(ij + u_left, l) + &
                                     wee(ij + u_ldown, 6)*le(ij + t_ldown + u_lup)*ut(ij + t_ldown + u_lup, l) + &
                                     wee(ij + u_ldown, 7)*le(ij + t_ldown + u_left)*ut(ij + t_ldown + u_left, l) + &
                                     wee(ij + u_ldown, 8)*le(ij + t_ldown + u_ldown)*ut(ij + t_ldown + u_ldown, l) + &
                                     wee(ij + u_ldown, 9)*le(ij + t_ldown + u_rdown)*ut(ij + t_ldown + u_rdown, l) + &
                                     wee(ij + u_ldown, 10)*le(ij + t_ldown + u_right)*ut(ij + t_ldown + u_right, l))

!           u(ij+u_right,l)=fact(ij+u_right)*ut(ij+u_right,l)
!           u(ij+u_lup,l)=fact(ij+u_lup)*ut(ij+u_lup,l)
!           u(ij+u_ldown,l)=fact(ij+u_ldown)*ut(ij+u_ldown,l)
!            ps(ij)=phi(ij,1)
            END DO
         END DO
      END DO

      DO l = 1, llm
         DO j = jj_begin, jj_end
            DO i = ii_begin, ii_end
               ij = (j - 1)*iim + i
               CALL RANDOM_NUMBER(r); r = 0
               theta(ij, l) = theta(ij, l)*(1 - 0.0005*r)
            END DO
         END DO
      END DO

      CALL compute_theta2theta_rhodz(ps, theta, theta_rhodz, 1, ondevice=.false.)

   END SUBROUTINE compute_etat0_academic

END MODULE etat0_academic_mod
