#include "Population.h"
#include "Logger.h"
#include "Parameters.h"
#include "Random.h"
#include "Statistics.h"

#include <cassert>
#include <queue>
#include <string>
#include <boost/format.hpp>

bool Population::ConvergedAllele::Comparator::operator()( const ConvergedAllele & lhs, const ConvergedAllele & rhs ) {
  return ( lhs.probability < rhs.probability );
}

Population::ConvergedAllele::ConvergedAllele(){ /* Empty */ }

Population::ConvergedAllele::ConvergedAllele( unsigned int l, AlleleType a, FloatType p ) :
    locus( l ), allele( a ), probability( p ) { /* Empty */ }

AlleleType Population::getRandomAllele( unsigned int locus, Random *random ) const {
  return random->getInt( 0, Parameters::getInstance()->getA( locus ) - 1 );
}

AlleleType Population::getProportionateAllele( unsigned int locus, Random *random ) const {
  const AlleleType a = Parameters::getInstance()->getA( locus );
  const FloatType total = ( *data )[locus][a];
  FloatType cumulative = static_cast<FloatType>( 0.0 );

  // Pick a random number and compare it to the cumulative probabilities
  const FloatType randVal = random->get01() * total; // Random value between 0 and total
  for ( AlleleType allele = 0; allele < a; allele++ ) {
    cumulative += ( *data )[locus][allele];
    if ( randVal < cumulative )
      return allele;
  }
  // It can happen that randVal is infinitesimally >= total. In that case, just return the max. allele.
  return ( a - 1 );
}

const Population::ConvergedAllele Population::getConvergedAllele( unsigned int locus ) const {
  const AlleleType a = Parameters::getInstance()->getA( locus );
  const FloatType cT = ( Parameters::getInstance()->getConvergenceThreshold() ) *
                       (( *data )[locus][a] );
  assert( convergedAlleles.test( locus ) );
  for ( AlleleType allele = 0; allele < a; allele++ ) {
    FloatType probability = ( *data )[locus][allele];
    if ( probability > cT )
      return ConvergedAllele( locus, allele, probability );
  }
  // If we get here, something is seriously wrong
  assert( false );
  return ConvergedAllele( 0, 0, static_cast<FloatType>( 0.0 ) );
}

void Population::adjustProbabilities( const AlleleType *chromosome, FloatType epsilon ) {
  const FloatType zero = static_cast<FloatType>( 0.0 );
  const unsigned int n = Parameters::getInstance()->getN();
  FloatType *cell, oldValue, newValue;
  for ( unsigned int locus = 0; locus < n; locus++ ) {
    cell = &(( *data )[ locus ][ chromosome[locus] ] );
    oldValue = *cell;
    newValue = oldValue + epsilon;
    if ( newValue < zero ) {
      /* If the new value is strictly less than zero it means that we are administering penalties.
       * This means that the probability-addition in the reward is not matched by a deduction here,
       * and therefore the sum of probabilities increases by epsilon.
       */
      assert( epsilon < zero );
      ( *data )[ locus ][ Parameters::getInstance()->getA( locus )] -= epsilon;
    } else {
      *cell = newValue;
    }
  }
}

Population::Population() :
    data( new MatrixType( boost::extents
                          [Parameters::getInstance()->getN()]
                          [Parameters::getInstance()->getMaxA()+1u] ) ),
    convergedAlleles( Parameters::getInstance()->getN() ) {
  // Local copies of parameters
  const unsigned int n = Parameters::getInstance()->getN();
  AlleleType a;
  const FloatType cT = Parameters::getInstance()->getConvergenceThreshold();
  // Fill the data structures
  for ( unsigned int locus = 0; locus < n; locus++ ) {
    a = Parameters::getInstance()->getA( locus );
    FloatType probability = static_cast<FloatType>( 1.0 ) / static_cast<FloatType>( a );
    for ( AlleleType allele = 0; allele < a; allele++ ) {
      ( *data )[locus][allele] = probability;
    }
    ( *data )[locus][a] = static_cast<FloatType>( 1.0 );
    if ( probability > cT )
      convergedAlleles.set( locus );
  }
}

Population::~Population() {
  delete data;
}

void Population::print() const {
  const unsigned int n = Parameters::getInstance()->getN();
  const unsigned int maxA = Parameters::getInstance()->getMaxA();
  const Logger::LogLevel logLevel = Logger::LOGLEVEL_DEBUG;
  AlleleType a;
  std::string buffer;

  LOG( logLevel, "Virtual Population\t%ux%u matrix", n, maxA );
  for ( unsigned int locus = 0; locus < n; locus++ ) {
    a = Parameters::getInstance()->getA( locus );
    buffer.clear();
    buffer += boost::str( boost::format( "Virtual Population\t%u" ) % locus );
    for ( AlleleType allele = 0; allele < a; allele++ )
      buffer += boost::str( boost::format( "\t%.3f" ) % (( *data )[locus][allele] ) );
    LOG( Logger::LOGLEVEL_DEBUG, buffer.c_str() );
  }
}

void Population::selectIndividual( AlleleType *destination ) const {
  const unsigned int n = Parameters::getInstance()->getN();
  const FloatType rho = Parameters::getInstance()->getRho();
  Random *random = Random::getInstance();
  for ( unsigned int locus = 0; locus < n; locus++ ) {
    if ( random->get01() < rho ) {
      // Do mutation
      destination[locus] = getRandomAllele( locus, random );
    } else {
      // Use probability
      destination[locus] = getProportionateAllele( locus, random );
    }
  }
}

void Population::rewardGenes( const AlleleType *chromosome ) {
  adjustProbabilities( chromosome, Parameters::getInstance()->getEpsilon() );
}

void Population::penalizeGenes( const AlleleType *chromosome ) {
  adjustProbabilities( chromosome, -Parameters::getInstance()->getEpsilon() );
}

void Population::updateConvergedAlleles() {
  const unsigned int n = Parameters::getInstance()->getN();
  const FloatType cT = Parameters::getInstance()->getConvergenceThreshold();
  unsigned int c = 0;
  for ( unsigned int locus = 0; locus < n; locus++ ) {
    const AlleleType a = Parameters::getInstance()->getA( locus );
    const FloatType threshold = cT * (( *data )[ locus ][ a ] );
    bool isConverged = false;
    for ( AlleleType i = 0; ( !isConverged ) && ( i < a ); i++ ) {
      FloatType p = ( *data )[ locus ][ i ];
      if ( p > threshold ) {
        c++;
        isConverged = true;
      }
    }
    convergedAlleles.set( locus, isConverged );
  }
  Statistics::getInstance()->setNumConvergedAlleles( c );
}

void Population::processConvergedAllele( unsigned int locus, AlleleType allele ) {
  const FloatType zero = static_cast<FloatType>( 0.0 );
  const AlleleType a = Parameters::getInstance()->getA( locus );
  const FloatType mu = Parameters::getInstance()->getMu();
  FloatType *cell, oldValue, newValue;
  for ( AlleleType i = 0; i < a; i++ ) {
    // Fetch the old value from the data cell
    cell = &(( *data )[ locus ][ i ] );
    oldValue = *cell;
    // Compute the new value
    if ( i == allele )
      newValue = ( oldValue + mu ) / ( static_cast<FloatType>( 1.0 ) + mu );
    else
      newValue = oldValue / ( static_cast<FloatType>( 1.0 ) + mu );
    // Make sure the new value is within range
    assert( newValue >= zero );
    // Assign the new value
    *cell = newValue;
  }
  LOG( Logger::LOGLEVEL_DEBUG, "processConvergedAllele\tlocus:%u,allele:%u", locus, allele );
}

unsigned int Population::prepareAllelesMessage( MessagePayloadType *buffer ) {
  const unsigned int x = Parameters::getInstance()->getEpochX();
  const unsigned int c = Statistics::getInstance()->getNumConvergedAlleles();

  // Make sure all is SNAFU
  assert( convergedAlleles.count() == c );

  // Create priority queue to hold converged alleles, sorted by probability (greatest first)
  std::priority_queue < ConvergedAllele,
  std::vector<ConvergedAllele>,
  ConvergedAllele::Comparator > queue;

  // Iterate through CA and add all set indices to queue
  for ( boost::dynamic_bitset<>::size_type itCA = convergedAlleles.find_first();
        itCA != convergedAlleles.npos;
        itCA = convergedAlleles.find_next( itCA ) ) {
    queue.push( getConvergedAllele( static_cast<unsigned int>( itCA ) ) );
  }

  // Extract ordered items from the queue, add them to buffer
  unsigned int count = 0;
  unsigned int itBuffer = 0;
  while ( count < x && !queue.empty() ) {
    const ConvergedAllele ca = queue.top();
    buffer[++itBuffer] = static_cast<MessagePayloadType>( ca.locus );
    buffer[++itBuffer] = static_cast<MessagePayloadType>( ca.allele );
    queue.pop();
    count++;
    LOG( Logger::LOGLEVEL_DEBUG, "prepareAllelesMessage\t%03u::locus:%u,allele:%u\tprobability:%.3f",
         count, ca.locus, ca.allele, ca.probability );
  }

// By now we should have added 2*count items to the buffer
  assert( itBuffer == 2u * count );

// Set the buffer's header
  buffer[0] = static_cast<MessagePayloadType>( count );
  return count;
}
