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

#ifdef BEGIN_DYSL

KERNEL(compute_omega)

  ! Pressure
  FORALL_CELLS_EXT()
    ON_PRIMAL
      p(CELL) = AP(CELL) + BP(CELL)*ps(HIDX(CELL))
    END_BLOCK
  END_BLOCK

  BARRIER

  ! Mass and grad(ps)
  FORALL_CELLS_EXT()
    ON_PRIMAL
      rhodz(CELL) = (p(CELL)-p(UP(CELL)))*(1./g)
    END_BLOCK
    ON_EDGES
      CST_IFTHEN(IS_BOTTOM_LEVEL)
        gradps(HIDX(EDGE)) = (ps(HIDX(CELL2))-ps(HIDX(CELL1)))*SIGN*LE
      CST_ENDIF
    END_BLOCK
  END_BLOCK
  
  ! Mass flux
  FORALL_CELLS_EXT()
    ON_EDGES
      Fe(EDGE)=0.5*(rhodz(CELL1)+rhodz(CELL2)*LE
    END_BLOCK
  END_BLOCK

  ! Mass flux divergence
  ! convm = +div(mass flux), sign convention as in Ringler et al. 2012, eq. 21
  FORALL_CELLS()
    ON_PRIMAL
      divflux=0.
      FORALL_EDGES
        divflux = divflux + SIGN*Fe(EDGE)
      END_BLOCK
      convm(CELL) = divflux*(1./AI)
    END_BLOCK
  END_BLOCK

  ! Barrier needed before and after doing a vertical recurrence
  BARRIER

  ! vertical integration from up to down
  SEQUENCE_C1
    PROLOGUE('llm')
      convm(CELL)=0.
    END_BLOCK
    BODY('llm-1,1,-1') 
      convm(CELL) = convm(CELL) + convm(UP(CELL))
    END_BLOCK
  END_BLOCK

  BARRIER

  ! omega = dp/dt = u.grad p + \pdiff{p}{t}
  FORALL_CELLS()
    ON_PRIMAL
      ugradps=0.
      FORALL_EDGES
        ugradps = ugradps + u(EDGE)*gradps(HIDX(EDGE))
      END_BLOCK
      ugradps = .5*(BP(CELL)+BP(UP(CELL)))*ugradps/(-4.*AI)  ! sign convention as in Ringler et al. 2010, Eq. 22 p.3072
      w(CELL) =  ugradps - g*.5*(convm(CELL)+convm(UP(CELL)))
    END_BLOCK
  END_BLOCK
END_BLOCK

#endif END_DYSL

  SUBROUTINE compute_omega(ps,u, w)
    USE disvert_mod, ONLY : ap,bp
    USE omp_para
    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)
            ENDDO
         ENDDO
      ENDDO
    
!!! 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
            ENDDO
         ENDDO
      ENDDO

      !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)
            ENDDO
         ENDDO
      ENDDO
    
!!! 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))
            ENDDO
         ENDDO
      ENDDO
    
      ! 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)
            ENDDO
         ENDDO
      ENDDO
      convm(:,llm+1)=0.

!!! Compute omega

       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)) 
          ENDDO
        ENDDO
      ENDDO
    ENDIF
!$OMP BARRIER
    
  END SUBROUTINE compute_omega
  
END MODULE compute_omega_mod