MODULE xios_mod
  
#ifdef CPP_USING_XIOS
  USE xios
#endif

  USE prec, ONLY : rstd
  USE field_mod, ONLY : t_field, field_T, field_U, field_Z
  USE domain_mod, ONLY : t_domain, t_cellset, domain, ndomain, mesh_loc
  IMPLICIT NONE   
  PRIVATE
  SAVE

  LOGICAL :: using_xios

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

  PUBLIC :: using_xios, xios_init, &
       xios_init_write_field,  xios_write_field_finalize, &
       xios_write_field, xios_read_field

#ifdef CPP_USING_XIOS

  PUBLIC :: xios_timestep,    & 
       xios_set_file_attr, xios_set_fieldgroup_attr, &
       xios_set_filegroup_attr, xios_get_axis_attr, &
       xios_send_field, xios_read_var, &
       xios_update_calendar, xios_set_context
  
CONTAINS
 
 SUBROUTINE xios_init
   USE mpipara, ONLY : comm_icosa
   TYPE(xios_context) :: ctx_hdl

   using_xios=.TRUE.
   CALL xios_context_initialize("icosagcm",comm_icosa)
   CALL xios_get_handle("icosagcm",ctx_hdl)
   CALL xios_set_current_context(ctx_hdl)   

 END SUBROUTINE xios_init

 SUBROUTINE xios_init_write_field
   USE time_mod, ONLY : dt, itau_out
   USE grid_param, ONLY : llm, nqtot
   USE mpi_mod, ONLY : MPI_INTEGER

   TYPE(xios_context) :: ctx_hdl
   TYPE(xios_duration)      :: dtime
   REAL(rstd) :: lev_value(llm)
   REAL(rstd) :: lev_valuep1(llm+1)
   REAL(rstd) :: nq_value(nqtot)
   INTEGER :: l, ncell, ncell_tot, displ
   REAL(rstd),ALLOCATABLE    :: lon(:), lat(:), bounds_lon(:,:), bounds_lat(:,:)
   INTEGER, ALLOCATABLE      :: ind_glo(:)
   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) /)     
   nq_value(:) = (/ (l,l=1,nqtot) /)     
   CALL xios_set_axis_attr("lev",n_glo=llm ,value=lev_value) ;
   CALL xios_set_axis_attr("levp1",n_glo=llm+1 ,value=lev_valuep1) ;
   CALL xios_set_axis_attr("nq",n_glo=nqtot, value=nq_value) ;
   
   !------------------ primal cells ------------------
   CALL collect_bounds(6, mesh_loc%primal_own)
   ncell_i=ncell   
   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, i_index=ind_glo)
   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,ind_glo) 
   
   !--------------------- dual cells ------------------
   
   CALL collect_bounds(3, mesh_loc%dual_own)
   ncell_v=ncell
   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)
   DEALLOCATE(lon, lat, bounds_lon, bounds_lat,ind_glo) 
   
   !---------------------- edges -----------------------
   
   CALL collect_bounds(2, mesh_loc%edge_own)
   ncell_e=ncell
   CALL xios_set_domain_attr("u",ni_glo=ncell_tot, ibegin=displ, ni=ncell)
   CALL xios_set_domain_attr("u", data_dim=1, type='unstructured' , nvertex=2, i_index=ind_glo)
   CALL xios_set_domain_attr("u",lonvalue_1d=lon, latvalue_1d=lat, bounds_lon_1d=bounds_lon, bounds_lat_1d=bounds_lat)
   DEALLOCATE(lon, lat, bounds_lon, bounds_lat,ind_glo) 
   
   dtime%second=dt
   CALL xios_set_timestep(dtime)
   
   CALL xios_set_fieldgroup_attr("standard_output", freq_op=itau_out*xios_timestep, freq_offset=(itau_out-1)*xios_timestep)
   
   CALL xios_close_context_definition()
   !$OMP END MASTER
   !$OMP BARRIER
   
 CONTAINS
   
   SUBROUTINE collect_bounds(nvert, cells)
     USE mpipara, ONLY : comm_icosa, mpi_size, mpi_rank
     INTEGER, INTENT(IN) :: nvert
     TYPE(t_cellset)     :: cells(:)
     INTEGER :: i, ind, n_beg, n_end, ierr, ncell_glo(0:mpi_size-1)
     ncell = SUM(cells%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), ind_glo(ncell))
     ALLOCATE(bounds_lon(nvert,ncell), bounds_lat(nvert,ncell))
     
     n_beg=0
     DO ind=1,ndomain
        n_end = n_beg + cells(ind)%ncell
        ind_glo(n_beg+1:n_end) = cells(ind)%ind_glo(:)
        lon(n_beg+1:n_end) = cells(ind)%lon(:)
        lat(n_beg+1:n_end) = cells(ind)%lat(:)
        bounds_lon(:,n_beg+1:n_end) = cells(ind)%bnds_lon(:,:)
        bounds_lat(:,n_beg+1:n_end) = cells(ind)%bnds_lat(:,:)
        n_beg = n_end
     END DO
   END SUBROUTINE collect_bounds
   
 END SUBROUTINE xios_init_write_field
 
 SUBROUTINE xios_write_field(name,field)
   CHARACTER(LEN=*),INTENT(IN) :: name
   TYPE(t_field), POINTER :: field(:)
   TYPE(t_cellset), POINTER :: cells(:)
   INTEGER :: ncells

!$OMP BARRIER
!$OMP MASTER

   SELECT CASE(field(1)%field_type)
   CASE(field_T)
      cells => mesh_loc%primal_own
      ncells = ncell_i
   CASE(field_U)
      cells => mesh_loc%edge_own
      ncells = ncell_e
   CASE(field_Z)
      cells => mesh_loc%dual_own
      ncells = ncell_v
   END SELECT

   IF (field(1)%ndim>4) THEN
      PRINT *, "xios_write_field : dimension > 4 are not supported for now"
   ELSE
      CALL xios_write_field_hex(name, field, cells, & 
           ncells, field(1)%dim3, field(1)%dim4)
   END IF
   
!$OMP END MASTER
!$OMP BARRIER
      
 END SUBROUTINE xios_write_field

 SUBROUTINE xios_read_field(name,field)
   CHARACTER(LEN=*),INTENT(IN) :: name
   TYPE(t_field), POINTER :: field(:)
   TYPE(t_cellset), POINTER :: cells(:)
   INTEGER :: ncells

!$OMP BARRIER
!$OMP MASTER

   SELECT CASE(field(1)%field_type)
   CASE(field_T)
      cells => mesh_loc%primal_own
      ncells = ncell_i
   CASE(field_U)
      cells => mesh_loc%edge_own
      ncells = ncell_e
   CASE(field_Z)
      cells => mesh_loc%dual_own
      ncells = ncell_v
   END SELECT

   IF (field(1)%ndim>4) THEN
      PRINT *, "xios_read_field : dimension > 4 are not supported for now"
   ELSE
      CALL xios_read_field_hex(name, field, cells, & 
           ncells, field(1)%dim3, field(1)%dim4)
   END IF
   
!$OMP END MASTER
!$OMP BARRIER
      
 END SUBROUTINE xios_read_field

 SUBROUTINE xios_write_field_hex(name, field, cells, ncell_tot, nlev, nq)
   CHARACTER(LEN=*),INTENT(IN) :: name
   TYPE(t_field) :: field(:)
   TYPE(t_cellset), TARGET :: cells(:)
   INTEGER,INTENT(IN) :: ncell_tot, nlev, nq

   REAL(rstd) :: field_tmp(ncell_tot,nlev,nq)
   TYPE(t_cellset), POINTER :: cellset
   INTEGER :: ind, n_beg, n_end, n, ij, sgn
   LOGICAL :: signed

   IF(ALLOCATED(cells(1)%sgn)) THEN
      signed=.TRUE.
   ELSE
      signed=.FALSE.
      sgn=1
   END IF

   n_beg=0
   DO ind=1,ndomain
      cellset => cells(ind)
      n_end = n_beg + cellset%ncell
      DO n=1, cellset%ncell
         ij = cellset%ij(n)
         IF(signed) sgn = cellset%sgn(n)
         SELECT CASE(field(1)%ndim)
         CASE(2)
            field_tmp(n_beg+n,1,1) = sgn*field(ind)%rval2d(ij)
         CASE(3)
            field_tmp(n_beg+n,:,1) = sgn*field(ind)%rval3d(ij,:)
         CASE(4)
            field_tmp(n_beg+n,:,:) = sgn*field(ind)%rval4d(ij,:,:)
         END SELECT
      END DO
   END DO
   CALL xios_send_field(name,field_tmp)  
 END SUBROUTINE xios_write_field_hex

SUBROUTINE xios_read_field_hex(name, field, cells, ncell_tot, nlev, nq)
   CHARACTER(LEN=*),INTENT(IN) :: name
   TYPE(t_field) :: field(:)
   TYPE(t_cellset), TARGET :: cells(:)
   INTEGER,INTENT(IN) :: ncell_tot, nlev, nq

   REAL(rstd) :: field_tmp(ncell_tot,nlev,nq)
   TYPE(t_cellset), POINTER :: cellset
   INTEGER :: ind, n_beg, n_end, n, ij, sgn
   LOGICAL :: signed

   CALL xios_recv_field(name,field_tmp)

   IF(ALLOCATED(cells(1)%sgn)) THEN
      signed=.TRUE.
   ELSE
      signed=.FALSE.
      sgn=1
   END IF

   n_beg=0
   DO ind=1,ndomain
      cellset => cells(ind)
      n_end = n_beg + cellset%ncell
      DO n=1, cellset%ncell
         ij = cellset%ij(n)
         IF(signed) sgn = cellset%sgn(n)
         SELECT CASE(field(1)%ndim)
         CASE(2)
            field(ind)%rval2d(ij) = sgn*field_tmp(n_beg+n,1,1) 
         CASE(3)
            field(ind)%rval3d(ij,:) = sgn*field_tmp(n_beg+n,:,1)
         CASE(4)
            field(ind)%rval4d(ij,:,:) = sgn*field_tmp(n_beg+n,:,:)
         END SELECT
      END DO
   END DO
 END SUBROUTINE xios_read_field_hex

 SUBROUTINE xios_read_var(name,field)
   USE prec
   USE transfert_mod
   CHARACTER(LEN=*),INTENT(IN) :: name
   REAL(rstd), INTENT(OUT) :: field
   !$OMP MASTER
   CALL xios_recv_field(name,field)
   !$OMP END MASTER
   CALL bcast_omp(field)
 END SUBROUTINE
  
 SUBROUTINE xios_write_field_finalize
!$OMP BARRIER
!$OMP MASTER 
   CALL xios_context_finalize
!$OMP END MASTER
!$OMP BARRIER
 END SUBROUTINE xios_write_field_finalize

 SUBROUTINE xios_set_context
  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
  

INTERFACE xios_send_field
  MODULE PROCEDURE xios_send_field_scalar, xios_send_field_1d
END INTERFACE  xios_send_field

INTEGER,PARAMETER :: xios_timestep=1

CONTAINS  
  
  
  SUBROUTINE xios_init
    using_xios=.FALSE.    
  END SUBROUTINE xios_init
  
  SUBROUTINE xios_send_field_scalar(name,field)
    CHARACTER(LEN=*),INTENT(IN) :: name
    REAL,INTENT(IN) :: field
  END SUBROUTINE xios_send_field_scalar  

  SUBROUTINE xios_send_field_1d(name,field)
    CHARACTER(LEN=*),INTENT(IN) :: name
    REAL,INTENT(IN) :: field(:)
  END SUBROUTINE xios_send_field_1d  
 
  SUBROUTINE xios_write_field(name,field)
  USE field_mod
   CHARACTER(LEN=*),INTENT(IN) :: name
   TYPE(t_field), POINTER :: field(:)
  END SUBROUTINE xios_write_field

  SUBROUTINE xios_read_field(name,field)
  USE field_mod
   CHARACTER(LEN=*),INTENT(IN) :: name
   TYPE(t_field), POINTER :: field(:)
  END SUBROUTINE xios_read_field

 SUBROUTINE xios_read_var(name,field)
   USE prec
   CHARACTER(LEN=*),INTENT(IN) :: name
   REAL(rstd), INTENT(OUT) :: field
 END SUBROUTINE
 
  SUBROUTINE xios_update_calendar(step)
   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
  
  SUBROUTINE xios_set_fieldgroup_attr(name,enabled,freq_op)
    CHARACTER(LEN=*) :: name
    LOGICAL,OPTIONAL          :: enabled
    INTEGER,OPTIONAL          :: freq_op
  END SUBROUTINE xios_set_fieldgroup_attr

  SUBROUTINE xios_set_filegroup_attr(name,enabled)
    CHARACTER(LEN=*) :: name
    LOGICAL,OPTIONAL          :: enabled
  END SUBROUTINE xios_set_filegroup_attr

  SUBROUTINE xios_set_file_attr(id,name,mode,enabled, output_freq)
    CHARACTER(LEN=*) :: id
    CHARACTER(LEN=*),OPTIONAL :: name, mode
    LOGICAL,OPTIONAL          :: enabled
    INTEGER,OPTIONAL          :: output_freq
  END SUBROUTINE xios_set_file_attr

  SUBROUTINE xios_get_axis_attr(name,n_glo,value)
    CHARACTER(LEN=*) :: name
    INTEGER,OPTIONAL          :: n_glo
    REAL,OPTIONAL             :: value(:)
  END SUBROUTINE xios_get_axis_attr

  SUBROUTINE xios_set_axis_attr(id,n_glo,value)
    CHARACTER(LEN=*) :: id
    INTEGER,OPTIONAL          :: n_glo
    REAL,OPTIONAL             :: value(:)
  END SUBROUTINE xios_set_axis_attr

#endif  

END MODULE xios_mod