MODULE xios_mod

#ifdef CPP_USING_XIOS
   USE xios
#endif

  PUBLIC
  LOGICAL,SAVE :: using_xios

  INTEGER,SAVE :: ncell_i
!$OMP THREADPRIVATE(ncell_i)
  INTEGER,SAVE :: ncell_v
!$OMP THREADPRIVATE(ncell_v)

  PRIVATE ncell_i,ncell_v
  
CONTAINS

#ifdef CPP_USING_XIOS
 
 SUBROUTINE xios_init
   USE getin_mod
   IMPLICIT NONE

   using_xios=.TRUE.
   
 END SUBROUTINE xios_init  
 
 SUBROUTINE xios_init_write_field
 USE genmod
 USE mpipara
 USE xios
 USE grid_param
 USE domain_mod
 USE dimensions
 USE spherical_geom_mod
 USE geometry
 USE mpi_mod
 USE time_mod
 USE metric, ONLY : vup,vdown
 IMPLICIT NONE
  TYPE(xios_context) :: ctx_hdl
  TYPE(xios_duration)      :: dtime
  REAL(rstd) :: lev_value(llm)
  REAL(rstd) :: lev_valuep1(llm+1)
  INTEGER :: ncell, ncell_tot, ncell_glo(0:mpi_size-1), displ
  INTEGER :: ind, i,j,k,l
  REAL(rstd),ALLOCATABLE    :: lon(:), lat(:), bounds_lon(:,:), bounds_lat(:,:)
  TYPE(t_domain),POINTER :: d

!$OMP BARRIER
!$OMP MASTER
   CALL xios_context_initialize("icosagcm",comm_icosa)
   CALL xios_get_handle("icosagcm",ctx_hdl)
   CALL xios_set_current_context(ctx_hdl)
   lev_value(:) = (/ (l,l=1,llm) /)     
   lev_valuep1(:) = (/ (l,l=1,llm+1) /)     
   CALL xios_set_axis_attr("lev",n_glo=llm ,value=lev_value) ;
   CALL xios_set_axis_attr("levp1",n_glo=llm+1 ,value=lev_value) ;
   
   ncell=0
   DO ind=1,ndomain
     d=>domain(ind)
        
     DO j=d%jj_begin,d%jj_end
       DO i=d%ii_begin,d%ii_end
         IF (domain(ind)%own(i,j)) ncell=ncell+1
       ENDDO
     ENDDO
   ENDDO      
   ncell_i=ncell
   
   CALL MPI_ALLGATHER(ncell,1,MPI_INTEGER,ncell_glo,1,MPI_INTEGER,comm_icosa,ierr)

   displ=0
   DO i=1,mpi_rank
     displ=displ+ncell_glo(i-1)
   ENDDO

   ncell_tot=sum(ncell_glo(:))
   
   ALLOCATE(lon(ncell), lat(ncell), bounds_lon(0:5,ncell), bounds_lat(0:5,ncell)) 
   
   ncell=0
   DO ind=1,ndomain
     d=>domain(ind)
        
     DO j=d%jj_begin,d%jj_end
       DO i=d%ii_begin,d%ii_end
         IF (domain(ind)%own(i,j)) THEN 
           ncell=ncell+1
           CALL xyz2lonlat(d%xyz(:,i,j),lon(ncell),lat(ncell))
           lon(ncell)=lon(ncell)*180/Pi
           lat(ncell)=lat(ncell)*180/Pi
           DO k=0,5
             CALL xyz2lonlat(d%vertex(:,k,i,j),bounds_lon(k,ncell), bounds_lat(k,ncell))
             bounds_lat(k,ncell)=bounds_lat(k,ncell)*180/Pi
             bounds_lon(k,ncell)=bounds_lon(k,ncell)*180/Pi
           ENDDO
         ENDIF
       ENDDO
     ENDDO
   ENDDO         

  
   CALL xios_set_domaingroup_attr("i",ni_glo=ncell_tot, ibegin=displ, ni=ncell)
   CALL xios_set_domaingroup_attr("i", data_dim=1, type='unstructured' , nvertex=6)
   CALL xios_set_domaingroup_attr("i",lonvalue_1d=lon, latvalue_1d=lat, bounds_lon_1d=bounds_lon, bounds_lat_1d=bounds_lat)
   
   DEALLOCATE(lon, lat, bounds_lon, bounds_lat) 
    
   ncell=0
   DO ind=1,ndomain
     d=>domain(ind)
        
     DO j=d%jj_begin+1,d%jj_end
       DO i=d%ii_begin,d%ii_end-1
         ncell=ncell+1
       ENDDO
     ENDDO

     DO j=d%jj_begin,d%jj_end-1
       DO i=d%ii_begin+1,d%ii_end
          ncell=ncell+1
        ENDDO
     ENDDO

   ENDDO      
   
   ncell_v=ncell
   
   CALL MPI_ALLGATHER(ncell,1,MPI_INTEGER,ncell_glo,1,MPI_INTEGER,comm_icosa,ierr)

   displ=0
   DO i=1,mpi_rank
     displ=displ+ncell_glo(i-1)
   ENDDO

   ncell_tot=sum(ncell_glo(:))
   
   ALLOCATE(lon(ncell), lat(ncell), bounds_lon(0:2,ncell), bounds_lat(0:2,ncell)) 
   
   ncell=0
   DO ind=1,ndomain
     d=>domain(ind)
 
     DO j=d%jj_begin+1,d%jj_end
       DO i=d%ii_begin,d%ii_end-1
           ncell=ncell+1
           CALL xyz2lonlat(d%vertex(:,vdown,i,j),lon(ncell),lat(ncell))
           lon(ncell)=lon(ncell)*180/Pi
           lat(ncell)=lat(ncell)*180/Pi

           CALL xyz2lonlat(d%xyz(:,i,j),bounds_lon(0,ncell), bounds_lat(0,ncell))
           CALL xyz2lonlat(d%xyz(:,i,j-1),bounds_lon(1,ncell), bounds_lat(1,ncell))
           CALL xyz2lonlat(d%xyz(:,i+1,j-1),bounds_lon(2,ncell), bounds_lat(2,ncell))

           DO k=0,2
             bounds_lat(k,ncell)=bounds_lat(k,ncell)*180/Pi
             bounds_lon(k,ncell)=bounds_lon(k,ncell)*180/Pi
           ENDDO
         ENDDO
       ENDDO 
 
       DO j=d%jj_begin,d%jj_end-1
         DO i=d%ii_begin+1,d%ii_end
           ncell=ncell+1
           CALL xyz2lonlat(d%vertex(:,vup,i,j),lon(ncell),lat(ncell))
           lon(ncell)=lon(ncell)*180/Pi
           lat(ncell)=lat(ncell)*180/Pi
           CALL xyz2lonlat(d%xyz(:,i,j),bounds_lon(0,ncell), bounds_lat(0,ncell))
           CALL xyz2lonlat(d%xyz(:,i,j+1),bounds_lon(1,ncell), bounds_lat(1,ncell))
           CALL xyz2lonlat(d%xyz(:,i-1,j+1),bounds_lon(2,ncell), bounds_lat(2,ncell))

           DO k=0,2
             bounds_lat(k,ncell)=bounds_lat(k,ncell)*180/Pi
             bounds_lon(k,ncell)=bounds_lon(k,ncell)*180/Pi
           ENDDO
         ENDDO
       ENDDO 
        
   ENDDO         

  
   CALL xios_set_domain_attr("v",ni_glo=ncell_tot, ibegin=displ, ni=ncell)
   CALL xios_set_domain_attr("v", data_dim=1, type='unstructured' , nvertex=3)
   CALL xios_set_domain_attr("v",lonvalue_1d=lon, latvalue_1d=lat, bounds_lon_1d=bounds_lon, bounds_lat_1d=bounds_lat)


   dtime%second=dt
   CALL xios_set_timestep(dtime) 
   CALL xios_close_context_definition()
!$OMP END MASTER
!$OMP BARRIER
   
 END SUBROUTINE xios_init_write_field
  
  
 SUBROUTINE xios_write_field(name,field)
 USE field_mod
 IMPLICIT NONE 
   CHARACTER(LEN=*),INTENT(IN) :: name
   TYPE(t_field), POINTER :: field(:)
   CHARACTER(LEN=10) :: str_number
   INTEGER :: iq

!$OMP BARRIER
!$OMP MASTER
   
   IF (Field(1)%field_type==field_T) THEN
     IF (field(1)%ndim==2) THEN
        CALL xios_write_field_scalar(name,field,1)
      ELSE IF (field(1)%ndim==3) THEN
        CALL xios_write_field_scalar(name,field,size(field(1)%rval3d,2))
      ELSE IF (field(1)%ndim==4) THEN
        DO iq=1,size(field(1)%rval4d,3)
          WRITE(str_number,'(i10)') iq
          CALL xios_write_field_scalar(name//TRIM(ADJUSTL(str_number)),field,size(field(1)%rval4d,2),iq)
        ENDDO
      ELSE
        PRINT *, "xios_write_field : dimension > 4 are not supported for now"
      ENDIF
    ELSE IF (Field(1)%field_type==field_Z) THEN
     IF (field(1)%ndim==2) THEN
        CALL xios_write_field_vort(name,field,1)
      ELSE IF (field(1)%ndim==3) THEN
        CALL xios_write_field_vort(name,field,size(field(1)%rval3d,2))
      ELSE IF (field(1)%ndim==4) THEN
        DO iq=1,size(field(1)%rval4d,3)
          WRITE(str_number,'(i10)') iq
          CALL xios_write_field_vort(name//TRIM(ADJUSTL(str_number)),field,size(field(1)%rval4d,2),iq)
        ENDDO
      ELSE
        PRINT *, "xios_write_field : dimension > 4 are not supported for now"
      ENDIF
    ENDIF
!$OMP END MASTER
!$OMP BARRIER
      
 END SUBROUTINE xios_write_field

 SUBROUTINE xios_read_field(name,field)
 USE field_mod
 IMPLICIT NONE 
   CHARACTER(LEN=*),INTENT(IN) :: name
   TYPE(t_field), POINTER :: field(:)
   CHARACTER(LEN=10) :: str_number
   INTEGER :: iq

!$OMP BARRIER
!$OMP MASTER
   
   IF (Field(1)%field_type==field_T) THEN
     IF (field(1)%ndim==2) THEN
        CALL xios_read_field_scalar(name,field,1)
      ELSE IF (field(1)%ndim==3) THEN
        CALL xios_read_field_scalar(name,field,size(field(1)%rval3d,2))
      ELSE IF (field(1)%ndim==4) THEN
        DO iq=1,size(field(1)%rval4d,3)
          WRITE(str_number,'(i10)') iq
          CALL xios_read_field_scalar(name//TRIM(ADJUSTL(str_number)),field,size(field(1)%rval4d,2),iq)
        ENDDO
      ELSE
        PRINT *, "xios_write_field : dimension > 4 are not supported for now"
      ENDIF
    ELSE IF (Field(1)%field_type==field_Z) THEN
     IF (field(1)%ndim==2) THEN
        CALL xios_read_field_vort(name,field,1)
      ELSE IF (field(1)%ndim==3) THEN
        CALL xios_read_field_vort(name,field,size(field(1)%rval3d,2))
      ELSE IF (field(1)%ndim==4) THEN
        DO iq=1,size(field(1)%rval4d,3)
          WRITE(str_number,'(i10)') iq
          CALL xios_read_field_vort(name//TRIM(ADJUSTL(str_number)),field,size(field(1)%rval4d,2),iq)
        ENDDO
      ELSE
        PRINT *, "xios_write_field : dimension > 4 are not supported for now"
      ENDIF
    ENDIF
!$OMP END MASTER
!$OMP BARRIER
      
 END SUBROUTINE xios_read_field


 
 SUBROUTINE xios_write_field_scalar(name,field,nlev,iq)
 USE genmod
 USE mpipara
 USE xios
 USE grid_param
 USE domain_mod
 USE dimensions
 USE spherical_geom_mod
 USE geometry
 USE mpi_mod
 IMPLICIT NONE 
   CHARACTER(LEN=*),INTENT(IN) :: name
   TYPE(t_field), POINTER :: field(:)
   INTEGER,INTENT(IN) :: nlev
   INTEGER,INTENT(IN),OPTIONAL :: iq
   
   REAL(rstd) :: field_tmp(ncell_i,nlev)
   TYPE(t_domain),POINTER :: d
   INTEGER :: n,i,j,ij,ind
   
   IF (field(1)%ndim==2) THEN
     n=0
     DO ind=1,ndomain
       
       d=>domain(ind)
        
       DO j=d%jj_begin,d%jj_end
         DO i=d%ii_begin,d%ii_end
           IF (d%own(i,j)) THEN
             n=n+1
             ij=d%iim*(j-1)+i
             field_tmp(n,1)=field(ind)%rval2d(ij)
           ENDIF
         ENDDO
       ENDDO
     ENDDO
   ELSE IF (field(1)%ndim==3) THEN
     n=0
     DO ind=1,ndomain
       d=>domain(ind)
        
       DO j=d%jj_begin,d%jj_end
         DO i=d%ii_begin,d%ii_end
           IF (d%own(i,j)) THEN
             n=n+1
             ij=d%iim*(j-1)+i
             field_tmp(n,:)=field(ind)%rval3d(ij,:)
           ENDIF
         ENDDO
       ENDDO
     ENDDO
   ELSE IF (field(1)%ndim==4) THEN
     n=0
     DO ind=1,ndomain
       d=>domain(ind)
        
       DO j=d%jj_begin,d%jj_end
         DO i=d%ii_begin,d%ii_end
           IF (d%own(i,j)) THEN
             n=n+1
             ij=d%iim*(j-1)+i
             field_tmp(n,:)=field(ind)%rval4d(ij,:,iq)
           ENDIF
         ENDDO
       ENDDO
     ENDDO     
   ENDIF
   
   CALL xios_send_field(name,field_tmp)
 
 END SUBROUTINE xios_write_field_scalar  


 SUBROUTINE xios_read_field_scalar(name,field,nlev,iq)
 USE genmod
 USE mpipara
 USE xios
 USE grid_param
 USE domain_mod
 USE dimensions
 USE spherical_geom_mod
 USE geometry
 USE mpi_mod
 IMPLICIT NONE 
   CHARACTER(LEN=*),INTENT(IN) :: name
   TYPE(t_field), POINTER :: field(:)
   INTEGER,INTENT(IN) :: nlev
   INTEGER,INTENT(IN),OPTIONAL :: iq
   
   REAL(rstd) :: field_tmp(ncell_i,nlev)
   TYPE(t_domain),POINTER :: d
   INTEGER :: n,i,j,ij,ind

   CALL xios_recv_field(name,field_tmp)
   
   IF (field(1)%ndim==2) THEN
     n=0
     DO ind=1,ndomain
       
       d=>domain(ind)
        
       DO j=d%jj_begin,d%jj_end
         DO i=d%ii_begin,d%ii_end
           IF (d%own(i,j)) THEN
             n=n+1
             ij=d%iim*(j-1)+i
             field(ind)%rval2d(ij)=field_tmp(n,1)
           ENDIF
         ENDDO
       ENDDO
     ENDDO
   ELSE IF (field(1)%ndim==3) THEN
     n=0
     DO ind=1,ndomain
       d=>domain(ind)
        
       DO j=d%jj_begin,d%jj_end
         DO i=d%ii_begin,d%ii_end
           IF (d%own(i,j)) THEN
             n=n+1
             ij=d%iim*(j-1)+i
             field(ind)%rval3d(ij,:)=field_tmp(n,:)
           ENDIF
         ENDDO
       ENDDO
     ENDDO
   ELSE IF (field(1)%ndim==4) THEN
     n=0
     DO ind=1,ndomain
       d=>domain(ind)
        
       DO j=d%jj_begin,d%jj_end
         DO i=d%ii_begin,d%ii_end
           IF (d%own(i,j)) THEN
             n=n+1
             ij=d%iim*(j-1)+i
             field(ind)%rval4d(ij,:,iq)=field_tmp(n,:)
           ENDIF
         ENDDO
       ENDDO
     ENDDO     
   ENDIF
 
 END SUBROUTINE xios_read_field_scalar


      
 SUBROUTINE xios_write_field_vort(name,field,nlev,iq)
 USE genmod
 USE mpipara
 USE xios
 USE grid_param
 USE domain_mod
 USE dimensions
 USE spherical_geom_mod
 USE geometry
 USE mpi_mod
 IMPLICIT NONE 
   CHARACTER(LEN=*),INTENT(IN) :: name
   TYPE(t_field), POINTER :: field(:)
   INTEGER,INTENT(IN) :: nlev
   INTEGER,INTENT(IN),OPTIONAL :: iq
   
   REAL(rstd) :: field_tmp(ncell_v,nlev)
   TYPE(t_domain),POINTER :: d
   INTEGER :: n,i,j,ij,ind
   
   IF (field(1)%ndim==2) THEN
     n=0
     DO ind=1,ndomain
       d=>domain(ind)
       CALL swap_dimensions(ind)  
       
       DO j=d%jj_begin+1,d%jj_end
         DO i=d%ii_begin,d%ii_end-1
           n=n+1
           ij=iim*(j-1)+i
           Field_tmp(n,1)=field(ind)%rval2d(ij+z_down)
         ENDDO
       ENDDO

       DO j=d%jj_begin,d%jj_end-1
         DO i=d%ii_begin+1,d%ii_end
           n=n+1
           ij=iim*(j-1)+i
           Field_tmp(n,1)=field(ind)%rval2d(ij+z_up)
          ENDDO
       ENDDO
          
     ENDDO

   ELSE IF (field(1)%ndim==3) THEN
     n=0
     DO ind=1,ndomain
       d=>domain(ind)
       CALL swap_dimensions(ind)    
             
       DO j=d%jj_begin+1,d%jj_end
         DO i=d%ii_begin,d%ii_end-1
           n=n+1
           ij=iim*(j-1)+i
           Field_tmp(n,:)=field(ind)%rval3d(ij+z_down,:)
         ENDDO
       ENDDO

       DO j=d%jj_begin,d%jj_end-1
         DO i=d%ii_begin+1,d%ii_end
           n=n+1
           ij=iim*(j-1)+i
           Field_tmp(n,:)=field(ind)%rval3d(ij+z_up,:)
          ENDDO
       ENDDO
          
     ENDDO

   ELSE IF (field(1)%ndim==4) THEN
     n=0
     DO ind=1,ndomain
       d=>domain(ind)
       CALL swap_dimensions(ind) 
               
       DO j=d%jj_begin+1,d%jj_end
         DO i=d%ii_begin,d%ii_end-1
           n=n+1
           ij=iim*(j-1)+i
           Field_tmp(n,:)=field(ind)%rval4d(ij+z_down,:,iq)
         ENDDO
       ENDDO

       DO j=d%jj_begin,d%jj_end-1
         DO i=d%ii_begin+1,d%ii_end
           n=n+1
           ij=iim*(j-1)+i
           Field_tmp(n,:)=field(ind)%rval4d(ij+z_up,:,iq)
          ENDDO
       ENDDO
          
     ENDDO

   ENDIF
   
   CALL xios_send_field(name,field_tmp)
 
 END SUBROUTINE xios_write_field_vort 

 SUBROUTINE xios_read_field_vort(name,field,nlev,iq)
 USE genmod
 USE mpipara
 USE xios
 USE grid_param
 USE domain_mod
 USE dimensions
 USE spherical_geom_mod
 USE geometry
 USE mpi_mod
 IMPLICIT NONE 
   CHARACTER(LEN=*),INTENT(IN) :: name
   TYPE(t_field), POINTER :: field(:)
   INTEGER,INTENT(IN) :: nlev
   INTEGER,INTENT(IN),OPTIONAL :: iq
   
   REAL(rstd) :: field_tmp(ncell_v,nlev)
   TYPE(t_domain),POINTER :: d
   INTEGER :: n,i,j,ij,ind

   CALL xios_recv_field(name,field_tmp)

   
   IF (field(1)%ndim==2) THEN
     n=0
     DO ind=1,ndomain
       d=>domain(ind)
       CALL swap_dimensions(ind)  
       
       DO j=d%jj_begin+1,d%jj_end
         DO i=d%ii_begin,d%ii_end-1
           n=n+1
           ij=iim*(j-1)+i
           field(ind)%rval2d(ij+z_down)=Field_tmp(n,1)
         ENDDO
       ENDDO

       DO j=d%jj_begin,d%jj_end-1
         DO i=d%ii_begin+1,d%ii_end
           n=n+1
           ij=iim*(j-1)+i
           Field_tmp(n,1)=field(ind)%rval2d(ij+z_up)
           field(ind)%rval2d(ij+z_up)=Field_tmp(n,1)
          ENDDO
       ENDDO
          
     ENDDO

   ELSE IF (field(1)%ndim==3) THEN
     n=0
     DO ind=1,ndomain
       d=>domain(ind)
       CALL swap_dimensions(ind)    
             
       DO j=d%jj_begin+1,d%jj_end
         DO i=d%ii_begin,d%ii_end-1
           n=n+1
           ij=iim*(j-1)+i
           field(ind)%rval3d(ij+z_down,:)=Field_tmp(n,:)
         ENDDO
       ENDDO

       DO j=d%jj_begin,d%jj_end-1
         DO i=d%ii_begin+1,d%ii_end
           n=n+1
           ij=iim*(j-1)+i
           field(ind)%rval3d(ij+z_up,:)=Field_tmp(n,:)
          ENDDO
       ENDDO
          
     ENDDO

   ELSE IF (field(1)%ndim==4) THEN
     n=0
     DO ind=1,ndomain
       d=>domain(ind)
       CALL swap_dimensions(ind) 
               
       DO j=d%jj_begin+1,d%jj_end
         DO i=d%ii_begin,d%ii_end-1
           n=n+1
           ij=iim*(j-1)+i
           field(ind)%rval4d(ij+z_down,:,iq)=Field_tmp(n,:)
         ENDDO
       ENDDO

       DO j=d%jj_begin,d%jj_end-1
         DO i=d%ii_begin+1,d%ii_end
           n=n+1
           ij=iim*(j-1)+i
           field(ind)%rval4d(ij+z_up,:,iq)=Field_tmp(n,:)
          ENDDO
       ENDDO
          
     ENDDO

   ENDIF
 
 END SUBROUTINE xios_read_field_vort 




  
 SUBROUTINE xios_write_field_finalize
 IMPLICIT NONE

!$OMP BARRIER
!$OMP MASTER 
   CALL xios_context_finalize
!$OMP END MASTER
!$OMP BARRIER

 END SUBROUTINE xios_write_field_finalize

 SUBROUTINE xios_set_context
 IMPLICIT NONE   
  TYPE(xios_context) :: ctx_hdl

!$OMP MASTER 
   CALL xios_get_handle("icosagcm",ctx_hdl)
   CALL xios_set_current_context(ctx_hdl)
!$OMP END MASTER

  END SUBROUTINE xios_set_context
#else
  
  SUBROUTINE xios_init
   IMPLICIT NONE
         
    using_xios=.FALSE.
    
  END SUBROUTINE xios_init
  
  SUBROUTINE xios_write_field(name,field)
  USE field_mod
  IMPLICIT NONE 
   CHARACTER(LEN=*),INTENT(IN) :: name
   TYPE(t_field), POINTER :: field(:)
  END SUBROUTINE xios_write_field
 
  SUBROUTINE xios_update_calendar(step)
  IMPLICIT NONE
   INTEGER, INTENT(IN):: step  
  END SUBROUTINE xios_update_calendar

  SUBROUTINE xios_write_field_finalize
  END SUBROUTINE xios_write_field_finalize
  
  SUBROUTINE xios_init_write_field
  END SUBROUTINE xios_init_write_field  
  
  SUBROUTINE xios_set_context
  END SUBROUTINE xios_set_context
  

#endif  
  
END MODULE xios_mod