MODULE kinetic_mod
   IMPLICIT NONE
   PRIVATE

   PUBLIC :: kinetic, kinetic_new, gradient

CONTAINS

   SUBROUTINE kinetic(f_ue, f_Ki)
      USE icosa
      IMPLICIT NONE
      TYPE(t_field), POINTER :: f_ue(:)
      TYPE(t_field), POINTER :: f_Ki(:)

      REAL(rstd), POINTER :: ue(:, :)
      REAL(rstd), POINTER :: Ki(:, :)
      INTEGER :: ind

      CALL transfert_request(f_ue, req_e1_vect)

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ue = f_ue(ind)
         Ki = f_Ki(ind)
         CALL compute_kinetic(ue, Ki)
      END DO
   END SUBROUTINE kinetic

   SUBROUTINE kinetic_new(f_ue, f_Kv, f_Ki)
      USE icosa
      IMPLICIT NONE
      TYPE(t_field), POINTER :: f_ue(:)
      TYPE(t_field), POINTER :: f_Kv(:)
      TYPE(t_field), POINTER :: f_Ki(:)

      REAL(rstd), POINTER :: ue(:, :)
      REAL(rstd), POINTER :: Kv(:, :)
      REAL(rstd), POINTER :: Ki(:, :)
      INTEGER :: ind

      CALL transfert_request(f_ue, req_e1_vect)

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         ue = f_ue(ind)
         Kv = f_Kv(ind)
         CALL compute_kv(ue, Kv)
      END DO

      CALL transfert_request(f_Kv, req_z1_scal)

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         Kv = f_Kv(ind)
         Ki = f_Ki(ind)
         CALL compute_Ki_from_Kv(Kv, Ki)
      END DO
   END SUBROUTINE kinetic_new

   SUBROUTINE compute_kinetic(ue, Ki)
      USE icosa
      USE omp_para
      IMPLICIT NONE
      REAL(rstd), INTENT(IN) :: ue(3*iim*jjm, llm)
      REAL(rstd), INTENT(OUT) :: Ki(iim*jjm, llm)
      INTEGER :: i, j, ij, l

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

               Ki(ij, l) = 1/(4*Ai(ij))*(le(ij + u_right)*de(ij + u_right)*ue(ij + u_right, l)**2 + &
                                         le(ij + u_rup)*de(ij + u_rup)*ue(ij + u_rup, l)**2 + &
                                         le(ij + u_lup)*de(ij + u_lup)*ue(ij + u_lup, l)**2 + &
                                         le(ij + u_left)*de(ij + u_left)*ue(ij + u_left, l)**2 + &
                                         le(ij + u_ldown)*de(ij + u_ldown)*ue(ij + u_ldown, l)**2 + &
                                         le(ij + u_rdown)*de(ij + u_rdown)*ue(ij + u_rdown, l)**2)

            END DO
         END DO
      END DO
   END SUBROUTINE compute_kinetic

   SUBROUTINE compute_kv(ue, Kv)
      USE icosa
      USE omp_para
      IMPLICIT NONE
      REAL(rstd), INTENT(IN) :: ue(3*iim*jjm, llm)
      REAL(rstd), INTENT(OUT) :: Kv(2*iim*jjm, llm)
      INTEGER :: ij, l, u_up, u_down

      u_up = t_lup + u_right
      u_down = t_rdown + u_left

      DO l = ll_begin, ll_end
         Kv(:, l) = 0.
         DO ij = ij_begin, ij_end
            Kv(ij + z_up, l) = (radius**2/Av(ij + z_up))*( &
                               S1(ij, vup)*ue(ij + u_rup, l)**2 + &
                               S2(ij, vup)*ue(ij + u_lup, l)**2 + &
                               S2(ij + t_lup, vrdown)*ue(ij + u_up, l)**2)

            Kv(ij + z_down, l) = (radius**2/Av(ij + z_down))*( &
                                 S1(ij, vdown)*ue(ij + u_ldown, l)**2 + &
                                 S2(ij, vdown)*ue(ij + u_rdown, l)**2 + &
                                 S2(ij + t_rdown, vlup)*ue(ij + u_down, l)**2)
         END DO
      END DO
   END SUBROUTINE compute_kv

   SUBROUTINE compute_Ki_from_Kv(Kv, Ki)
      USE icosa
      USE omp_para
      IMPLICIT NONE
      REAL(rstd), INTENT(OUT):: Ki(iim*jjm, llm)
      REAL(rstd), INTENT(IN) :: Kv(2*iim*jjm, llm)
      INTEGER :: ij, l

      DO l = ll_begin, ll_end
         DO ij = ij_begin, ij_end
            Ki(ij, l) = Riv(ij, vup)*Kv(ij + z_up, l) + &
                        Riv(ij, vlup)*Kv(ij + z_lup, l) + &
                        Riv(ij, vldown)*Kv(ij + z_ldown, l) + &
                        Riv(ij, vdown)*Kv(ij + z_down, l) + &
                        Riv(ij, vrdown)*Kv(ij + z_rdown, l) + &
                        Riv(ij, vrup)*Kv(ij + z_rup, l)
         END DO
      END DO
   END SUBROUTINE compute_Ki_from_Kv

   SUBROUTINE gradient(f_berni, f_du)
      USE icosa
      IMPLICIT NONE
      TYPE(t_field), POINTER :: f_berni(:)
      TYPE(t_field), POINTER :: f_du(:)

      REAL(rstd), POINTER :: du(:, :)
      REAL(rstd), POINTER :: berni(:, :)
      INTEGER :: ind

      CALL transfert_request(f_du, req_e1_vect)
      CALL transfert_request(f_du, req_e1_vect)

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         berni = f_berni(ind)
         du = f_du(ind)
         CALL compute_grad(berni, du)
      END DO

   END SUBROUTINE gradient

   SUBROUTINE compute_grad(berni, du)
      USE icosa
      USE omp_para
      IMPLICIT NONE
      REAL(rstd), INTENT(IN) :: berni(iim*jjm, llm)
      REAL(rstd), INTENT(OUT) :: du(3*iim*jjm, llm)
      INTEGER :: i, j, ij, l

      DO l = ll_begin, ll_end
         DO j = jj_begin, jj_end
            DO i = ii_begin, ii_end
               ij = (j - 1)*iim + i
               du(ij + u_right, l) = ne_right*(berni(ij, l) - berni(ij + t_right, l))/de(ij + u_right)
               du(ij + u_lup, l) = ne_lup*(berni(ij, l) - berni(ij + t_lup, l))/de(ij + u_right)
               du(ij + u_ldown, l) = ne_ldown*(berni(ij, l) - berni(ij + t_ldown, l))/de(ij + u_right)
            END DO
         END DO
      END DO

   END SUBROUTINE compute_grad

END MODULE kinetic_mod
