! $Id: iophy.F90 5135 2024-07-26 17:20:23Z abarral $

module iophy

  ! abd  REAL,PRIVATE,ALLOCATABLE,DIMENSION(:),save :: io_lat
  ! abd  REAL,PRIVATE,ALLOCATABLE,DIMENSION(:),save :: io_lon
  REAL, ALLOCATABLE, DIMENSION(:), save :: io_lat
  REAL, ALLOCATABLE, DIMENSION(:), save :: io_lon
  INTEGER, save :: phys_domain_id
  INTEGER, save :: npstn
  INTEGER, ALLOCATABLE, DIMENSION(:), save :: nptabij


  ! interfaces for both IOIPSL and XIOS
  INTERFACE histwrite_phy
    MODULE PROCEDURE histwrite2d_phy, histwrite3d_phy, histwrite2d_xios, histwrite3d_xios
  END INTERFACE

  ! interfaces for both IOIPSL and XIOS
  INTERFACE histbeg_phy_all
    MODULE PROCEDURE histbeg_phy, histbeg_phyxios
  END INTERFACE

CONTAINS

  SUBROUTINE init_iophy_new(rlat, rlon)
    USE dimphy, ONLY: klon
    USE lmdz_phys_para, ONLY: gather, bcast, &
            jj_nb, jj_begin, jj_end, ii_begin, ii_end, &
            mpi_size, mpi_rank, klon_mpi, &
            is_sequential, is_south_pole_dyn
    USE lmdz_grid_phy, ONLY: nbp_lon, nbp_lat, klon_glo
    USE lmdz_print_control, ONLY: lunout, prt_level
    USE lmdz_grid_phy, ONLY: nbp_lon, nbp_lat
    USE ioipsl, ONLY: flio_dom_set
    USE lmdz_wxios, ONLY: wxios_domain_param, using_xios
    IMPLICIT NONE
    REAL, DIMENSION(klon), INTENT(IN) :: rlon
    REAL, DIMENSION(klon), INTENT(IN) :: rlat

    REAL, DIMENSION(klon_glo) :: rlat_glo
    REAL, DIMENSION(klon_glo) :: rlon_glo

    INTEGER, DIMENSION(2) :: ddid
    INTEGER, DIMENSION(2) :: dsg
    INTEGER, DIMENSION(2) :: dsl
    INTEGER, DIMENSION(2) :: dpf
    INTEGER, DIMENSION(2) :: dpl
    INTEGER, DIMENSION(2) :: dhs
    INTEGER, DIMENSION(2) :: dhe
    INTEGER :: i
    INTEGER :: data_ibegin, data_iend

    CALL gather(rlat, rlat_glo)
    CALL bcast(rlat_glo)
    CALL gather(rlon, rlon_glo)
    CALL bcast(rlon_glo)

    !$OMP MASTER
    ALLOCATE(io_lat(nbp_lat))
    io_lat(1) = rlat_glo(1)
    io_lat(nbp_lat) = rlat_glo(klon_glo)
    IF ((nbp_lon * nbp_lat) > 1) THEN
      DO i = 2, nbp_lat - 1
        io_lat(i) = rlat_glo(2 + (i - 2) * nbp_lon)
      ENDDO
    ENDIF

    ALLOCATE(io_lon(nbp_lon))
    IF ((nbp_lon * nbp_lat) > 1) THEN
      io_lon(:) = rlon_glo(2:nbp_lon + 1)
    ELSE
      io_lon(1) = rlon_glo(1)
    ENDIF
    !! (I) dtnb   : total number of domains
    !! (I) dnb    : domain number
    !! (I) did(:) : distributed dimensions identifiers
    !!              (up to 5 dimensions are supported)
    !! (I) dsg(:) : total number of points for each dimension
    !! (I) dsl(:) : local number of points for each dimension
    !! (I) dpf(:) : position of first local point for each dimension
    !! (I) dpl(:) : position of last local point for each dimension
    !! (I) dhs(:) : start halo size for each dimension
    !! (I) dhe(:) : end halo size for each dimension
    !! (C) cdnm   : Model domain definition name.
    !!              The names actually supported are :
    !!              "BOX", "APPLE", "ORANGE".
    !!              These names are case insensitive.
    ddid = (/ 1, 2 /)
    dsg = (/ nbp_lon, nbp_lat /)
    dsl = (/ nbp_lon, jj_nb /)
    dpf = (/ 1, jj_begin /)
    dpl = (/ nbp_lon, jj_end /)
    dhs = (/ ii_begin - 1, 0 /)
    IF (mpi_rank==mpi_size - 1) THEN
      dhe = (/0, 0/)
    else
      dhe = (/ nbp_lon - ii_end, 0 /)
    endif

#ifndef CPP_IOIPSL_NO_OUTPUT
    CALL flio_dom_set(mpi_size, mpi_rank, ddid, dsg, dsl, dpf, dpl, dhs, dhe, &
            'APPLE', phys_domain_id)
#endif
    IF (using_xios) THEN
      ! Set values for the mask:
      IF (mpi_rank == 0) THEN
        data_ibegin = 0
      ELSE
        data_ibegin = ii_begin - 1
      END IF

      IF (mpi_rank == mpi_size - 1) THEN
        data_iend = nbp_lon
      ELSE
        data_iend = ii_end + 1
      END IF

      IF (prt_level>=10) THEN
        WRITE(lunout, *) "init_iophy_new: mpirank=", mpi_rank, " iibegin=", ii_begin, " ii_end=", ii_end, " jjbegin=", jj_begin, " jj_nb=", jj_nb, " jj_end=", jj_end
        WRITE(lunout, *) "init_iophy_new: mpirank=", mpi_rank, " nbp_lon=", nbp_lon, " nbp_lat=", nbp_lat
        WRITE(lunout, *) "init_iophy_new: mpirank=", mpi_rank, " data_ibegin=", data_ibegin, " data_iend=", data_iend
        WRITE(lunout, *) "init_iophy_new: mpirank=", mpi_rank, " data_ibegin=", data_ibegin, " data_iend=", data_iend
        WRITE(lunout, *) "init_iophy_new: mpirank=", mpi_rank, " is_south_pole=", is_south_pole_dyn
      endif

      ! Initialize the XIOS domain coreesponding to this process:
      CALL wxios_domain_param("dom_glo")
      !      CALL wxios_domain_param("dom_glo", is_sequential, nbp_lon, jj_nb, nbp_lon, nbp_lat, &
      !                              1, nbp_lon, ii_begin, ii_end, jj_begin, jj_end,             &
      !                              klon_mpi+2*(nbp_lon-1), data_ibegin, data_iend,             &
      !                              io_lat, io_lon,is_south_pole_dyn,mpi_rank)
    ENDIF
    !$OMP END MASTER

  END SUBROUTINE init_iophy_new

  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  SUBROUTINE histbeg_phy(name, itau0, zjulian, dtime, nhori, nid_day)
    USE lmdz_phys_para, ONLY: is_sequential, jj_begin, jj_end, jj_nb
    USE ioipsl, ONLY: histbeg
    USE lmdz_print_control, ONLY: prt_level, lunout
    USE lmdz_grid_phy, ONLY: nbp_lon
    IMPLICIT NONE

    CHARACTER*(*), INTENT(IN) :: name
    INTEGER, INTENT(IN) :: itau0
    REAL, INTENT(IN) :: zjulian
    REAL, INTENT(IN) :: dtime
    INTEGER, INTENT(OUT) :: nhori
    INTEGER, INTENT(OUT) :: nid_day

    !$OMP MASTER
    IF (is_sequential) THEN
      CALL histbeg(name, nbp_lon, io_lon, jj_nb, io_lat(jj_begin:jj_end), &
              1, nbp_lon, 1, jj_nb, itau0, zjulian, dtime, nhori, nid_day)
    else
      CALL histbeg(name, nbp_lon, io_lon, jj_nb, io_lat(jj_begin:jj_end), &
              1, nbp_lon, 1, jj_nb, itau0, zjulian, dtime, nhori, nid_day, phys_domain_id)
    endif
    !$OMP END MASTER

  END SUBROUTINE  histbeg_phy

  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!


  ! SUBROUTINE histbeg_phyxios(name,itau0,zjulian,dtime,ffreq,lev,nhori,nid_day)
  SUBROUTINE histbeg_phyxios(name, ffreq, lev)
    USE lmdz_phys_para, ONLY: is_using_mpi, is_mpi_root
    USE lmdz_wxios, ONLY: wxios_add_file
    IMPLICIT NONE

    CHARACTER*(*), INTENT(IN) :: name
    !    INTEGER, INTENT(IN) :: itau0
    !    REAL,INTENT(IN) :: zjulian
    !    REAL,INTENT(IN) :: dtime
    character(LEN = *), INTENT(IN) :: ffreq
    INTEGER, INTENT(IN) :: lev
    !    INTEGER,INTENT(OUT) :: nhori
    !    INTEGER,INTENT(OUT) :: nid_day

    !$OMP MASTER

    ! ug OMP en chantier...
    IF((.NOT. is_using_mpi) .OR. is_mpi_root) THEN
      ! ug Création du fichier
      CALL wxios_add_file(name, ffreq, lev)
    END IF

    !$OMP END MASTER

  END SUBROUTINE histbeg_phyxios


  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  SUBROUTINE histwrite2d_phy(nid, lpoint, name, itau, field)
    USE dimphy, ONLY: klon
    USE lmdz_phys_para, ONLY: Gather_omp, grid1Dto2D_mpi, &
            is_sequential, klon_mpi_begin, klon_mpi_end, &
            jj_nb, klon_mpi
    USE ioipsl, ONLY: histwrite
    USE lmdz_grid_phy, ONLY: nbp_lon
    USE lmdz_abort_physic, ONLY: abort_physic
    IMPLICIT NONE

    INTEGER, INTENT(IN) :: nid
    logical, INTENT(IN) :: lpoint
    CHARACTER*(*), INTENT(IN) :: name
    INTEGER, INTENT(IN) :: itau
    REAL, DIMENSION(:), INTENT(IN) :: field
    REAL, DIMENSION(klon_mpi) :: buffer_omp
    INTEGER, ALLOCATABLE, DIMENSION(:) :: index2d
    REAL :: Field2d(nbp_lon, jj_nb)

    INTEGER :: ip
    REAL, ALLOCATABLE, DIMENSION(:) :: fieldok

    IF (size(field)/=klon) CALL abort_physic('iophy::histwrite2d', 'Field first dimension not equal to klon', 1)

    CALL Gather_omp(field, buffer_omp)
    !$OMP MASTER
    CALL grid1Dto2D_mpi(buffer_omp, Field2d)
    IF(.NOT.lpoint) THEN
      ALLOCATE(index2d(nbp_lon * jj_nb))
      ALLOCATE(fieldok(nbp_lon * jj_nb))
      CALL histwrite(nid, name, itau, Field2d, nbp_lon * jj_nb, index2d)
    else
      ALLOCATE(fieldok(npstn))
      ALLOCATE(index2d(npstn))

      IF(is_sequential) THEN
        !     klon_mpi_begin=1
        !     klon_mpi_end=klon
        DO ip = 1, npstn
          fieldok(ip) = buffer_omp(nptabij(ip))
        ENDDO
      else
        DO ip = 1, npstn
          !     PRINT*,'histwrite2d is_sequential npstn ip name nptabij',npstn,ip,name,nptabij(ip)
          IF(nptabij(ip)>=klon_mpi_begin.AND. &
                  nptabij(ip)<=klon_mpi_end) THEN
            fieldok(ip) = buffer_omp(nptabij(ip) - klon_mpi_begin + 1)
          ENDIF
        ENDDO
      endif
      CALL histwrite(nid, name, itau, fieldok, npstn, index2d)

    endif
    deallocate(index2d)
    deallocate(fieldok)
    !$OMP END MASTER
  END SUBROUTINE  histwrite2d_phy

  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  SUBROUTINE histwrite3d_phy(nid, lpoint, name, itau, field)
    USE dimphy, ONLY: klon
    USE lmdz_phys_para, ONLY: Gather_omp, grid1Dto2D_mpi, &
            is_sequential, klon_mpi_begin, klon_mpi_end, &
            jj_nb, klon_mpi
    USE ioipsl, ONLY: histwrite
    USE lmdz_grid_phy, ONLY: nbp_lon
    USE lmdz_abort_physic, ONLY: abort_physic
    IMPLICIT NONE

    INTEGER, INTENT(IN) :: nid
    logical, INTENT(IN) :: lpoint
    CHARACTER*(*), INTENT(IN) :: name
    INTEGER, INTENT(IN) :: itau
    REAL, DIMENSION(:, :), INTENT(IN) :: field  ! --> field(klon,:)
    REAL, DIMENSION(klon_mpi, size(field, 2)) :: buffer_omp
    REAL :: Field3d(nbp_lon, jj_nb, size(field, 2))
    INTEGER :: ip, n, nlev
    INTEGER, ALLOCATABLE, DIMENSION(:) :: index3d
    REAL, ALLOCATABLE, DIMENSION(:, :) :: fieldok

    IF (size(field, 1)/=klon) CALL abort_physic('iophy::histwrite3d', 'Field first dimension not equal to klon', 1)
    nlev = size(field, 2)

    CALL Gather_omp(field, buffer_omp)
    !$OMP MASTER
    CALL grid1Dto2D_mpi(buffer_omp, field3d)
    IF(.NOT.lpoint) THEN
      ALLOCATE(index3d(nbp_lon * jj_nb * nlev))
      ALLOCATE(fieldok(nbp_lon * jj_nb, nlev))
      CALL histwrite(nid, name, itau, Field3d, nbp_lon * jj_nb * nlev, index3d)
    else
      nlev = size(field, 2)
      ALLOCATE(index3d(npstn * nlev))
      ALLOCATE(fieldok(npstn, nlev))

      IF(is_sequential) THEN
        !      klon_mpi_begin=1
        !      klon_mpi_end=klon
        DO n = 1, nlev
          DO ip = 1, npstn
            fieldok(ip, n) = buffer_omp(nptabij(ip), n)
          ENDDO
        ENDDO
      else
        DO n = 1, nlev
          DO ip = 1, npstn
            IF(nptabij(ip)>=klon_mpi_begin.AND. &
                    nptabij(ip)<=klon_mpi_end) THEN
              fieldok(ip, n) = buffer_omp(nptabij(ip) - klon_mpi_begin + 1, n)
            ENDIF
          ENDDO
        ENDDO
      endif
      CALL histwrite(nid, name, itau, fieldok, npstn * nlev, index3d)
    endif
    deallocate(index3d)
    deallocate(fieldok)
    !$OMP END MASTER
  END SUBROUTINE  histwrite3d_phy

  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  ! VERSION DES HISTWRITE DEDIEES AU TOUT-XIOS-XML DEJA UTILISEE DANS PHYDEV

  SUBROUTINE histwrite2d_xios(field_name, field)
    USE dimphy, ONLY: klon
    USE lmdz_phys_para, ONLY: gather_omp, grid1Dto2D_mpi, &
            jj_nb, klon_mpi
    USE lmdz_xios, ONLY: xios_send_field
    USE lmdz_print_control, ONLY: prt_level, lunout
    USE lmdz_grid_phy, ONLY: nbp_lon
    USE lmdz_abort_physic, ONLY: abort_physic
    IMPLICIT NONE

    CHARACTER(LEN = *), INTENT(IN) :: field_name
    REAL, DIMENSION(:), INTENT(IN) :: field

    REAL, DIMENSION(klon_mpi) :: buffer_omp
    REAL :: Field2d(nbp_lon, jj_nb)

    IF (prt_level >= 10) WRITE(lunout, *)'Begin histrwrite2d_xios ', trim(field_name)

    IF (SIZE(field)/=klon) CALL abort_physic('iophy::histwrite2d_xios', 'Field first DIMENSION not equal to klon', 1)

    CALL Gather_omp(field, buffer_omp)
    !$OMP MASTER
    CALL grid1Dto2D_mpi(buffer_omp, Field2d)

    CALL xios_send_field(field_name, Field2d)
    !$OMP END MASTER

    IF (prt_level >= 10) WRITE(lunout, *)'End histrwrite2d_xios ', trim(field_name)
  END SUBROUTINE histwrite2d_xios

  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  ! VERSION DES HISTWRITE DEDIEES AU TOUT-XIOS-XML DEJA UTILISEE DANS PHYDEV

  SUBROUTINE histwrite3d_xios(field_name, field)
    USE dimphy, ONLY: klon, klev
    USE lmdz_phys_para, ONLY: gather_omp, grid1Dto2D_mpi, &
            jj_nb, klon_mpi
    USE lmdz_xios, ONLY: xios_send_field
    USE lmdz_print_control, ONLY: prt_level, lunout
    USE lmdz_grid_phy, ONLY: nbp_lon
    USE lmdz_abort_physic, ONLY: abort_physic
    IMPLICIT NONE

    CHARACTER(LEN = *), INTENT(IN) :: field_name
    REAL, DIMENSION(:, :), INTENT(IN) :: field ! --> field(klon,:)

    REAL, DIMENSION(klon_mpi, SIZE(field, 2)) :: buffer_omp
    REAL :: Field3d(nbp_lon, jj_nb, SIZE(field, 2))
    INTEGER :: ip, n, nlev

    IF (prt_level >= 10) WRITE(lunout, *)'Begin histrwrite3d_xios ', trim(field_name)

    !Et on.... écrit
    IF (SIZE(field, 1)/=klon) CALL abort_physic('iophy::histwrite3d', 'Field first DIMENSION not equal to klon', 1)
    nlev = SIZE(field, 2)

    CALL Gather_omp(field, buffer_omp)
    !$OMP MASTER
    CALL grid1Dto2D_mpi(buffer_omp, field3d)

    CALL xios_send_field(field_name, Field3d(:, :, 1:nlev))
    !$OMP END MASTER

    IF (prt_level >= 10) WRITE(lunout, *)'End histrwrite3d_xios ', trim(field_name)
  END SUBROUTINE histwrite3d_xios

END MODULE iophy