MODULE wind_mod
  USE omp_para
  USE icosa
  IMPLICIT NONE
  PRIVATE

  PUBLIC :: compute_wind_centered, compute_flux_centered, &
       compute_wind_centered_lonlat_compound, compute_wind2d_perp_from_lonlat_compound, &
       compute_wind_centered_from_lonlat_compound, compute_wind_perp_from_lonlat_compound, &
       un2ulonlat, ulonlat2un
       
CONTAINS

  SUBROUTINE un2ulonlat(f_u, f_ulon, f_ulat)
    TYPE(t_field), POINTER :: f_u(:) ! IN  : normal velocity components on edges
    TYPE(t_field), POINTER :: f_ulon(:), f_ulat(:) ! OUT : velocity reconstructed at hexagons    
    REAL(rstd),POINTER :: u(:,:),  ulon(:,:), ulat(:,:)
    INTEGER :: ind

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       u=f_u(ind)
       ulon=f_ulon(ind)
       ulat=f_ulat(ind)
       CALL compute_un2ulonlat(u,ulon, ulat)
    END DO

  END SUBROUTINE un2ulonlat

  SUBROUTINE ulonlat2un(f_ulon, f_ulat,f_u)
    TYPE(t_field), POINTER :: f_ulon(:), f_ulat(:) ! IN : velocity reconstructed at hexagons
    TYPE(t_field), POINTER :: f_u(:) ! OUT  : normal velocity components on edges
    
    REAL(rstd),POINTER :: u(:,:),  ulon(:,:), ulat(:,:)
    INTEGER :: ind

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       u=f_u(ind)
       ulon=f_ulon(ind)
       ulat=f_ulat(ind)
       CALL compute_ulonlat2un(ulon, ulat,u)
    END DO

  END SUBROUTINE ulonlat2un
 
  SUBROUTINE compute_wind_centered(ue,ucenter)
  REAL(rstd) :: ue(3*iim*jjm,llm)
  REAL(rstd) :: ucenter(iim*jjm,llm,3)
  INTEGER :: ij,l
  REAL(rstd), PARAMETER :: scale=1.
  REAL(rstd) :: fac, ue_le, cx,cy,cz, ux,uy,uz
#include "../kernels_hex/wind_centered.k90"
 END SUBROUTINE compute_wind_centered
 
  SUBROUTINE compute_flux_centered(scale,ue,ucenter)
  REAL(rstd), INTENT(IN) :: scale
  REAL(rstd) :: ue(3*iim*jjm,llm)
  REAL(rstd) :: ucenter(iim*jjm,llm,3)
  INTEGER :: ij,l
  REAL(rstd) :: fac, ue_le, cx,cy,cz, ux,uy,uz
#include "../kernels_hex/flux_centered.k90"
  END SUBROUTINE compute_flux_centered
 
  
 SUBROUTINE compute_wind_on_edge(ue,uedge)
  REAL(rstd) :: ue(3*iim*jjm,llm)
  REAL(rstd) :: uedge(3*iim*jjm,llm,3)

  REAL(rstd) :: ut(3*iim*jjm,llm)
  INTEGER :: i,j,ij,l     
    
    CALL compute_tangential_compound(ue,ut)
  
    DO l=ll_begin,ll_end
      DO j=jj_begin,jj_end
        DO i=ii_begin,ii_end
          ij=(j-1)*iim+i
          uedge(ij+u_right,l,:)=ue(ij+u_right,l)*ep_e(ij+u_right,:)*ne(ij,right) + ut(ij+u_right,l)*et_e(ij+u_right,:)*ne(ij,right) 
          uedge(ij+u_lup,l,:)=ue(ij+u_lup,l)*ep_e(ij+u_lup,:)*ne(ij,lup) + ut(ij+u_lup,l)*et_e(ij+u_lup,:)*ne(ij,lup)
          uedge(ij+u_ldown,l,:)=ue(ij+u_ldown,l)*ep_e(ij+u_ldown,:)*ne(ij,ldown) + ut(ij+u_ldown,l)*et_e(ij+u_ldown,:)*ne(ij,ldown)
        ENDDO
      ENDDO
    ENDDO
 
 END SUBROUTINE compute_wind_on_edge
 
 
 
 SUBROUTINE compute_tangential_compound(ue,ut)
  REAL(rstd) :: ue(3*iim*jjm,llm)
  REAL(rstd) :: ut(3*iim*jjm,llm)
  INTEGER :: i,j,l,ij
    
  DO l=ll_begin,ll_end
    DO j=jj_begin,jj_end
      DO i=ii_begin,ii_end
        ij=(j-1)*iim+i
   
        ut(ij+u_right,l) = 1/de(ij+u_right) *                                            &  
                         ( wee(ij+u_right,1,1)*ue(ij+u_rup,l)+                           &
                           wee(ij+u_right,2,1)*ue(ij+u_lup,l)+                           &
                           wee(ij+u_right,3,1)*ue(ij+u_left,l)+                          &
                           wee(ij+u_right,4,1)*ue(ij+u_ldown,l)+                         &
                           wee(ij+u_right,5,1)*ue(ij+u_rdown,l)+                         & 
                           wee(ij+u_right,1,2)*ue(ij+t_right+u_ldown,l)+                 &
                           wee(ij+u_right,2,2)*ue(ij+t_right+u_rdown,l)+                 &
                           wee(ij+u_right,3,2)*ue(ij+t_right+u_right,l)+                 &
                           wee(ij+u_right,4,2)*ue(ij+t_right+u_rup,l)+                   &
                           wee(ij+u_right,5,2)*ue(ij+t_right+u_lup,l) )   
     
        ut(ij+u_lup,l) =  1/de(ij+u_lup) *                                           &  
                         ( wee(ij+u_lup,1,1)*ue(ij+u_left,l)+                        &
                           wee(ij+u_lup,2,1)*ue(ij+u_ldown,l)+                       &
                           wee(ij+u_lup,3,1)*ue(ij+u_rdown,l)+                       &
                           wee(ij+u_lup,4,1)*ue(ij+u_right,l)+                       &
                           wee(ij+u_lup,5,1)*ue(ij+u_rup,l)+                         & 
                           wee(ij+u_lup,1,2)*ue(ij+t_lup+u_right,l)+                 &
                           wee(ij+u_lup,2,2)*ue(ij+t_lup+u_rup,l)+                   &
                           wee(ij+u_lup,3,2)*ue(ij+t_lup+u_lup,l)+                   &
                           wee(ij+u_lup,4,2)*ue(ij+t_lup+u_left,l)+                  &
                           wee(ij+u_lup,5,2)*ue(ij+t_lup+u_ldown,l) )

    
        ut(ij+u_ldown,l) = 1/de(ij+u_ldown) *   &  
                         ( wee(ij+u_ldown,1,1)*ue(ij+u_rdown,l)+                      &
                           wee(ij+u_ldown,2,1)*ue(ij+u_right,l)+                      &
                           wee(ij+u_ldown,3,1)*ue(ij+u_rup,l)+                        &
                           wee(ij+u_ldown,4,1)*ue(ij+u_lup,l)+                        &
                           wee(ij+u_ldown,5,1)*ue(ij+u_left,l)+                       & 
                           wee(ij+u_ldown,1,2)*ue(ij+t_ldown+u_lup,l)+                &
                           wee(ij+u_ldown,2,2)*ue(ij+t_ldown+u_left,l)+               &
                           wee(ij+u_ldown,3,2)*ue(ij+t_ldown+u_ldown,l)+              &
                           wee(ij+u_ldown,4,2)*ue(ij+t_ldown+u_rdown,l)+              &
                           wee(ij+u_ldown,5,2)*ue(ij+t_ldown+u_right,l) )  
        
        ENDDO
      ENDDO
    ENDDO
                        
 END SUBROUTINE compute_tangential_compound
 
! SUBROUTINE compute_wind_lonlat_compound(u, ulon, ulat)
!  REAL(rstd) :: u(3*iim*jjm,3,llm)
!  REAL(rstd) :: ulon(3*iim*jjm,3,llm)
!  REAL(rstd) :: ulat(3*iim*jjm,3,llm)
!
!  INTEGER :: i,j,ij,l     
!    
!  
!    DO l=ll_begin,ll_end
!      DO j=jj_begin-1,jj_end+1
!        DO i=ii_begin-1,ii_end+1
!          ij=(j-1)*iim+i
!          ulon(ij+u_right,:,l)=sum(u(ij+u_right,:,l)*elon_e(ij+u_right,:))*elon_e(ij+u_right,:) 
!          ulon(ij+u_lup,:,l)=sum(u(ij+u_lup,:,l)*elon_e(ij+u_lup,:))*elon_e(ij+u_lup,:)
!          ulon(ij+u_ldown,:,l)=sum(u(ij+u_ldown,:,l)*elon_e(ij+u_ldown,:))*elon_e(ij+u_ldown,:)
!          
!          ulat(ij+u_right,:,l)=sum(u(ij+u_right,:,l)*elat_e(ij+u_right,:))*elat_e(ij+u_right,:) 
!          ulat(ij+u_lup,:,l)=sum(u(ij+u_lup,:,l)*elat_e(ij+u_lup,:))*elat_e(ij+u_lup,:) 
!          ulat(ij+u_ldown,:,l)=sum(u(ij+u_ldown,:,l)*elat_e(ij+u_ldown,:))*elat_e(ij+u_ldown,:) 
!          
!        ENDDO
!      ENDDO
!    ENDDO
! 
! END SUBROUTINE compute_wind_lonlat_compound
 
  SUBROUTINE compute_wind_from_lonlat_compound(ulon, ulat, u)
  REAL(rstd) :: u(3*iim*jjm,llm,3)
  REAL(rstd) :: ulon(3*iim*jjm,llm)
  REAL(rstd) :: ulat(3*iim*jjm,llm)

  INTEGER :: i,j,ij,l     
  
    DO l=ll_begin,ll_end
      DO j=jj_begin-1,jj_end+1
        DO i=ii_begin-1,ii_end+1
          ij=(j-1)*iim+i
          u(ij+u_right,l,:)=ulon(ij+u_right,l)*elon_e(ij+u_right,:)+ ulat(ij+u_right,l)*elat_e(ij+u_right,:)
          u(ij+u_lup,l,:)=ulon(ij+u_lup,l)*elon_e(ij+u_lup,:) + ulat(ij+u_lup,l)*elat_e(ij+u_lup,:)
          u(ij+u_ldown,l,:)=ulon(ij+u_ldown,l)*elon_e(ij+u_ldown,:) + ulat(ij+u_ldown,l)*elat_e(ij+u_ldown,:)
        ENDDO
      ENDDO
    ENDDO
 
  END SUBROUTINE compute_wind_from_lonlat_compound
 
  SUBROUTINE compute_wind_centered_from_lonlat_compound(ulon, ulat, u)
  REAL(rstd) :: u(iim*jjm,llm,3)
  REAL(rstd) :: ulon(iim*jjm,llm)
  REAL(rstd) :: ulat(iim*jjm,llm)
  INTEGER :: i,j,ij,l     
  DO l=ll_begin,ll_end
      DO j=jj_begin-1,jj_end+1
        DO i=ii_begin-1,ii_end+1
          ij=(j-1)*iim+i
          u(ij,l,:)=ulon(ij,l)*elon_i(ij,:)+ ulat(ij,l)*elat_i(ij,:)
        ENDDO
      ENDDO
    ENDDO 
  END SUBROUTINE compute_wind_centered_from_lonlat_compound
 
  SUBROUTINE compute_wind2D_from_lonlat_compound(ulon, ulat, u)
  REAL(rstd) :: u(3*iim*jjm,3)
  REAL(rstd) :: ulon(3*iim*jjm)
  REAL(rstd) :: ulat(3*iim*jjm)
  
  INTEGER :: i,j,ij
  
  DO j=jj_begin-1,jj_end+1
     DO i=ii_begin-1,ii_end+1
        ij=(j-1)*iim+i
        u(ij+u_right,:)=ulon(ij+u_right)*elon_e(ij+u_right,:)+ ulat(ij+u_right)*elat_e(ij+u_right,:)
        u(ij+u_lup,:)=ulon(ij+u_lup)*elon_e(ij+u_lup,:) + ulat(ij+u_lup)*elat_e(ij+u_lup,:)
        u(ij+u_ldown,:)=ulon(ij+u_ldown)*elon_e(ij+u_ldown,:) + ulat(ij+u_ldown)*elat_e(ij+u_ldown,:)
     ENDDO
  ENDDO
    
  END SUBROUTINE compute_wind2D_from_lonlat_compound
 
  SUBROUTINE compute_wind_perp_from_lonlat_compound(ulon, ulat, up)
  REAL(rstd) :: up(3*iim*jjm,llm)
  REAL(rstd) :: ulon(3*iim*jjm,llm)
  REAL(rstd) :: ulat(3*iim*jjm,llm)
  REAL(rstd) :: u(3*iim*jjm,llm,3)

  INTEGER :: i,j,ij,l     
  
   CALL compute_wind_from_lonlat_compound(ulon, ulat, u)

    DO l=ll_begin,ll_end
      DO j=jj_begin-1,jj_end+1
        DO i=ii_begin-1,ii_end+1
          ij=(j-1)*iim+i
          up(ij+u_right,l)=sum(u(ij+u_right,l,:)*ep_e(ij+u_right,:))
          up(ij+u_lup,l)=sum(u(ij+u_lup,l,:)*ep_e(ij+u_lup,:))
          up(ij+u_ldown,l)=sum(u(ij+u_ldown,l,:)*ep_e(ij+u_ldown,:))
        ENDDO
      ENDDO
    ENDDO
 
  END SUBROUTINE compute_wind_perp_from_lonlat_compound
   
  SUBROUTINE compute_wind2D_perp_from_lonlat_compound(ulon, ulat, up)
  REAL(rstd) :: up(3*iim*jjm)
  REAL(rstd) :: ulon(3*iim*jjm)
  REAL(rstd) :: ulat(3*iim*jjm)
  REAL(rstd) :: u(3*iim*jjm,3)

  INTEGER :: i,j,ij 
  
  CALL compute_wind2D_from_lonlat_compound(ulon, ulat, u)
  DO j=jj_begin-1,jj_end+1
     DO i=ii_begin-1,ii_end+1
        ij=(j-1)*iim+i
        up(ij+u_right)=sum(u(ij+u_right,:)*ep_e(ij+u_right,:))
        up(ij+u_lup)=sum(u(ij+u_lup,:)*ep_e(ij+u_lup,:))
        up(ij+u_ldown)=sum(u(ij+u_ldown,:)*ep_e(ij+u_ldown,:))
     ENDDO
  ENDDO
   
  END SUBROUTINE compute_wind2D_perp_from_lonlat_compound
   
 SUBROUTINE compute_wind_centered_lonlat_compound(uc, ulon, ulat)
  REAL(rstd) :: uc(iim*jjm,llm,3)
  REAL(rstd) :: ulon(iim*jjm,llm)
  REAL(rstd) :: ulat(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
          ulon(ij,l)=sum(uc(ij,l,:)*elon_i(ij,:))
          ulat(ij,l)=sum(uc(ij,l,:)*elat_i(ij,:)) 
        ENDDO
      ENDDO
    ENDDO
 
 END SUBROUTINE compute_wind_centered_lonlat_compound

 SUBROUTINE compute_wind_centered_from_wind_lonlat_centered(ulon, ulat,uc)
  REAL(rstd) :: ulon(iim*jjm,llm)
  REAL(rstd) :: ulat(iim*jjm,llm)
  REAL(rstd) :: uc(iim*jjm,llm,3)

  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
          uc(ij,l,:)=ulon(ij,l)*elon_i(ij,:)+ulat(ij,l)*elat_i(ij,:)
        ENDDO
      ENDDO
    ENDDO
 
 END SUBROUTINE compute_wind_centered_from_wind_lonlat_centered

 SUBROUTINE compute_wind_perp_from_wind_centered(uc,un)

  IMPLICIT NONE
  REAL(rstd),INTENT(IN)   :: uc(iim*jjm,llm,3)
  REAL(rstd),INTENT(OUT)  :: un(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
          un(ij+u_right,l) = sum(0.5*(uc(ij,l,:) + uc(ij+t_right,l,:))*ep_e(ij+u_right,:))
          un(ij+u_lup,l) = sum(0.5*(uc(ij,l,:) + uc(ij+t_lup,l,:))*ep_e(ij+u_lup,:))
          un(ij+u_ldown,l) = sum(0.5*(uc(ij,l,:) + uc(ij+t_ldown,l,:))*ep_e(ij+u_ldown,:))
         ENDDO
      ENDDO
    ENDDO
 
 END SUBROUTINE compute_wind_perp_from_wind_centered


 SUBROUTINE compute_un2ulonlat(un, ulon, ulat)
  REAL(rstd),INTENT(IN)  :: un(3*iim*jjm,llm)
  REAL(rstd),INTENT(OUT) :: ulon(iim*jjm,llm)
  REAL(rstd),INTENT(OUT) :: ulat(iim*jjm,llm)

  REAL(rstd)             :: uc(iim*jjm,llm,3)
    
  CALL compute_wind_centered(un,uc)  
  CALL compute_wind_centered_lonlat_compound(uc, ulon, ulat)
 
 END SUBROUTINE compute_un2ulonlat

 SUBROUTINE compute_ulonlat2un(ulon, ulat,un)
  REAL(rstd),INTENT(IN) :: ulon(iim*jjm,llm)
  REAL(rstd),INTENT(IN) :: ulat(iim*jjm,llm)
  REAL(rstd),INTENT(OUT)  :: un(3*iim*jjm,llm)

  REAL(rstd)             :: uc(iim*jjm,llm,3)
    
    CALL compute_wind_centered_from_wind_lonlat_centered(ulon, ulat, uc)  
    CALL compute_wind_perp_from_wind_centered(uc, un)
 
 END SUBROUTINE compute_ulonlat2un


END MODULE wind_mod