! See copyright notice in the COPYRIGHT file.
!mpi!nprocs = 5
program tem_sparta_test
  use mpi
  use env_module,               only: rk, init_env, fin_env, long_k
  use tem_utestEnv_module,  only: cubeconf
  use tem_sparta_module, only: tem_sparta_type, tem_balance_sparta, &
    &                          tem_init_sparta, tem_destroy_sparta, &
    &                          tem_output_sparta
  use tem_logging_module,    only: logUnit, tem_logging_load_primary
  use tem_general_module,    only: tem_general_type, tem_start

  use aotus_module,          only: open_config_chunk
  use flu_binding,           only: flu_state

  implicit none

  ! MPI variables
  integer         :: iError
  integer         :: myrank
  integer         :: nprocs
  integer         :: comm
  type(flu_state) :: conf
  type(tem_general_type) :: general

  logical :: OK = .false.
  logical :: correct = .false.

  ! main variables
  real(kind=rk), allocatable :: weight(:)
  integer :: myElems
  integer(kind=long_k) :: offset
  type( tem_sparta_type ) :: sparta

  call tem_start('TREELM unit test', 'utest', general)
  comm = general%proc%comm
  myrank = general%proc%rank
  nprocs = general%proc%comm_size
  if ( nprocs /= 5 ) stop

  ! Open the configuration file 
  call open_config_chunk(L = conf, chunk = trim(cubeconf))
  ! load and initialize logUnit
  call tem_logging_load_primary(conf = conf,  &
    &                           rank = myrank )

  allocate( weight(5) )
  myElems = 5
  select case ( myrank )
    case (0)
      weight = [ 5.0, 3.0, 1.0, 2.0, 1.0 ]
    case (1)
      weight = [ 4.0, 6.0, 1.0, 3.0, 2.0 ]
    case (2)
      weight = [ 1.0, 3.0, 1.0, 1.0, 1.0 ]
    case (3)
      weight = [ 1.0, 9.0, 1.0, 1.0, 1.0 ]
    case (4)
      weight = [ 1.0, 1.0, 1.0, 1.0, 1.0 ]
    case default
      stop
  end select

  ! write(*,"(A,I0,A,F7.1)") "Before balance, rank: ", myrank, ", my workload: ", sum(weight)
  call tem_init_sparta( sparta, nprocs )
  call tem_balance_sparta(weight, myrank, nprocs, comm, myElems, offset, &
    &                    sparta )
  call tem_output_sparta( sparta, logUnit(1) )
  write(*,"(3(A,I2),A,F5.1)") "After balance, rank: ", myrank, &
    &                      ", myElems: ", myElems, &
    &                      ", offset: ", offset, &
    &                      ", my workload: ", sum(weight)

  if (myrank == 0 .and. myElems == 4 .and. offset == 0) then
    OK = .true.
  else if (myrank == 1 .and. myElems == 3 .and. offset == 4) then
    OK = .true.
  else if (myrank == 2 .and. myElems == 5 .and. offset == 7) then
    OK = .true.
  else if (myrank == 3 .and. myElems == 5 .and. offset == 12) then
    OK = .true.
  else if (myrank == 4 .and. myElems == 8 .and. offset == 17) then
    OK = .true.
  else
    OK = .false.
  end if

  call mpi_reduce( OK, correct, 1, mpi_logical, mpi_land, 0, comm, iError )

  call tem_destroy_sparta( sparta )

  call fin_env()

  if ( myrank == 0 ) then
    if ( .not. correct ) then
      write(*,"(3(A,I0),A,F7.1)") "After  balance, rank: ", myrank, &
        &                      ", myElems: ", myElems, &
        &                      ", offset: ", offset, &
        &                      ", my workload: ", sum(weight)
      write(*,*) "FAILED"
    else
      write(*,*) "PASSED"
    end if
  end if

  deallocate( weight )

end program tem_sparta_test
