! (C) Crown Copyright 2025- Met Office.
! 
! This software is licensed under the terms of the Apache Licence Version 2.0
! which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
!

PROGRAM TEST_SPLIT_MPI_COMM

USE PARKIND1,        ONLY : JPRB, JPIM
USE MPL_MODULE,      ONLY : MPL_INIT, MPL_END, MPL_MYRANK, MPL_NPROC, MPL_COMM_SPLIT, MPL_SETDFLT_COMM, MPL_COMM
USE ABORT_TRANS_MOD, ONLY : ABORT_TRANS

USE MPL_MPIF

IMPLICIT NONE

#include "setup_trans0.h"
#include "setup_trans.h"
#include "inv_trans.h"
#include "dir_trans.h"
#include "trans_inq.h"
#include "gath_grid.h"

INTEGER(KIND=JPIM), PARAMETER, DIMENSION(2) :: TRUNCATIONS = [79, 188]

! MODE
INTEGER(KIND=JPIM), PARAMETER, DIMENSION(2) :: MS = [1, 3]
INTEGER(KIND=JPIM), PARAMETER, DIMENSION(2) :: NS = [2, 4]


INTEGER(KIND=JPIM) :: NUM_SPECTRAL_ELEMENTS, NUM_GRID_POINTS
INTEGER(KIND=JPIM) :: G_NUM_SPECTRAL_ELEMENTS, G_NUM_GRID_POINTS  ! GLOBAL
INTEGER(KIND=JPIM) :: LOCAL_SPECTRAL_COEFFICIENT_INDEX
INTEGER(KIND=JPIM) :: IERROR
INTEGER(KIND=JPIM) :: I
INTEGER(KIND=JPIM) :: SPLIT_NUM_RANKS, SPLIT_RANK
INTEGER(KIND=JPIM) :: WORLD_NUM_RANKS, WORLD_RANK
INTEGER(KIND=JPIM) :: SPLIT_COLOUR, SPLIT_KEY
INTEGER(KIND=JPIM) :: SPLIT_COMM, DUMMY_COMM
INTEGER(KIND=JPIM) :: TRUNCATION
INTEGER(KIND=JPIM) :: M, N
INTEGER(KIND=JPIM) :: NUM_LATITUDES, NUM_LONGITUDES

INTEGER(KIND=JPIM), ALLOCATABLE :: SPECTRAL_INDICES(:)

! Fields
REAL(KIND=JPRB), ALLOCATABLE :: SPECTRAL_FIELD(:,:)
REAL(KIND=JPRB), ALLOCATABLE :: GRID_POINT_FIELD(:,:,:)

! NOTE: 1 dimensional global field used to write to file output.
REAL(KIND=JPRB), ALLOCATABLE :: G_GRID_POINT_FIELD(:,:)

CHARACTER(LEN=1024) :: FILENAME

CALL MPL_INIT()
WORLD_NUM_RANKS = MPL_NPROC()
WORLD_RANK = MPL_MYRANK()

SPLIT_COLOUR = GET_SPLIT_GROUP(WORLD_RANK, WORLD_NUM_RANKS)
SPLIT_KEY = WORLD_RANK
CALL MPL_COMM_SPLIT(MPL_COMM, SPLIT_COLOUR, SPLIT_KEY, SPLIT_COMM, IERROR)
CALL MPL_SETDFLT_COMM(SPLIT_COMM, DUMMY_COMM)

SPLIT_RANK = MPL_MYRANK()
SPLIT_NUM_RANKS = MPL_NPROC()

PRINT*,"=== RANK ", WORLD_RANK, ", SETUP ON GROUP", SPLIT_COLOUR, "NUM RANKS = ", SPLIT_NUM_RANKS, "==="

! Assert that the split comm is smaller than world.
IF (SPLIT_NUM_RANKS >= WORLD_NUM_RANKS) THEN
  PRINT*, "SPLIT = ", SPLIT_NUM_RANKS, "TOTAL = ", WORLD_NUM_RANKS
  CALL ABORT_TRANS("ERROR: SPLIT COMMUNICATOR NOT SMALLER THAN MPI_COMM_WORLD.")
END IF

PRINT*,"=== LOCAL RANK ", SPLIT_RANK, ", ON GROUP", SPLIT_COLOUR, "SIZE", SPLIT_NUM_RANKS, "==="

CALL SETUP_TRANS0(KPRINTLEV=0, LDMPOFF=.FALSE., &
!                 SPLIT GRID NS            SPLIT SPECTRAL
                  KPRGPNS=SPLIT_NUM_RANKS, KPRTRW=SPLIT_NUM_RANKS)

! Different transform based on the colour.
TRUNCATION = TRUNCATIONS(SPLIT_COLOUR + 1)

NUM_LATITUDES = 2*(TRUNCATION + 1)
NUM_LONGITUDES = NUM_LATITUDES*2

IF (SPLIT_RANK == 1) PRINT*,"COLOUR ", SPLIT_COLOUR, &
  " GLOBAL NUM LON =>", NUM_LONGITUDES, "| LAT =>", NUM_LATITUDES

CALL SETUP_TRANS(KSMAX=TRUNCATION, KDGL=NUM_LATITUDES)

! Get function space sizes
CALL TRANS_INQ(KSPEC2=NUM_SPECTRAL_ELEMENTS, KGPTOT=NUM_GRID_POINTS)
CALL TRANS_INQ(KSPEC2G=G_NUM_SPECTRAL_ELEMENTS, KGPTOTG=G_NUM_GRID_POINTS)

ALLOCATE(SPECTRAL_FIELD(1, NUM_SPECTRAL_ELEMENTS))
ALLOCATE(GRID_POINT_FIELD(NUM_GRID_POINTS, 1, 1))

! Get spectral indices
ALLOCATE(SPECTRAL_INDICES(0:TRUNCATION))
CALL TRANS_INQ(KASM0=SPECTRAL_INDICES)

! Select mode
M = MS(SPLIT_COLOUR + 1)
N = NS(SPLIT_COLOUR + 1)
LOCAL_SPECTRAL_COEFFICIENT_INDEX = SPECTRAL_INDICES(M) + 2*(N - M) + 1

SPECTRAL_FIELD(:,:) = 0.0

IF (LOCAL_SPECTRAL_COEFFICIENT_INDEX > 0) THEN
  SPECTRAL_FIELD(1,LOCAL_SPECTRAL_COEFFICIENT_INDEX) = 1.0
END IF

CALL INV_TRANS(PSPSCALAR=SPECTRAL_FIELD, PGP=GRID_POINT_FIELD)

! -------------- Gather the result on the root (0) and write to file -----------------

IF (SPLIT_RANK == 1) THEN
  ! Allocate a global field.
  ALLOCATE(G_GRID_POINT_FIELD(G_NUM_GRID_POINTS,1))

  CALL GATH_GRID(PGPG=G_GRID_POINT_FIELD, KFGATHG=1, KTO=[1], PGP=GRID_POINT_FIELD)

  ! Write to file. can then be plotted using a python script
  ! Such as in the docs: https://sites.ecmwf.int/docs/ectrans/page/usage.html.
  WRITE(FILENAME, "(A,I0,A4)") "grid_point_field_trunc_", TRUNCATION, ".dat"
  OPEN(7, FILE=FILENAME, FORM="unformatted")
  WRITE(7) G_GRID_POINT_FIELD(:,1)
  CLOSE(7)

  PRINT*,"Colour", SPLIT_COLOUR, "finished and written to file: "//TRIM(FILENAME)
ELSE
  CALL GATH_GRID(KFGATHG=1, KTO=[1], PGP=GRID_POINT_FIELD)
END IF

CALL MPL_END()

CONTAINS

! Get the colour of comm for this rank.
FUNCTION GET_SPLIT_GROUP(RANK, WORLD_SIZE) RESULT(GROUP)
  IMPLICIT NONE

  INTEGER(KIND=JPIM), INTENT(IN) :: RANK
  INTEGER(KIND=JPIM), INTENT(IN) :: WORLD_SIZE
  ! RETURN
  INTEGER(KIND=JPIM) :: GROUP

  REAL(KIND=JPRB) :: RANK_RATIO

  ! ----------------------------------------------
  ! Uneven splitting based on a ratio 1:3.
  ! ----------------------------------------------
  RANK_RATIO = REAL(RANK, KIND=JPRB) / REAL(WORLD_SIZE, KIND=JPRB)

  ! Split X%
  IF (RANK_RATIO <= 0.25_JPRB) THEN
    GROUP = 0
  ELSE
    GROUP = 1
  END IF

END FUNCTION GET_SPLIT_GROUP

END PROGRAM
