MODULE kinetic_mod
IMPLICIT NONE
PRIVATE

PUBLIC :: kinetic, kinetic_v, 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)
    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)
    ENDDO  
  END SUBROUTINE kinetic
  
  SUBROUTINE kinetic_new(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)
    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_Ki_new(ue, Ki)
    ENDDO
  END SUBROUTINE kinetic_new
  
  SUBROUTINE kinetic_v(f_ue,f_Kv)
  USE icosa
  IMPLICIT NONE
    TYPE(t_field), POINTER :: f_ue(:)
    TYPE(t_field), POINTER :: f_Kv(:)
  
    REAL(rstd), POINTER :: ue(:,:)
    REAL(rstd), POINTER :: Kv(:,:)
    INTEGER :: ind
  
    CALL transfert_request(f_ue,req_e1_vect)
    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)
    ENDDO  
  END SUBROUTINE kinetic_v
  
  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 )  
       
        ENDDO
      ENDDO
    ENDDO
  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
      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 )
       ENDDO
    ENDDO
  END SUBROUTINE compute_kv
  
  SUBROUTINE compute_Ki_new(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)
    REAL(rstd) :: Kv(2*iim*jjm,llm)
    INTEGER :: ij,l, u_up, u_down
    
    CALL compute_kv(ue,Kv)
    
    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_new
  
  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)
    ENDDO
  
  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)
         ENDDO
      ENDDO
    ENDDO

  END SUBROUTINE compute_grad


END MODULE kinetic_mod