MODULE read_pmesh_mod
  ! Read local part of partitioned mesh
  USE netcdf_mod
  IMPLICIT NONE
  SAVE
  PRIVATE

  CHARACTER(LEN=*),PARAMETER :: meshfile="input/pmesh.nc"

  INTERFACE read_from_pmesh
     MODULE PROCEDURE read_from_pmesh_i0
     MODULE PROCEDURE read_from_pmesh_i1
     MODULE PROCEDURE read_from_pmesh_i2
     MODULE PROCEDURE read_from_pmesh_r1
     MODULE PROCEDURE read_from_pmesh_r2
     MODULE PROCEDURE read_from_pmesh_r3
  END INTERFACE

  TYPE t_pmesh
     ! this type describes the local piece of unstructured mesh managed by a single MPI process
     INTEGER :: &              ! number of geometric elements, halo included
          primal_num, &        !     primal cells
          edge_num,   &        !     edges
          dual_num             !     dual cells

     INTEGER, ALLOCATABLE ::  &
          primal_glo(:),      &  ! global indices for : primal cells
          dual_glo(:),        &  !                      dual cells
          edge_glo(:),        &  !                      edges
          primal_deg(:),      &  ! primal cells : degree (number of edges)
          primal_vertex(:,:), &  !                indices of vertices
          primal_edge(:,:),   &  !                indices of edges
          primal_ne(:,:),     &  !                sign to apply when computing div
          dual_deg(:),        &  ! dual cells :   degree
          dual_edge(:,:),     &  !                edges
          dual_vertex(:,:),   &  !                vertices
          dual_ne(:,:),       &  !                sign to apply when computing curl
          left(:),right(:),   &  ! edges :        left and right cells
          up(:), down(:),     &  !                up and down dual cells
          trisk_deg(:),       &  !                size of TRiSK stencil
          trisk(:,:)             !                edges in TRiSK stencil

     INTEGER, ALLOCATABLE :: & ! halo exchange descriptor for :
          primal_recv(:), &  !     receiving primal cells
          primal_send(:), &  !     sending primal cells
          edge_recv(:),   &  !     receiving edges
          edge_send(:),   &  !     sending edges
          dual_recv(:),   &  !     receiving vertices
          dual_send(:)       !     sending vertices

     REAL, ALLOCATABLE :: & ! for dynamics
          Ai(:),      & ! primal cells : area
          lon_i(:),   & !                longitude
          lat_i(:),   & !                latitude
          Av(:),      & ! dual cells   : area
          lon_v(:),   & !                longitude
          lat_v(:),   & !                latitude
          fv(:),      & !                Coriolis factor * area
          Riv2(:,:),  & !                weights to average from primal cells to dual cells
          de(:),      & ! Edges :        length (dual cell edge)
          le(:),      & !                length (primal cell edge)
          le_de(:),   & !                ratio le/de
          lon_e(:),   & !                longitude
          lat_e(:),   & !                latitude
          angle_e(:), & !                angle between zonal direction and dual cell edge
          wee(:,:),   & !                weights of TRiSK perp operator
          perot(:,:,:), & !              weights of Perot reconstruction operator
          grad(:,:,:)     !              weights of vector gradient operator

     REAL, ALLOCATABLE :: & ! for XIOS : lon/lat of cell vertices (boundaries)
          primal_bounds_lon(:,:), primal_bounds_lat(:,:) , & ! primal mesh
          dual_bounds_lon(:,:),   dual_bounds_lat(:,:)       ! dual mesh

  END type t_pmesh

  PUBLIC :: t_pmesh, read_pmesh, angle_to_vector, pmesh_to_cellsets

CONTAINS

  SUBROUTINE get_shape(mpi_rank,id_nc,name,local_size, id_var, shp, start)
    INTEGER,      INTENT(IN)  :: mpi_rank, local_size
    INTEGER,      INTENT(IN)  :: id_nc   ! NetCDF ID of file
    CHARACTER(*), INTENT(IN)  :: name    ! variable name
    INTEGER,      INTENT(OUT) :: id_var  !          NetCDF ID
    INTEGER,      INTENT(OUT) :: shp(:)  !          shape, including last dimension (mpi_rank)
    INTEGER,      INTENT(OUT) :: start(:)
    INTEGER :: dimids(20), rank, status, ndims, i
    rank=SIZE(shp)
    status = nf90_inq_varid(id_nc, name,id_var)
    IF(status /= 0) THEN
       PRINT *, "Routine read_from_pmesh unable able to read ', TRIM(name), 'from mesh file ", TRIM(meshfile)
       STOP "Exit"
    ENDIF
    ! inquire dimensions
    status = nf90_inquire_variable(id_nc,id_var,dimids=dimIDs,ndims=ndims)
    IF(ndims/=rank) THEN
       PRINT *, 'Variable ', TRIM(name), ' from mesh file ',TRIM(meshfile), &
            ' has ', ndims, ' dimensions while we expect ', rank
       STOP
    END IF
    DO i = 1,rank
       status = nf90_inquire_dimension(id_nc, dimIDs(i), len = shp(i))
    END DO
    PRINT *,"NetCDF shape of array, ",name,":", shp
    ! Note : indices given by ncdump are in C order
    !   while indices in the Fortran NetCDF interface are in Fortran order
    !   The unlimited dimension is the last one in Fortran order
    !   In the NetCDF Fortran interface, indices start at 1
    start(:)=1
    ! The last dimension (i=rank) corresponds to the MPI dimension => let size:=1, start:=mpi_rank+1
    shp(rank)=1
    start(rank)=mpi_rank+1
    IF(rank>1) THEN
       ! let rank be now the index of the dimension distributed across MPI process
       SELECT CASE(rank-1)
       CASE(1,2)
          rank=rank-1 
       CASE(3)
          rank=2
       CASE DEFAULT
          PRINT *, 'rank = ', rank-1, ' while read_pmesh supports reading arrays up to 3D'
          STOP
       END SELECT
       PRINT *, "  ", TRIM(name), " : local size = ", ", NetCDF size = ", shp(rank)
       IF(shp(rank)<local_size) THEN
          PRINT *, "Fatal : local size ", local_size, " of ", TRIM(name), " is larger than NetCDF size ", shp(rank)
          STOP
       END IF
       shp(rank)=local_size
    END IF
  END SUBROUTINE get_shape

  SUBROUTINE NoError(name, status)
    INTEGER, INTENT(IN)      :: status
    CHARACTER(*), INTENT(IN) :: name
    IF(status /= nf90_NoErr) THEN
       print *, "Error when reading ", name, "from grid file ", TRIM(meshfile)
       STOP "Exit"
    ENDIF
  END SUBROUTINE NoError

  SUBROUTINE read_from_pmesh_i0(id_nc, mpi_rank, name, data)
    INTEGER,      INTENT(IN)       :: id_nc, mpi_rank
    CHARACTER(*), INTENT(IN)       :: name
    INTEGER,      INTENT(OUT)      :: data
    INTEGER :: shp(1), start(1), id_var
    CALL get_shape(mpi_rank,id_nc,name,0, id_var,shp,start)
    CALL NoError(name, nf90_get_var(id_nc,id_var, data, start=start))
    PRINT *,"  First value of ",name,":", data
  END SUBROUTINE read_from_pmesh_i0

  SUBROUTINE read_from_pmesh_i1(id_nc, mpi_rank, name, local_size, data)
    INTEGER,      INTENT(IN)          :: id_nc, mpi_rank, local_size
    CHARACTER(*), INTENT(IN)          :: name
    INTEGER, ALLOCATABLE, INTENT(OUT) :: data(:)
    INTEGER :: shp(2), start(2), id_var
    CALL get_shape(mpi_rank,id_nc,name,local_size, id_var,shp,start)
    ALLOCATE(data(shp(1)))
    CALL NoError(name,nf90_get_var(id_nc,id_var, data,start=start,count=shp))
    PRINT *,"  First value of array, ",name,":", data(1)
  END SUBROUTINE read_from_pmesh_i1

  SUBROUTINE read_from_pmesh_i2(id_nc, mpi_rank, name, local_size, data)
    INTEGER,      INTENT(IN)          :: id_nc, mpi_rank, local_size
    CHARACTER(*), INTENT(IN)          :: name
    INTEGER, ALLOCATABLE, INTENT(OUT) :: data(:,:)
    INTEGER :: shp(3), start(3), id_var
    CALL get_shape(mpi_rank,id_nc,name,local_size, id_var,shp,start)
    ALLOCATE(data(shp(1),shp(2)))
    CALL NoError(name,nf90_get_var(id_nc,id_var, data,start=start,count=shp))
    PRINT *,"  First value of array, ",name,":", data(1,1)
  END SUBROUTINE read_from_pmesh_i2

  SUBROUTINE read_from_pmesh_r1(id_nc, mpi_rank, name, local_size, data)
    INTEGER,      INTENT(IN)       :: id_nc, mpi_rank, local_size
    CHARACTER(*), INTENT(IN)       :: name
    REAL, ALLOCATABLE, INTENT(OUT) :: data(:)
    INTEGER :: shp(2), start(2), id_var
    CALL get_shape(mpi_rank,id_nc,name,local_size, id_var,shp,start)
    ALLOCATE(data(shp(1)))
    CALL NoError(name,nf90_get_var(id_nc,id_var, data,start=start,count=shp))
    PRINT *,"  First value of array, ",name,":", data(1)
  END SUBROUTINE read_from_pmesh_r1

  SUBROUTINE read_from_pmesh_r2(id_nc, mpi_rank, name, local_size, data)
    INTEGER,      INTENT(IN)       :: id_nc, mpi_rank, local_size
    CHARACTER(*), INTENT(IN)       :: name
    REAL, ALLOCATABLE, INTENT(OUT) :: data(:,:)
    INTEGER :: shp(3), start(3), id_var
    CALL get_shape(mpi_rank,id_nc,name,local_size, id_var,shp,start)
    ALLOCATE(data(shp(1),shp(2)))
    CALL NoError(name,nf90_get_var(id_nc,id_var, data,start=start,count=shp))
    PRINT *,"  First value of array, ",name,":", data(1,1)
  END SUBROUTINE read_from_pmesh_r2

  SUBROUTINE read_from_pmesh_r3(id_nc, mpi_rank, name, local_size, data)
    INTEGER,      INTENT(IN)       :: id_nc, mpi_rank, local_size
    CHARACTER(*), INTENT(IN)       :: name
    REAL, ALLOCATABLE, INTENT(OUT) :: data(:,:,:)
    INTEGER :: shp(4), start(4), id_var
    CALL get_shape(mpi_rank,id_nc,name,local_size, id_var,shp,start)
    ALLOCATE(data(shp(1),shp(2),shp(3)))
    CALL NoError(name,nf90_get_var(id_nc,id_var, data,start=start,count=shp))
    PRINT *,"  First value of array, ",name,":", data(1,1,1)
  END SUBROUTINE read_from_pmesh_r3

  SUBROUTINE read_pmesh(mpi_rank, pmesh)
    INTEGER,       INTENT(IN)  :: mpi_rank
    TYPE(t_pmesh), INTENT(OUT) :: pmesh
    CHARACTER(LEN=80) :: description
    INTEGER :: primal_num, edge_num, dual_num, &
         status, id_nc, descriptionLength, local_size

    PRINT *,"------------------ READING FILE " , meshfile, "----------------------- "

    status = nf90_open(meshfile, NF90_NOWRITE, id_nc)
    if (status /= nf90_noerr) then
       print *, trim(nf90_strerror(status))
       stop "Error reading file"
    end if

    !------------------- Global attributes -----------------------

    status = nf90_inquire_attribute(id_nc, nf90_global, "description", len=descriptionLength)
    IF(status /= nf90_noerr .or. len(description) < descriptionLength) THEN
       print *, "Not enough space to put NetCDF attribute values."
       STOP "Error reading file"
    ENDIF

    status = nf90_get_att(id_nc, nf90_global, "description", description)
    print *,"Data file description :",description

    !------------------- Indexing and halo exchange --------------

    CALL read_from_pmesh(id_nc,mpi_rank,'primal_num', primal_num)
    CALL read_from_pmesh(id_nc,mpi_rank,'dual_num',   dual_num)
    CALL read_from_pmesh(id_nc,mpi_rank,'edge_num',   edge_num)
    CALL read_from_pmesh(id_nc,mpi_rank,'primal_glo', primal_num, pmesh%primal_glo)
    CALL read_from_pmesh(id_nc,mpi_rank,'dual_glo',   dual_num,  pmesh%dual_glo)
    CALL read_from_pmesh(id_nc,mpi_rank,'edge_glo',   edge_num,  pmesh%edge_glo)

    CALL read_from_pmesh(id_nc,mpi_rank,'primal_recv_num', local_size)
    CALL read_from_pmesh(id_nc,mpi_rank,'primal_recv',     local_size, pmesh%primal_recv)
    CALL read_from_pmesh(id_nc,mpi_rank,'primal_send_num', local_size)
    CALL read_from_pmesh(id_nc,mpi_rank,'primal_send',     local_size, pmesh%primal_send)

    CALL read_from_pmesh(id_nc,mpi_rank,'edge_recv_num', local_size)
    CALL read_from_pmesh(id_nc,mpi_rank,'edge_recv',     local_size, pmesh%edge_recv)
    CALL read_from_pmesh(id_nc,mpi_rank,'edge_send_num', local_size)
    CALL read_from_pmesh(id_nc,mpi_rank,'edge_send',     local_size, pmesh%edge_send)

    CALL read_from_pmesh(id_nc,mpi_rank,'dual_recv_num', local_size)
    CALL read_from_pmesh(id_nc,mpi_rank,'dual_recv',     local_size, pmesh%dual_recv)
    CALL read_from_pmesh(id_nc,mpi_rank,'dual_send_num', local_size)
    CALL read_from_pmesh(id_nc,mpi_rank,'dual_send',     local_size, pmesh%dual_send)

    !----------------------- Stencils ---------------------------

    CALL read_from_pmesh(id_nc,mpi_rank,'primal_deg',    primal_num, pmesh%primal_deg)
    CALL read_from_pmesh(id_nc,mpi_rank,'primal_vertex', primal_num, pmesh%primal_vertex)
    CALL read_from_pmesh(id_nc,mpi_rank,'primal_edge',   primal_num, pmesh%primal_edge)
    CALL read_from_pmesh(id_nc,mpi_rank,'primal_ne',     primal_num, pmesh%primal_ne)

    CALL read_from_pmesh(id_nc,mpi_rank,'dual_deg',    dual_num,  pmesh%dual_deg)
    CALL read_from_pmesh(id_nc,mpi_rank,'dual_edge',   dual_num,  pmesh%dual_edge)
    CALL read_from_pmesh(id_nc,mpi_rank,'dual_vertex', dual_num,  pmesh%dual_vertex)
    CALL read_from_pmesh(id_nc,mpi_rank,'dual_ne',     dual_num,  pmesh%dual_ne)

    CALL read_from_pmesh(id_nc,mpi_rank,'left',     edge_num, pmesh%left)
    CALL read_from_pmesh(id_nc,mpi_rank,'right',    edge_num, pmesh%right)
    CALL read_from_pmesh(id_nc,mpi_rank,'up',       edge_num, pmesh%up)
    CALL read_from_pmesh(id_nc,mpi_rank,'down',     edge_num, pmesh%down)
    CALL read_from_pmesh(id_nc,mpi_rank,'trisk_deg',edge_num, pmesh%trisk_deg)
    CALL read_from_pmesh(id_nc,mpi_rank,'trisk',    edge_num, pmesh%trisk)

    !------------------- Lengths, areas, angles and weights ---------------

    CALL read_from_pmesh(id_nc,mpi_rank,'Ai',    primal_num, pmesh%Ai)
    CALL read_from_pmesh(id_nc,mpi_rank,'lon_i', primal_num, pmesh%lon_i)
    CALL read_from_pmesh(id_nc,mpi_rank,'lat_i', primal_num, pmesh%lat_i)
    CALL read_from_pmesh(id_nc,mpi_rank,'primal_grad',  primal_num, pmesh%grad)
    CALL read_from_pmesh(id_nc,mpi_rank,'primal_perot', primal_num, pmesh%perot)

    CALL read_from_pmesh(id_nc,mpi_rank,'Av',    dual_num, pmesh%Av)
    CALL read_from_pmesh(id_nc,mpi_rank,'lon_v', dual_num, pmesh%lon_v)
    CALL read_from_pmesh(id_nc,mpi_rank,'lat_v', dual_num, pmesh%lat_v)
    CALL read_from_pmesh(id_nc,mpi_rank,'fv',    dual_num, pmesh%fv)
    CALL read_from_pmesh(id_nc,mpi_rank,'Riv2',  dual_num, pmesh%Riv2)

    CALL read_from_pmesh(id_nc,mpi_rank,'de',      edge_num, pmesh%de)
    CALL read_from_pmesh(id_nc,mpi_rank,'le',      edge_num, pmesh%le)
    CALL read_from_pmesh(id_nc,mpi_rank,'le_de',   edge_num, pmesh%le_de)
    CALL read_from_pmesh(id_nc,mpi_rank,'lon_e',   edge_num, pmesh%lon_e)
    CALL read_from_pmesh(id_nc,mpi_rank,'lat_e',   edge_num, pmesh%lat_e)
    CALL read_from_pmesh(id_nc,mpi_rank,'angle_e', edge_num, pmesh%angle_e)
    CALL read_from_pmesh(id_nc,mpi_rank,'wee',     edge_num, pmesh%wee)

    !-------------------- Cell boundaries => XIOS --------------------

    CALL read_from_pmesh(id_nc,mpi_rank,'primal_bounds_lon', primal_num, pmesh%primal_bounds_lon)
    CALL read_from_pmesh(id_nc,mpi_rank,'primal_bounds_lat', primal_num, pmesh%primal_bounds_lat)
    CALL read_from_pmesh(id_nc,mpi_rank,'dual_bounds_lon', dual_num, pmesh%dual_bounds_lon)
    CALL read_from_pmesh(id_nc,mpi_rank,'dual_bounds_lat', dual_num, pmesh%dual_bounds_lat)

    pmesh%primal_num = primal_num
    pmesh%dual_num   = dual_num
    pmesh%edge_num   = edge_num

  END SUBROUTINE read_pmesh

  SUBROUTINE angle_to_vector(lon,lat,angle, vec)
    REAL, INTENT(IN) :: lon(:), lat(:), angle(:)
    REAL, INTENT(OUT) :: vec(:,:)
    REAL :: clon, slon, clat, slat, & ! COS/SIN of lon/lat
         elon(3), elat(3) ! lon/lat unit vectors
    INTEGER :: ij
    DO ij=1,SIZE(angle)
       clon = COS(lon(ij))
       slon = SIN(lon(ij))
       clat = COS(lat(ij))
       slat = SIN(lat(ij))
       ! x,y,z = clat*clon, clat*slon, slat
       elon = (/ -slon, clon, 0. /)
       elat = (/ -slat*clon, -slat*slon, clat /)
       vec(:,ij) = COS(angle(ij))*elon + SIN(angle(ij))*elat
    END DO
  END SUBROUTINE angle_to_vector

  SUBROUTINE pmesh_to_cellsets(mesh, primal, dual, edge)
    USE math_const,  ONLY : radian_to_degree
    USE cellset_mod, ONLY : t_cellset
    TYPE(t_pmesh),   INTENT(IN)  :: mesh
    TYPE(t_cellset), INTENT(OUT) :: primal, dual, edge

    REAL :: edge_bounds_lon(2,SIZE(mesh%lat_e)), edge_bounds_lat(2,SIZE(mesh%lat_e))
    INTEGER :: ij, vup, vdown

    DO ij=1, SIZE(mesh%lat_e)
       vdown = mesh%down(ij)
       vup = mesh%up(ij)
       edge_bounds_lon(1,ij) = mesh%lon_v(vdown)*radian_to_degree
       edge_bounds_lat(1,ij) = mesh%lat_v(vdown)*radian_to_degree
       edge_bounds_lon(2,ij) = mesh%lon_v(vup)*radian_to_degree
       edge_bounds_lat(2,ij) = mesh%lat_v(vup)*radian_to_degree
    END DO
    CALL pmesh_to_cellset(primal, mesh%primal_glo, mesh%primal_recv, &
         mesh%lon_i, mesh%lat_i, mesh%primal_bounds_lon, mesh%primal_bounds_lat)
    CALL pmesh_to_cellset(dual, mesh%dual_glo, mesh%dual_recv, &
         mesh%lon_v, mesh%lat_v, mesh%dual_bounds_lon, mesh%dual_bounds_lat)
    CALL pmesh_to_cellset(edge, mesh%edge_glo, mesh%edge_recv, &
         mesh%lon_e, mesh%lat_e, edge_bounds_lon, edge_bounds_lat)
  END SUBROUTINE pmesh_to_cellsets

  SUBROUTINE pmesh_to_cellset(cells, glo, recv, lon, lat, bnds_lon, bnds_lat)
    USE cellset_mod, ONLY : t_cellset
    TYPE(t_cellset), INTENT(OUT) :: cells
    INTEGER, INTENT(IN)  :: glo(:), recv(:) ! halo exchange descriptor
    REAL,    INTENT(IN)  :: lon(:), lat(:), bnds_lon(:,:), bnds_lat(:,:)
    ! We first  determines which cells are our own cells. The assumption is that
    ! any cells received during halo exchange are not our own.
    LOGICAL :: own(SIZE(lon))
    ! halo exchange descriptor stucture (see Python/test/py/gen_mpi_mesh_info.py) :
    !     recv(1) is the number of subsequent chunks
    !     recv(2:N+3) is the first chunk, received from MPI rank recv(2) with N=recv(3) the size of the chunk
    !     recv(i+3) (i=1,...N) is the local index of the cell received from rank halo(2)
    !     the next chunk start at N+4
    INTEGER :: nchunk, recv_rank, chunk_len, ichunk, icell, pos, ij
    own(:)=.TRUE.
    pos=1
    nchunk=next()
    DO ichunk = 1, nchunk
       recv_rank = next()
       chunk_len = next()
       DO icell = 1, chunk_len
          ij=next()
          own(ij)=.FALSE.
       END DO
    END DO

    CALL create_own_cellset(cells, COUNT(own), SIZE(bnds_lat,1), own, glo, lon, lat, bnds_lon, bnds_lat)

  CONTAINS
    FUNCTION next() RESULT(val)
      INTEGER :: val
      val = recv(pos)
      pos = pos+1
    END FUNCTION next

  END SUBROUTINE pmesh_to_cellset
  
  SUBROUTINE create_own_cellset(cells, n_own, max_deg, own, glo, lon, lat, bnds_lon, bnds_lat)
    USE cellset_mod, ONLY : t_cellset, create_cellset
    TYPE(t_cellset), INTENT(OUT) :: cells
    INTEGER, INTENT(IN) :: n_own, max_deg
    LOGICAL, INTENT(IN) :: own(:)
    INTEGER, INTENT(IN) :: glo(:)
    REAL,    INTENT(IN) :: lon(:), lat(:), bnds_lon(:,:), bnds_lat(:,:)
    ! arrays below collect values from arrays above only where own(ij) is .TRUE.
    INTEGER :: loc_own(n_own), glo_own(n_own)
    REAL    :: lon_own(n_own), lat_own(n_own), &
    &          bnds_lon_own(max_deg, n_own), bnds_lat_own(max_deg, n_own)
    INTEGER :: pos, ij
    pos=1
    DO ij=1, SIZE(lon)
       IF(own(ij)) THEN
          loc_own(pos)  = ij
          glo_own(pos) = glo(ij)
          lon_own(pos) = lon(ij)
          lat_own(pos) = lat(ij)
          bnds_lon_own(:,pos) = bnds_lon(:,ij)
          bnds_lat_own(:,pos) = bnds_lat(:,ij)
          pos=pos+1
       END IF
    END DO

    CALL create_cellset(cells, loc_own, glo_own, lon_own, lat_own, bnds_lon_own, bnds_lat_own)
    
  END SUBROUTINE create_own_cellset
  
END MODULE read_pmesh_mod
