MODULE compute_caldyn_kv_mod
  USE grid_param, ONLY : llm
  IMPLICIT NONE
  PRIVATE

  PUBLIC :: compute_caldyn_kv

CONTAINS

  SUBROUTINE compute_caldyn_kv(ue, Kv)
    USE icosa
    USE omp_para, ONLY : ll_begin, ll_end 
    REAL(rstd),INTENT(IN) :: ue(3*iim*jjm,llm)
    REAL(rstd),INTENT(OUT) :: Kv(2*iim*jjm,llm)
    REAL(rstd) :: ue2(3*iim*jjm), dem2(3*iim*jjm), r2_Av(2*iim*jjm), rad2
    INTEGER :: ij,l, u_up, u_down
    
    u_up    = t_lup + u_right
    u_down  = t_rdown + u_left
    
    rad2=radius**2

    !DIR$ SIMD
    DO ij=ij_begin_ext,ij_end_ext
       dem2(ij+u_right) = de(ij+u_right)**(-2)
       dem2(ij+u_lup)   = de(ij+u_lup)**(-2)
       dem2(ij+u_ldown) = de(ij+u_ldown)**(-2)
       r2_Av(ij+z_up)   = rad2*(1./Av(ij+z_up))
       r2_Av(ij+z_down) = rad2*(1./Av(ij+z_down))
    END DO

    DO l=ll_begin,ll_end
       ! compute squared normal component from 1-form
       !DIR$ SIMD
       DO ij=ij_begin_ext,ij_end_ext
          ue2(ij+u_right) = dem2(ij+u_right)* (ue(ij+u_right,l)**2)
          ue2(ij+u_lup)   = dem2(ij+u_lup)  * (ue(ij+u_lup,l)**2)
          ue2(ij+u_ldown) = dem2(ij+u_ldown)* (ue(ij+u_ldown,l)**2)
       END DO
       ! average squared normal component to vertices
       !DIR$ SIMD
       DO ij=ij_begin_ext,ij_end_ext
          Kv(ij+z_up,l) = r2_Av(ij+z_up)*(                       &
                               S1(ij,vup)*ue2(ij+u_rup) +        &
                               S2(ij,vup)*ue2(ij+u_lup) +        &
                               S2(ij+t_lup,vrdown)*ue2(ij+u_up))

          Kv(ij+z_down,l) = r2_Av(ij+z_down)*(                   &
                               S1(ij,vdown)*ue2(ij+u_ldown) +    &   
                               S2(ij,vdown)*ue2(ij+u_rdown) +    &
                               S2(ij+t_rdown,vlup)*ue2(ij+u_down) )
       ENDDO
    ENDDO
  END SUBROUTINE compute_caldyn_kv

END MODULE compute_caldyn_kv_mod