MODULE compute_omega_mod
   USE icosa
   IMPLICIT NONE
   PRIVATE

   PUBLIC :: w_omega, compute_omega

CONTAINS

   SUBROUTINE w_omega(f_ps, f_u, f_omega) ! Compute omega = Dp/Dt
      TYPE(t_field), POINTER :: f_ps(:), f_u(:), f_omega(:)
      INTEGER :: ind
      REAL(rstd), POINTER :: ps(:), u(:, :), om(:, :)
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ps = f_ps(ind)
         u = f_u(ind)
         om = f_omega(ind)
         CALL compute_omega(ps, u, om)
      END DO
   END SUBROUTINE W_omega

   SUBROUTINE compute_omega(ps, u, w)
      USE disvert_mod, ONLY: ap, bp
      USE omp_para
      IMPLICIT NONE
      REAL(rstd), INTENT(IN) :: u(iim*3*jjm, llm), ps(iim*jjm)
      REAL(rstd), INTENT(OUT):: w(iim*jjm, llm)
      REAL(rstd):: convm(iim*jjm, llm + 1)
      REAL(rstd):: p(iim*jjm, llm + 1), rhodz(iim*jjm, llm), Fe(iim*3*jjm, llm)
      REAL(rstd):: gradps(3*iim*jjm)
      REAL(rstd):: ugradps

      INTEGER :: i, j, l, ij

!$OMP BARRIER
      IF (is_omp_level_master) THEN
         DO l = 1, llm + 1
            DO j = jj_begin - 1, jj_end + 1
               DO i = ii_begin - 1, ii_end + 1
                  ij = (j - 1)*iim + i
                  p(ij, l) = ap(l) + bp(l)*ps(ij)
               END DO
            END DO
         END DO

!!! Compute mass
         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
                  rhodz(ij, l) = (p(ij, l) - p(ij, l + 1))/g
               END DO
            END DO
         END DO

         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            gradps(ij + u_right) = (ps(ij) - ps(ij + t_right))*ne_right*le(ij + u_right)
            gradps(ij + u_lup) = (ps(ij) - ps(ij + t_lup))*ne_lup*le(ij + u_lup)
            gradps(ij + u_ldown) = (ps(ij) - ps(ij + t_ldown))*ne_ldown*le(ij + u_ldown)
         END DO

!!!  Compute mass flux
         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
                  Fe(ij + u_right, l) = 0.5*(rhodz(ij, l) + rhodz(ij + t_right, l))*u(ij + u_right, l)*le(ij + u_right)
                  Fe(ij + u_lup, l) = 0.5*(rhodz(ij, l) + rhodz(ij + t_lup, l))*u(ij + u_lup, l)*le(ij + u_lup)
                  Fe(ij + u_ldown, l) = 0.5*(rhodz(ij, l) + rhodz(ij + t_ldown, l))*u(ij + u_ldown, l)*le(ij + u_ldown)
               END DO
            END DO
         END DO

!!! mass flux convergence computation

         ! horizontal convergence
         DO l = 1, llm
            DO j = jj_begin, jj_end
               DO i = ii_begin, ii_end
                  ij = (j - 1)*iim + i
                  ! convm = +div(mass flux), sign convention as in Ringler et al. 2012, eq. 21
                  convm(ij, l) = 1./Ai(ij)*(ne(ij, right)*Fe(ij + u_right, l) + &
                                            ne(ij, rup)*Fe(ij + u_rup, l) + &
                                            ne(ij, lup)*Fe(ij + u_lup, l) + &
                                            ne(ij, left)*Fe(ij + u_left, l) + &
                                            ne(ij, ldown)*Fe(ij + u_ldown, l) + &
                                            ne(ij, rdown)*Fe(ij + u_rdown, l))
               END DO
            END DO
         END DO

         ! vertical integration from up to down
         DO l = llm - 1, 1, -1
            DO j = jj_begin, jj_end
               DO i = ii_begin, ii_end
                  ij = (j - 1)*iim + i
                  convm(ij, l) = convm(ij, l) + convm(ij, l + 1)
               END DO
            END DO
         END DO
         convm(:, llm + 1) = 0.

!!! Compute dps
         !  DO j=jj_begin,jj_end
         !    DO i=ii_begin,ii_end
         !      ij=(j-1)*iim+i
         !      ! dps/dt = -int(div flux)dz
         !      dps(ij)=-convm(ij,1) * g
         !      convm(ij,llm+1)=0.
         !    ENDDO
         !  ENDDO

         ! Compute Omega = Dp/Dt
         ! with p = A(eta)+B(eta)ps
         !      Dp/Dt = dp/deta.Deta/Dt + B(eta)Dps/Dt
         !            = -mg.Deta/Dt + B.Dps/Dt
         ! By definition the mass flux through model levels is W=m.Deta/Dt with m=-1/g dp/deta
         ! therefore
         !  Dp/Dt = -g.W + B.dps/dt + Bu.grad ps
         !        = B.u.grad ps - g*convm

!!! Compute vertical flux through model layers
         !  DO l = 1,llm-1
         !    DO j=jj_begin,jj_end
         !      DO i=ii_begin,ii_end
         !        ij=(j-1)*iim+i
         !        ! w = int(z,ztop,div(flux)dz) + B(eta)dps/dt
         !        ! => w>0 for upward transport
         !        w( ij, l+1 ) = convm( ij, l+1 ) - bp(l+1) * convm( ij, 1 ) ! g.W = g.convm + B dps/dt
         !      ENDDO
         !    ENDDO
         !  ENDDO

!!! Compute omega
         ! -grad ps : ( ne(ij,ldown)*ps(ij,l) + ne(ij+t_ldown,rup)*ps(ij+t_ldown,l) ) ) / de(ij+u_ldown)

         DO l = 1, llm
            DO j = jj_begin, jj_end
               DO i = ii_begin, ii_end
                  ij = (j - 1)*iim + i
                  ugradps = &
                     u(ij + u_rup, l)*gradps(ij + u_rup) &
                     + u(ij + u_lup, l)*gradps(ij + u_lup) &
                     + u(ij + u_left, l)*gradps(ij + u_left) &
                     + u(ij + u_ldown, l)*gradps(ij + u_ldown) &
                     + u(ij + u_rdown, l)*gradps(ij + u_rdown) &
                     + u(ij + u_right, l)*gradps(ij + u_right)

                  ugradps = .5*(bp(l) + bp(l + 1))*ugradps/(-4*Ai(ij))  ! sign convention as in Ringler et al. 2010, Eq. 22 p.3072
                  w(ij, l) = ugradps - g*.5*(convm(ij, l + 1) + convm(ij, l))
               END DO
            END DO
         END DO
      END IF
!$OMP BARRIER

   END SUBROUTINE compute_omega

END MODULE compute_omega_mod
