/*=========================================================================

  Program:   Visualization Toolkit
  Module:    vtkMPIUtilities.cxx

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notice for more information.

=========================================================================*/
#include "vtkMPIUtilities.h"

// VTK includes
#include "vtkMPICommunicator.h"
#include "vtkMPIController.h"

// C/C++ includes
#include <cassert>
#include <cstdarg>
#include <cstdio>

namespace vtkMPIUtilities
{

void Printf(vtkMPIController* comm, const char* format, ...)
{
  // Sanity checks
  assert("pre: MPI controller is nullptr!" && (comm != nullptr));
  assert("pre: format argument is nullptr!" && (format != nullptr));

  if (comm->GetLocalProcessId() == 0)
  {
    va_list argptr;
    va_start(argptr, format);
    vprintf(format, argptr);
    fflush(stdout);
    va_end(argptr);
  }

  comm->Barrier();
}

//------------------------------------------------------------------------------
void SynchronizedPrintf(vtkMPIController* comm, const char* format, ...)
{
  // Sanity checks
  assert("pre: MPI controller is nullptr!" && (comm != nullptr));
  assert("pre: format argument is nullptr!" && (format != nullptr));

  int rank = comm->GetLocalProcessId();
  int numRanks = comm->GetNumberOfProcesses();

  vtkMPICommunicator::Request rqst;
  int* nullmsg = nullptr;

  if (rank == 0)
  {
    // STEP 0: print message
    printf("[%d]: ", rank);
    fflush(stdout);

    va_list argptr;
    va_start(argptr, format);
    vprintf(format, argptr);
    fflush(stdout);
    va_end(argptr);

    // STEP 1: signal next process (if any) to print
    if (numRanks > 1)
    {
      comm->NoBlockSend(nullmsg, 0, rank + 1, 0, rqst);
    } // END if
  }   // END first rank
  else if (rank == numRanks - 1)
  {
    // STEP 0: Block until previous process completes
    comm->Receive(nullmsg, 0, rank - 1, 0);

    // STEP 1: print message
    printf("[%d]: ", rank);

    va_list argptr;
    va_start(argptr, format);
    vprintf(format, argptr);
    fflush(stdout);
    va_end(argptr);
  } // END last rank
  else
  {
    // STEP 0: Block until previous process completes
    comm->Receive(nullmsg, 0, rank - 1, 0);

    // STEP 1: print message
    printf("[%d]: ", rank);

    va_list argptr;
    va_start(argptr, format);
    vprintf(format, argptr);
    fflush(stdout);
    va_end(argptr);

    // STEP 2: signal next process to print
    comm->NoBlockSend(nullmsg, 0, rank + 1, 0, rqst);
  }

  comm->Barrier();
}

} // END namespace vtkMPIUtilities