PROGRAM make_partition
   USE iso_c_binding ! For command-line interface
   USE partition_mod, ONLY: jpiglo, jpjglo, write_partition, nemo_partition, &
                            allocate_arrays, partition_mca_rk, cyclic_bc
   USE mask_mod, ONLY: imask, generate_mask
   IMPLICIT none

   INTEGER :: narg, iarg, arglen
   INTEGER :: ji, ierr
   INTEGER :: num_pes
   CHARACTER(len=100) :: argval
   INTEGER :: cores_per_node ! No. of cores per compute node on the machine
                             ! that NEMO will run on.
   LOGICAL :: regular ! True for regular decomp, false otherwise

   !!----------------------------------------------------------------------

   ! Initialise domain size
   jpiglo  = 0 ; jpjglo = 0
   num_pes = 0
   cores_per_node = 16
   regular   = .FALSE.        ! By default we do recursive partitioning
   cyclic_bc = .FALSE.        ! Cyclic boundary conditions off by default

   ! Parse the command-line arguments
   narg = COMMAND_ARGUMENT_COUNT()

   IF(narg == 0)THEN
      CALL print_usage()
      STOP
   END IF

   iarg = 1
   parse_args: DO WHILE(iarg <= narg)

      CALL GET_COMMAND_ARGUMENT(iarg,argval,arglen,ierr)

      IF(TRIM(argval) == '-d')THEN
         iarg = iarg + 1
         CALL GET_COMMAND_ARGUMENT(iarg,argval,arglen,ierr)

         ! Parse this argument. We expect it to be of the form <nx>x<ny>.
         ji = INDEX(argval,"x")
         IF(ji < 2 .OR. ji > (arglen-1))THEN
            WRITE(*,"('ERROR: unrecognised argument format: ',(A))") TRIM(argval)
         END IF
         READ(argval(1:ji-1),*,IOSTAT=ierr) jpiglo
         IF(ierr /= 0)THEN
            WRITE(*,"('ERROR: failed to parse jpiglo argument')")
            STOP
         END IF
         READ(argval(ji+1:),*,IOSTAT=ierr) jpjglo
         IF(ierr /= 0)THEN
            WRITE(*,"('ERROR: failed to parse jpjglo argument')")
            STOP
         END IF

      ELSE IF(TRIM(argval) == '-np')THEN
         iarg = iarg + 1
         CALL GET_COMMAND_ARGUMENT(iarg,argval,arglen,ierr)

         ! Parse this argument. We expect it to be of the form <np>.
         READ(argval,*,IOSTAT=ierr) num_pes
         IF(ierr /= 0)THEN
            WRITE(*,"('ERROR: failed to parse num_pes argument')")
            STOP
         END IF

      ELSE IF(TRIM(argval) == '-c')THEN
         iarg = iarg + 1
         CALL GET_COMMAND_ARGUMENT(iarg,argval,arglen,ierr)
         ! Parse this argument. We expect it to be of the form <cpn>.
         READ(argval,*,IOSTAT=ierr) cores_per_node

      ELSE IF(TRIM(argval) == '-r')THEN
         ! Switch-on regular partitioning
         regular = .TRUE.

      ELSE IF(TRIM(argval) == '-cyclic')THEN
         ! Switch-on cyclic boundary conditions
         cyclic_bc = .TRUE.

      ELSE IF(TRIM(argval) == '-h')THEN
         CALL print_usage()

      ELSE
         WRITE(*,"('ERROR: unrecognised argument: ',(A))") TRIM(argval)
         STOP

      END IF

      iarg = iarg + 1
   END DO parse_args

   ! Check that essential parameters have been supplied by user
   IF( num_pes == 0 )THEN
      WRITE (*,*) 'ERROR: you must specify num. PEs to decompose over!'
      CALL print_usage()
      STOP
   END IF

   ! Check whether domain dimensions have been specified
   IF( jpiglo < 1 .OR. jpjglo < 1 )THEN
      IF( regular )THEN
         ! Actually, I could read the bathy file if available to get the domain
         ! dimensions for this case.
         WRITE (*,*) 'ERROR: you must specify domain extent when doing regular partition!'
         CALL print_usage()
         STOP
      ELSE
         WRITE(*,*) 'Domain dimensions will be read from bathymetry file'
      END IF
   END IF

   WRITE (*,FMT="(/'No. of MPI Processes = ',I4)") num_pes
   IF(.NOT. regular)THEN
      WRITE (*,FMT="('No. of cores per compute node = ',I4)") cores_per_node
   END IF

   ! Generate the land/sea mask from the bathymetry which is read
   ! from file
   CALL generate_mask(regular, ierr)
   IF(ierr /= 0)THEN
      STOP 'Failed to generate land/sea mask'
   END IF

   CALL allocate_arrays(num_pes, ierr)

   IF(regular)THEN

      ! Imitate standard NEMO regular partitioning
      CALL nemo_partition(num_pes)

   ELSE

      ! Multi-core aware version of recursive k-section partitioning. 
      ! Currently only accounts for whether a grid point is wet or dry.
      ! It has no knowledge of the number of wet levels at a point.
      CALL partition_mca_rk ( imask, num_pes, cores_per_node, ierr )

      ! Check the error code from partitioning.
      IF ( ierr /= 0 ) THEN
         STOP 'nemo_recursive_partition: Partitioning failed'
      ENDIF

   END IF

   CALL write_partition(imask)

   WRITE(*,FMT="(/'Partitioning done and written to partition.dat.new. Also'/ &
        &  'written in postscript to domain_decomp.ps.'/)")

CONTAINS

  SUBROUTINE print_usage()
     IMPLICIT None
     WRITE(*,*)
     WRITE(*,*) 'Usage: make_partition.exe -d <jpiglo>x<jpjglo> -np <nproc> -c <cores_per_node> [-r] [-cyclic] [-h]'
     WRITE(*,*) 'Where:  domain is jpiglo x jpjglo in extent'
     WRITE(*,*) '        nproc is no. of sub-domains/MPI procs required'
     WRITE(*,*) '        cores_per_node is no. of cores per compute node on machine'
     WRITE(*,*) '           on which NEMO is to be executed'
     WRITE(*,*) '        -r      : generate regular (old style) partition'
     WRITE(*,*) '        -cyclic : enable cyclic E/W boundary conditions'
     WRITE(*,*) '        -h       : print this message'
     WRITE(*,*)
     WRITE(*,*) 'Vertical coordinate options are read from the NEMO namelist file and'
     WRITE(*,*) 'used to determine whether to smooth the bathy.'
     WRITE(*,*)
  END SUBROUTINE print_usage


END PROGRAM make_partition