#include "Communication.h"
#include "GlobalStatistics.h"
#include "Logger.h"
#include "MessageStatistics.h"
#include "Parameters.h"
#include "Statistics.h"
#include <cassert>
#include <limits>
#include <algorithm>

//#define VERBOSE

Communication::Communication( SgEngine *e ) :
    engine( e ),
    bufferSize( std::max( 5u, Parameters::getInstance()->getN()*2u + 1u ) ),
    bufferIn( new MessagePayloadType[bufferSize] ),
    bufferOut( new MessagePayloadType[bufferSize] ),
    localSolutionsReceived( 1 ) { /* Empty. */ }

Communication::~ Communication() {
  delete [] bufferIn;
  delete [] bufferOut;
}

int Communication::getNeighbor( int distance ) const {
  assert( distance > 0 );
  const int rank = Parameters::getInstance()->getRank();
  const int p = Parameters::getInstance()->getP();
  int offset, neighbor;
  if ( distance % 2 == 0 )
    offset = -distance / 2;
  else
    offset = ( distance + 1 ) / 2;
  neighbor = ( rank + offset + p ) % p;
#ifdef VERBOSE
  LOG( Logger::LOGLEVEL_DEBUG, "getNeighbor\trank:%d\tdistance:%d\toffset:%d\tneighbor:%d",
       rank, distance, offset, neighbor );
#endif
  return neighbor;
}

bool Communication::isMessageWaiting( MPI_Status *mpiStatus ) {
  int flag = 0;
  MPI_Iprobe( MPI_ANY_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, &flag, mpiStatus );
  return ( flag != 0 );
}

template<>
void Communication::processMessage<Messages::ALLELES>( int source )  {
  unsigned int howMany = static_cast<unsigned int>( bufferIn[0] );
  unsigned int locus;
  AlleleType allele;
  for ( unsigned int i = 0; i < howMany; i++ ) {
    locus = static_cast<unsigned int>( bufferIn[2u*i+1u] );
    allele = static_cast<AlleleType>( bufferIn[2u*i+2u] );
    engine->getPopulation()->processConvergedAllele( locus, allele );
  }
  LOG( Logger::LOGLEVEL_MEDIUM, "processMessage<ALLELES>\tsource:%d\tx:%u", source, howMany );
}

template<>
void Communication::processMessage<Messages::LOCAL_STOP>( int source ) {
  const unsigned int maxValue = static_cast<unsigned int>( std::numeric_limits<MessagePayloadType>::max() );
  GlobalStatistics *gs = GlobalStatistics::getInstance();
  unsigned int convergedAlleles = static_cast<unsigned int>( bufferIn[0] );
  unsigned int epochs = static_cast<unsigned int>( bufferIn[1] ) * maxValue +
                        static_cast<unsigned int>( bufferIn[2] );
  unsigned int generations = static_cast<unsigned int>( bufferIn[3] ) * maxValue +
                             static_cast<unsigned int>( bufferIn[4] );
  gs->setDemeStopped( source );
  gs->setNumConvergedAlleles( source, convergedAlleles );
  gs->setNumEpochs( source, epochs );
  gs->setNumGenerations( source, generations );
  LOG( Logger::LOGLEVEL_MEDIUM,
       "processMessage<LOCAL_STOP>\tsource:%d\tconvergedAlleles:%u\tepochs:%u\tgenerations:%u",
       source, convergedAlleles, epochs, generations );
}

template<>
void Communication::processMessage<Messages::GLOBAL_STOP>( int source ) {
  assert( Parameters::getInstance()->getRank() != 0 );
  // sendMessage( LOCAL_SOLUTION ); // Moved to solveEnd()
  Statistics::getInstance()->setExecutionFinished();
  LOG( Logger::LOGLEVEL_MEDIUM, "processMessage<GLOBAL_STOP>\tsource:%d", source );
}

template<>
void Communication::processMessage<Messages::LOCAL_SOLUTION>( int source ) {
  assert( Parameters::getInstance()->getRank() == 0 );
  if ( bufferIn[0] != std::numeric_limits<MessagePayloadType>::max() ) {
    const unsigned int n = Parameters::getInstance()->getN();
    AlleleType *storage = engine->getSpaceForSolutionCandidate();
    for ( unsigned int i = 0; i < n; i++ )
      storage[i] = static_cast<AlleleType>( bufferIn[i] );
    if ( engine->testSolutionCandidate() )
      LOG( Logger::LOGLEVEL_MEDIUM,
           "processMessage<LOCAL_SOLUTION>\tsource:%d\tnew best individual\tfitness:%.3f",
           source, engine->getBestFitness() );
    else
      LOG( Logger::LOGLEVEL_MEDIUM,
           "processMessage<LOCAL_SOLUTION>\tsource:%d\tinferior solution", source );
  } else {
    LOG( Logger::LOGLEVEL_MEDIUM, "processMessage<LOCAL_SOLUTION>\tsource:%d\tno solution", source );
  }
  if ( ++localSolutionsReceived == Parameters::getInstance()->getP() )
    Statistics::getInstance()->setExecutionFinished();
  LOG( Logger::LOGLEVEL_HIGH,
       "processMessage<LOCAL_SOLUTION>\tlocalSolutionsReceived:%d\tGSS:%c,EF:%c",
       localSolutionsReceived,
       (( Statistics::getInstance()->isGlobalStopSent() ) ? 'T' : 'F' ),
       (( Statistics::getInstance()->isExecutionFinished() ) ? 'T' : 'F' ) );
}

void Communication::checkMessages() {
  MPI_Status status;
  while ( isMessageWaiting( &status ) ) {
    MPI_Recv( bufferIn, bufferSize, MPI_UNSIGNED_SHORT, MPI_ANY_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, &status );
    MessageStatistics *ms = MessageStatistics::getInstance();
    switch ( status.MPI_TAG ) {
      case ALLELES:
        processMessage<ALLELES>( status.MPI_SOURCE );
        ms->increment( ALLELES, RCVD );
        break;
      case LOCAL_STOP:
        processMessage<LOCAL_STOP>( status.MPI_SOURCE );
        ms->increment( LOCAL_STOP, RCVD );
        break;
      case GLOBAL_STOP:
        processMessage<GLOBAL_STOP>( status.MPI_SOURCE );
        ms->increment( GLOBAL_STOP, RCVD );
        break;
      case LOCAL_SOLUTION:
        processMessage<LOCAL_SOLUTION>( status.MPI_SOURCE );
        ms->increment( LOCAL_SOLUTION, RCVD );
        break;
      default:
        LOG( Logger::LOGLEVEL_LOW, "Incoming message\tunknown tag\t%d", status.MPI_TAG );
        assert( false );
    }
  }
}

template<>
void Communication::sendMessage<Messages::ALLELES>() {
  const int delta = Parameters::getInstance()->getDelta();
  if ( delta > 0 ) {
    const unsigned int x = engine->getPopulation()->prepareAllelesMessage( bufferOut );
    for ( int distance = 1; distance <= delta; distance++ )
      MPI_Send( bufferOut, ( x * 2u + 1u ), MPI_UNSIGNED_SHORT, getNeighbor( distance ), ALLELES, MPI_COMM_WORLD );
    LOG( Logger::LOGLEVEL_HIGH, "sendMessage<ALLELES>\tx:%u", x );
  }
}

template<>
void Communication::sendMessage<Messages::LOCAL_STOP>() {
  const unsigned int maxValue = static_cast<unsigned int>( std::numeric_limits<MessagePayloadType>::max() );
  Statistics *statistics = Statistics::getInstance();
  const unsigned int convergedAlleles = statistics->getNumConvergedAlleles();
  const unsigned int epochs = statistics->getNumEpochs();
  const unsigned int generations = statistics->getNumGenerationsTotal();
  MessagePayloadType *buffer = ( Parameters::getInstance()->getRank() == 0 ) ? bufferIn : bufferOut;
  buffer[0] = static_cast<MessagePayloadType>( convergedAlleles );
  buffer[1] = static_cast<MessagePayloadType>( epochs / maxValue );
  buffer[2] = static_cast<MessagePayloadType>( epochs % maxValue );
  buffer[3] = static_cast<MessagePayloadType>( generations / maxValue );
  buffer[4] = static_cast<MessagePayloadType>( generations % maxValue );
  if ( Parameters::getInstance()->getRank() == 0 )
    processMessage<LOCAL_STOP>( 0 );
  else
    MPI_Send( buffer, 5, MPI_UNSIGNED_SHORT, 0, LOCAL_STOP, MPI_COMM_WORLD );
  LOG( Logger::LOGLEVEL_HIGH, "sendMessage<LOCAL_STOP>\tconvergedAlleles:%u\tepochs:%u\tgenerations:%u",
       convergedAlleles, epochs, generations );
}

template<>
void Communication::sendMessage<Messages::GLOBAL_STOP>() {
  assert( Parameters::getInstance()->getRank() == 0 );
  const int p = Parameters::getInstance()->getP();
  Statistics::getInstance()->setGlobalStopSent();
  if ( p == 1 )
    Statistics::getInstance()->setExecutionFinished();
  else
    for ( int neighbor = 1; neighbor < p; neighbor++ )
      MPI_Send( bufferOut, 0, MPI_UNSIGNED_SHORT, neighbor, GLOBAL_STOP, MPI_COMM_WORLD );
  LOG( Logger::LOGLEVEL_HIGH, "sendMessage<GLOBAL_STOP>\tCPU:%u",
       Statistics::getInstance()->getCpuTime() );
}

template<>
void Communication::sendMessage<Messages::LOCAL_SOLUTION>() {
  assert( Parameters::getInstance()->getRank() != 0 );
  const unsigned int n = Parameters::getInstance()->getN();
  unsigned int length;
  if ( engine->hasBestIndividual() ) {
    const AlleleType *solution = engine->getBestIndividual();
    for ( unsigned int i = 0; i < n; i++ )
      bufferOut[i] = static_cast<MessagePayloadType>( solution[i] );
    length = n;
  } else {
    bufferOut[0] = std::numeric_limits<MessagePayloadType>::max();
    length = 1;
  }
  MPI_Send( bufferOut, length, MPI_UNSIGNED_SHORT, 0, LOCAL_SOLUTION, MPI_COMM_WORLD );
  Statistics::getInstance()->setExecutionFinished();
  LOG( Logger::LOGLEVEL_HIGH, "sendMessage<LOCAL_SOLUTION>\tlength:%u", length );
}

template<>
void Communication::sendMessage<Messages::TERMINATION_TOKEN>() {
  assert( Parameters::getInstance()->getRank() == 0 );
  const int p = Parameters::getInstance()->getP();
  for ( int neighbor = 1; neighbor < p; neighbor++ )
    MPI_Send( bufferOut, 0, MPI_UNSIGNED_SHORT, neighbor, TERMINATION_TOKEN, MPI_COMM_WORLD );
  LOG( Logger::LOGLEVEL_HIGH, "sendMessage<TERMINATION_TOKEN>" );
}

void Communication::sendMessage( MessageType messageType ) {
  MessageStatistics *ms = MessageStatistics::getInstance();
  switch ( messageType ) {
    case ALLELES:
      sendMessage<ALLELES>();
      ms->increment( ALLELES, SENT );
      break;
    case LOCAL_STOP:
      sendMessage<LOCAL_STOP>();
      ms->increment( LOCAL_STOP, SENT );
      break;
    case GLOBAL_STOP:
      sendMessage<GLOBAL_STOP>();
      ms->increment( GLOBAL_STOP, SENT );
      break;
    case LOCAL_SOLUTION:
      sendMessage<LOCAL_SOLUTION>();
      ms->increment( LOCAL_SOLUTION, SENT );
      break;
    case TERMINATION_TOKEN:
      sendMessage<TERMINATION_TOKEN>();
      ms->increment( TERMINATION_TOKEN, SENT );
      break;
    default:
      assert( false );
  }
}

void Communication::solveBegin() {
  MPI_Barrier( MPI_COMM_WORLD );
  Statistics::getInstance()->startTimer();
}

void Communication::solveEnd() {
  LOG( Logger::LOGLEVEL_HIGH, "Communication::solveEnd\tCPU:%u\tAt start",
       Statistics::getInstance()->getCpuTime() );
  if ( Parameters::getInstance()->getRank() == 0 ) {
    sendMessage( TERMINATION_TOKEN );
    LOG( Logger::LOGLEVEL_HIGH, "Communication::solveEnd\tCPU:%u\tBefore barrier",
         Statistics::getInstance()->getCpuTime() );
  } else {
    sendMessage( LOCAL_SOLUTION );
    MPI_Status status;
    bool proceedToBarrier = false;
    unsigned int numMessages = 0;
    while ( !proceedToBarrier ) {
      MPI_Recv( bufferIn, bufferSize, MPI_UNSIGNED_SHORT, MPI_ANY_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, &status );
      if ( status.MPI_TAG == TERMINATION_TOKEN ) {
        proceedToBarrier = true;
          MessageStatistics::getInstance()->increment( TERMINATION_TOKEN, RCVD );
      }
      numMessages++;
    }
    LOG( Logger::LOGLEVEL_HIGH, "Communication::solveEnd\tCPU:%u\tBefore barrier\t%u messages received",
         Statistics::getInstance()->getCpuTime(), numMessages );
  }
  MPI_Barrier( MPI_COMM_WORLD );
  LOG( Logger::LOGLEVEL_HIGH, "Communication::solveEnd\tCPU:%u\tAfter barrier",
       Statistics::getInstance()->getCpuTime() );
}

void Communication::programStart( int *argc, char ***argv, int *myRank, int *myP ) {
  MPI_Init( argc, argv );
  MPI_Comm_rank( MPI_COMM_WORLD, myRank );
  MPI_Comm_size( MPI_COMM_WORLD, myP );
}

void Communication::programTerminate() {
  MPI_Finalize();
}
