module write_field_mod
   USE genmod
   USE write_field_vars_mod
   IMPLICIT NONE
   PRIVATE

   PUBLIC init_writeField, writefield, close_files

CONTAINS

   SUBROUTINE init_writeField
      USE ioipsl
      use netcdf_mod
      CHARACTER(LEN=255) :: netcdf_prec

      !$OMP CRITICAL
      netcdf_prec = 'float'
      CALL getin("netcdf_prec", netcdf_prec)
      SELECT CASE (TRIM(netcdf_prec))
      CASE ('float')
         ncprec = NF90_FLOAT
      CASE ('double')
         ncprec = NF90_DOUBLE
      CASE default
         PRINT *, 'Bad selector for variable netcdf_prec : <', TRIM(netcdf_prec), "> options are <float>, <double>"
         STOP
      END SELECT
      !$OMP END CRITICAL
   END SUBROUTINE init_writeField

   SUBROUTINE Writefield(name_in, field, nind, once)
      USE domain_mod
      USE field_mod
      USE transfert_mpi_mod
      USE transfert_mod, ONLY: gather_field
      USE dimensions
      USE mpipara
      USE netcdf_mod
      USE grid_param
      CHARACTER(LEN=*), INTENT(IN) :: name_in
      TYPE(t_field), POINTER :: field(:)
      INTEGER, OPTIONAL, INTENT(IN) :: nind
      LOGICAL, OPTIONAL, INTENT(IN) :: once
      LOGICAL                     :: once_

      TYPE(t_field), POINTER :: field_glo(:)

      ! NB : WriteField is not implemented for unstructured mesh
      IF (no_io .OR. grid_type == grid_unst) RETURN

!$OMP BARRIER
!$OMP MASTER
      IF (PRESENT(once)) THEN
         once_ = once
      ELSE
         once_ = .FALSE.
      END IF

      IF (field(1)%ndim == 2) THEN
         CALL allocate_field_glo(field_glo, field(1)%field_type, field(1)%data_type)
      ELSE IF (field(1)%ndim == 3) THEN
         CALL allocate_field_glo(field_glo, field(1)%field_type, field(1)%data_type, field(1)%dim3)
      ELSE IF (field(1)%ndim == 4) THEN
         CALL allocate_field_glo(field_glo, field(1)%field_type, field(1)%data_type, field(1)%dim3, field(1)%dim4)
      END IF

      CALL gather_field(field, field_glo)

      IF (mpi_rank == 0) THEN
         IF (PRESENT(nind)) THEN
            CALL writefield_gen(name_in, field_glo, domain_glo, nind, once=once_)
         ELSE
            CALL writefield_gen(name_in, field_glo, domain_glo, 1, ndomain_glo, once=once_)
         END IF
      END IF

      CALL deallocate_field_glo(field_glo)
!$OMP END MASTER
!$OMP BARRIER

   END SUBROUTINE writefield

   SUBROUTINE Writefield_gen(name_in, field, domain_type, ind_b_in, ind_e_in, once)
      USE netcdf_mod
      USE domain_mod
      USE field_mod
      CHARACTER(LEN=*), INTENT(IN) :: name_in
      TYPE(t_field), POINTER :: field(:)
      TYPE(t_domain), INTENT(IN), TARGET :: domain_type(:)
      INTEGER, OPTIONAL, INTENT(IN) :: ind_b_in
      INTEGER, OPTIONAL, INTENT(IN) :: ind_e_in
      LOGICAL, INTENT(IN) :: once
      ! local variables
      TYPE(t_cellset), POINTER :: cells(:)
      TYPE(t_domain), POINTER :: d
      LOGICAL :: single
      INTEGER :: ind_b, ind_e, ind
      CHARACTER(len=255) :: name
      INTEGER :: index, ncell, nvert, n
      ! for embedded routines
      REAL(r8), ALLOCATABLE :: field_val2d(:)
      REAL(r8), ALLOCATABLE :: field_val3d(:, :)
      INTEGER :: status, n_begin, ij, q, dim3

      name = TRIM(ADJUSTL(name_in))

      IF (PRESENT(ind_b_in) .AND. .NOT. PRESENT(ind_e_in)) THEN
         name = TRIM(name)//"_"//TRIM(int2str(ind_b))
         ind_b = ind_b_in
         ind_e = ind_b_in
         single = .TRUE.
      ELSE IF (.NOT. PRESENT(ind_b_in) .AND. PRESENT(ind_e_in)) THEN
         name = TRIM(name)//"_"//TRIM(int2str(ind_e))
         ind_b = ind_e_in
         ind_e = ind_e_in
         single = .TRUE.
      ELSE IF (PRESENT(ind_b_in) .AND. PRESENT(ind_e_in)) THEN
         ind_b = ind_b_in
         ind_e = ind_e_in
         single = .FALSE.
      ELSE
         ind_b = 1
         ind_e = ndomain
         single = .FALSE.
      END IF

      index = GetFieldIndex(name)
      if (index == -1) then
         call create_header_gen(name_in, field, domain_type, ind_b_in, ind_e_in, once)
         index = GetFieldIndex(name)
      else
         FieldIndex(index) = FieldIndex(index) + 1.
      end if

      SELECT CASE (Field(ind_b)%field_type)
      CASE (field_T)
         IF (single) THEN ! include halos
            cells => mesh_glo%primal_all
         ELSE
            cells => mesh_glo%primal_own
         END IF
      CASE (field_Z)
         IF (single) THEN ! include halos
            cells => mesh_glo%dual_all
         ELSE
            cells => mesh_glo%dual_own
         END IF
      END SELECT

      ncell = 0
      DO ind = ind_b, ind_e
         nvert = SIZE(cells(ind)%bnds_lon, 1)
         ncell = ncell + cells(ind)%ncell
      END DO

      SELECT CASE (field(1)%ndim)
      CASE (2)
         CALL write_2d
      CASE (3)
         CALL write_3d
      CASE (4)
         CALL write_4d
      END SELECT

      status = NF90_SYNC(FieldId(index))

   CONTAINS

      SUBROUTINE write_2d
         ALLOCATE (Field_val2d(ncell))
         n_begin = 0
         DO ind = ind_b, ind_e
            DO n = 1, cells(ind)%ncell
               ij = cells(ind)%ij(n)
               field_val2d(n_begin + n) = field(ind)%rval2d(ij)
            END DO
            n_begin = n_begin + cells(ind)%ncell
         END DO
         IF (once) THEN
            status = NF90_PUT_VAR(FieldId(index), FieldVarId(index)%nc_id(1), &
                                  Field_val2d, start=(/1/), count=(/ncell/))
         ELSE
            status = NF90_PUT_VAR(FieldId(Index), FieldVarId(index)%nc_id(1), &
                                  Field_val2d, start=(/1, FieldIndex(Index)/), count=(/ncell, 1/))
         END IF
         DEALLOCATE (field_val2d)
      END SUBROUTINE write_2d

      SUBROUTINE write_3d
         dim3 = SIZE(field(1)%rval3d, 2)
         ALLOCATE (field_val3d(ncell, dim3))
         n_begin = 0
         DO ind = ind_b, ind_e
            DO n = 1, cells(ind)%ncell
               ij = cells(ind)%ij(n)
               field_val3d(n_begin + n, :) = field(ind)%rval3d(ij, :)
            END DO
         END DO
         IF (once) THEN
            status = NF90_PUT_VAR(FieldId(Index), FieldVarId(index)%nc_id(1), &
                                  field_val3d, start=(/1, 1/), &
                                  count=(/ncell, dim3/))
         ELSE
            status = NF90_PUT_VAR(FieldId(Index), FieldVarId(index)%nc_id(1), &
                                  field_val3d, start=(/1, 1, FieldIndex(Index)/), &
                                  count=(/ncell, size(field(1)%rval3d, 2), 1/))
         END IF
         DEALLOCATE (field_val3d)
      END SUBROUTINE write_3d

      SUBROUTINE write_4d
         dim3 = SIZE(field(1)%rval4d, 2)
         ALLOCATE (field_val3d(ncell, dim3))
         DO q = 1, FieldVarId(index)%size
            n_begin = 0
            DO ind = ind_b, ind_e
               DO n = 1, cells(ind)%ncell
                  ij = cells(ind)%ij(n)
                  field_val3d(n_begin + n, :) = field(ind)%rval4d(ij, :, q)
               END DO
            END DO
            IF (once) THEN
               status = NF90_PUT_VAR(FieldId(Index), FieldVarId(index)%nc_id(q), &
                                     field_val3d, start=(/1, 1/), &
                                     count=(/ncell, dim3, 1/))
            ELSE
               status = NF90_PUT_VAR(FieldId(Index), FieldVarId(index)%nc_id(q), &
                                     field_val3d, start=(/1, 1, FieldIndex(Index)/), &
                                     count=(/ncell, dim3, 1/))
            END IF
         END DO
         DEALLOCATE (field_val3d)
      END SUBROUTINE write_4d

   END SUBROUTINE Writefield_gen

   SUBROUTINE Create_header_gen(name_in, field, domain_type, ind_b_in, ind_e_in, once)
      USE netcdf_mod
      USE field_mod
      USE domain_mod
      USE metric
      USE spherical_geom_mod
      CHARACTER(LEN=*), INTENT(IN) :: name_in
      TYPE(t_field), POINTER :: field(:)
      TYPE(t_domain), INTENT(IN), TARGET :: domain_type(:)
      INTEGER, OPTIONAL, INTENT(IN) :: ind_b_in
      INTEGER, OPTIONAL, INTENT(IN) :: ind_e_in
      LOGICAL, INTENT(IN) :: once

      TYPE(t_cellset), POINTER :: cells(:)
      INTEGER :: ncell
      INTEGER :: nvert
      REAL(rstd), ALLOCATABLE :: lon(:), lat(:), bounds_lon(:, :), bounds_lat(:, :)
      TYPE(t_domain), POINTER :: d
      INTEGER :: nvertId, ncid, lonId, latId, bounds_lonId, bounds_latId, timeId, ncellId
      INTEGER :: dim3id, dim4id
      INTEGER :: status
      INTEGER :: ind, i, j, k, n, q
      INTEGER :: iie, jje, iin, jjn
      INTEGER :: ind_b, ind_e, n_begin, n_end
      INTEGER :: halo_size
      LOGICAL :: single
      INTEGER :: nij
      CHARACTER(LEN=255) :: name
      INTEGER :: l, level_size, levId, dimlevId

      name = TRIM(ADJUSTL(name_in))

      IF (PRESENT(ind_b_in) .AND. .NOT. PRESENT(ind_e_in)) THEN
         name = TRIM(name)//"_"//TRIM(int2str(ind_b))
         ind_b = ind_b_in
         ind_e = ind_b_in
         halo_size = 1
         single = .TRUE.
      ELSE IF (.NOT. PRESENT(ind_b_in) .AND. PRESENT(ind_e_in)) THEN
         name = TRIM(name)//"_"//TRIM(int2str(ind_e))
         ind_b = ind_e_in
         ind_e = ind_e_in
         halo_size = 1
         single = .TRUE.
      ELSE IF (PRESENT(ind_b_in) .AND. PRESENT(ind_e_in)) THEN
         ind_b = ind_b_in
         ind_e = ind_e_in
         halo_size = 0
         single = .FALSE.
      ELSE
         halo_size = 0
         ind_b = 1
         ind_e = ndomain
         single = .FALSE.
      END IF

      NbField = NbField + 1
      FieldName(NbField) = TRIM(ADJUSTL(name))
      FieldIndex(NbField) = 1

!     PRINT *, 'Creating header for writefield : ', TRIM(FieldName(NbField)), ndomain, ind_b, ind_e ! FIXME
!     PRINT *, mesh_glo%ndomain, SIZE(mesh_glo%primal_own)
      SELECT CASE (Field(ind_b)%field_type)
      CASE (field_T)
!        PRINT *, '    field_type == field_T' ! FIXME
         IF (single) THEN ! include halos
            cells => mesh_glo%primal_all
         ELSE
            cells => mesh_glo%primal_own
         END IF
      CASE (field_Z)
!        PRINT *, '    field_type == field_Z' ! FIXME
         IF (single) THEN ! include halos
            cells => mesh_glo%dual_all
         ELSE
            cells => mesh_glo%dual_own
         END IF
      END SELECT

!     PRINT *, 'writefield : ind_b, ind_e :', ind_b, ind_e

      ncell = 0
      DO ind = ind_b, ind_e
         nvert = SIZE(cells(ind)%bnds_lon, 1)
         ncell = ncell + cells(ind)%ncell
      END DO
!     PRINT *, 'writefield : found ',ncell, ' cells.' ! FIXME

      status = NF90_CREATE(TRIM(ADJUSTL(name))//'.nc', NF90_CLOBBER, ncid)
      FieldId(NbField) = ncid
      status = NF90_DEF_DIM(ncid, 'time_counter', NF90_UNLIMITED, timeId)

      level_size = 0
      SELECT CASE (Field(ind_b)%ndim)
      CASE (2)
         FieldVarId(NbField)%size = 1
         ALLOCATE (FieldVarId(NbField)%nc_id(1))
      CASE (3)
         FieldVarId(NbField)%size = 1
         ALLOCATE (FieldVarId(NbField)%nc_id(1))
         status = NF90_DEF_DIM(ncid, 'lev', size(field(ind_b)%rval3d, 2), dim3id)
         level_size = size(field(ind_b)%rval3d, 2)
      CASE (4)
         FieldVarId(NbField)%size = size(field(ind_b)%rval4d, 3)
         ALLOCATE (FieldVarId(NbField)%nc_id(FieldVarId(NbField)%size))
         status = NF90_DEF_DIM(ncid, 'lev', size(field(ind_b)%rval4d, 2), dim3id)
         level_size = size(field(ind_b)%rval4d, 2)
         !          status = NF90_DEF_DIM(ncid,'Q',size(field(ind_b)%rval4d,3),dim4id)
      END SELECT

      PRINT *, "create_header_gen : LEVEL_SIZE=", level_size
      IF (level_size > 0) THEN
         status = NF90_DEF_VAR(ncid, 'lev', NF90_DOUBLE, (/dim3id/), levId)
         status = NF90_PUT_ATT(ncid, levId, "axis", "Z")
      END IF

      SELECT CASE (Field(ind_b)%field_type)
      CASE (field_T)
         status = NF90_DEF_DIM(ncid, 'cell_i', ncell, ncellId)
         status = NF90_DEF_DIM(ncid, 'nvert_i', nvert, nvertid)
         status = NF90_DEF_VAR(ncid, 'lon_i', NF90_DOUBLE, (/ncellId/), lonId)
         status = NF90_DEF_VAR(ncid, 'lat_i', NF90_DOUBLE, (/ncellId/), latId)
         status = NF90_PUT_ATT(ncid, lonId, "bounds", "bounds_lon_i")
         status = NF90_PUT_ATT(ncid, latId, "bounds", "bounds_lat_i")
         status = NF90_DEF_VAR(ncid, 'bounds_lon_i', NF90_DOUBLE, (/nvertid, ncellId/), bounds_lonId)
         status = NF90_DEF_VAR(ncid, 'bounds_lat_i', NF90_DOUBLE, (/nvertid, ncellId/), bounds_latId)
      CASE (field_Z)
         status = NF90_DEF_DIM(ncid, 'cell_v', ncell, ncellId)
         status = NF90_DEF_DIM(ncid, 'nvert_v', nvert, nvertid)
         status = NF90_DEF_VAR(ncid, 'lon_v', NF90_DOUBLE, (/ncellId/), lonId)
         status = NF90_DEF_VAR(ncid, 'lat_v', NF90_DOUBLE, (/ncellId/), latId)
         status = NF90_PUT_ATT(ncid, lonId, "bounds", "bounds_lon_v")
         status = NF90_PUT_ATT(ncid, latId, "bounds", "bounds_lat_v")
         status = NF90_DEF_VAR(ncid, 'bounds_lon_v', NF90_DOUBLE, (/nvertid, ncellId/), bounds_lonId)
         status = NF90_DEF_VAR(ncid, 'bounds_lat_v', NF90_DOUBLE, (/nvertid, ncellId/), bounds_latId)
      END SELECT

      status = NF90_PUT_ATT(ncid, lonId, "long_name", "longitude")
      status = NF90_PUT_ATT(ncid, lonId, "units", "degrees_east")
      status = NF90_PUT_ATT(ncid, latId, "long_name", "latitude")
      status = NF90_PUT_ATT(ncid, latId, "units", "degrees_north")

      SELECT CASE (Field(ind_b)%ndim)
      CASE (2)
         IF (once) THEN
            status = NF90_DEF_VAR(ncid, TRIM(ADJUSTL(name)), ncprec, &
                                  (/ncellId/), FieldVarId(NbField)%nc_id(1))
         ELSE
            status = NF90_DEF_VAR(ncid, TRIM(ADJUSTL(name)), ncprec, &
                                  (/ncellId, timeId/), FieldVarId(NbField)%nc_id(1))
         END IF
         CALL put_att_coordinates(1)
      CASE (3)
         IF (once) THEN
            status = NF90_DEF_VAR(ncid, TRIM(ADJUSTL(name)), ncprec, &
                                  (/ncellId, dim3id/), FieldVarId(NbField)%nc_id(1))
         ELSE
            status = NF90_DEF_VAR(ncid, TRIM(ADJUSTL(name)), ncprec, &
                                  (/ncellId, dim3id, timeId/), FieldVarId(NbField)%nc_id(1))
         END IF
         CALL put_att_coordinates(1)
      CASE (4)
         DO i = 1, FieldVarId(NbField)%size
            IF (once) THEN
               status = NF90_DEF_VAR(ncid, TRIM(ADJUSTL(name))//int2str(i), ncprec, (/ncellId, dim3id/), &
                                     FieldVarId(NbField)%nc_id(i))
            ELSE
               status = NF90_DEF_VAR(ncid, TRIM(ADJUSTL(name))//int2str(i), ncprec, (/ncellId, dim3id, timeId/), &
                                     FieldVarId(NbField)%nc_id(i))
            END IF
            CALL put_att_coordinates(i)
            status = NF90_PUT_ATT(ncid, FieldVarId(NbField)%nc_id(i), "coordinates", "lon_i lat_i")
         END DO
      END SELECT

      status = NF90_ENDDEF(ncid)

      if (level_size > 0) status = NF90_PUT_VAR(ncid, levId, (/(l, l=1, level_size)/))

      ALLOCATE (lon(ncell), lat(ncell))
      ALLOCATE (bounds_lon(nvert, ncell), bounds_lat(nvert, ncell))
      n_begin = 0
      DO ind = ind_b, ind_e
         n_end = n_begin + cells(ind)%ncell
!        PRINT *, 'create_header_gen ', n_begin, n_end, SHAPE(cells(ind)%bnds_lon), SHAPE(bounds_lon)
         lat(n_begin + 1:n_end) = cells(ind)%lat(:)*180./Pi
         lon(n_begin + 1:n_end) = cells(ind)%lon(:)*180./Pi
         bounds_lon(:, n_begin + 1:n_end) = cells(ind)%bnds_lon(:, :)*180./Pi
         bounds_lat(:, n_begin + 1:n_end) = cells(ind)%bnds_lat(:, :)*180./Pi
         n_begin = n_end
      END DO
      status = NF90_PUT_VAR(ncid, lonid, REAL(lon, r8), start=(/1/), count=(/ncell/))
      status = NF90_PUT_VAR(ncid, latid, REAL(lat, r8), start=(/1/), count=(/ncell/))
      status = NF90_PUT_VAR(ncid, bounds_lonId, REAL(bounds_lon, r8), start=(/1, 1/), count=(/nvert, ncell/))
      status = NF90_PUT_VAR(ncid, bounds_latId, REAL(bounds_lat, r8), start=(/1, 1/), count=(/nvert, ncell/))

      DEALLOCATE (lon, lat, bounds_lon, bounds_lat)

   CONTAINS

      SUBROUTINE put_att_coordinates(ind)
         INTEGER :: ind
         SELECT CASE (Field(ind_b)%field_type)
         CASE (field_T)
            status = NF90_PUT_ATT(ncid, FieldVarId(NbField)%nc_id(ind), &
                                  "coordinates", "lon_i lat_i")
         CASE (field_Z)
            status = NF90_PUT_ATT(ncid, FieldVarId(NbField)%nc_id(ind), &
                                  "coordinates", "lon_v lat_v")
         END SELECT
      END SUBROUTINE put_att_coordinates

   END SUBROUTINE Create_Header_gen

   SUBROUTINE Close_files
      USE netcdf_mod
      INTEGER :: i, k, status
!$OMP MASTER
      DO i = 1, NbField
         status = NF90_CLOSE(FieldId(i))
      END DO
!$OMP END MASTER
   END SUBROUTINE Close_files

end module write_field_mod
