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

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

const FloatType Population::limitProbabilityMin = static_cast<FloatType>( 0.0 );
const FloatType Population::limitProbabilityMax = static_cast<FloatType>( 1.0 );

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 );
  AlleleType allele;
  FloatType randVal;
  FloatType total = static_cast<FloatType>( 0.0 );
  std::vector<FloatType> cumulative( a );

  // Build a vector of cumulative probabilities for the alleles at this locus
  // TODO: Maybe cache these arrays...test the app with a profiler to see if that's necessary.
  for ( allele = 0; allele < a; allele++ ) {
    total += ( *data )[locus][allele];
    cumulative[allele] = total;
  }

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

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

void Population::adjustProbabilities( const AlleleType * chromosome, FloatType epsilon ) {
  const unsigned int n = Parameters::getInstance()->getN();
  const FloatType cT = Parameters::getInstance()->getConvergenceThreshold();
  for ( unsigned int locus = 0; locus < n; locus++ ) {
    FloatType *cell = &(( *data )[ locus ][ chromosome[locus] ] );
    FloatType oldValue = *cell, newValue = oldValue + epsilon;
    if ( newValue < limitProbabilityMax && newValue > limitProbabilityMin ) {
      *cell = newValue; // Don't update probability if new value is out of range
      if ( oldValue <= cT && newValue > cT )
        convergeAllele( locus );
      else if ( oldValue > cT && newValue <= cT )
        divergeAllele( locus );
    }
  }
}

void Population::convergeAllele( unsigned int locus ) {
  assert( convergedAlleles.test( locus ) == false );
  const unsigned int x = Parameters::getInstance()->getEpochX();
  const unsigned int c = Statistics::getInstance()->adjustNumConvergedAlleles( 1 );
  convergedAlleles.set( locus );
  if ( c > x )
    recentConvergedAlleles.push_front( locus );
  resizeRecentConvergedAlleles( c, x );

  LOG( Logger::LOGLEVEL_DEBUG, "convergeAllele(%u)\tc:%u,x:%u\tca.count:%u,rca.size:%u",
       locus, c, x,
       convergedAlleles.count(),
       recentConvergedAlleles.size() );
}

void Population::divergeAllele( unsigned int locus ) {
  assert( convergedAlleles.test( locus ) == true );
  const unsigned int x = Parameters::getInstance()->getEpochX();
  const unsigned int c = Statistics::getInstance()->adjustNumConvergedAlleles( -1 );
  convergedAlleles.reset( locus );
  // If the list is nonempty, remove one element: remove 'locus' if we can, otherwise the oldest element.
  if ( !recentConvergedAlleles.empty() ) {
    bool foundIt = false;
    std::list<unsigned int>::iterator it = recentConvergedAlleles.begin();
    for ( ; !foundIt && it != recentConvergedAlleles.end(); it++ ) {
      if ( *it == locus ) {
        it = recentConvergedAlleles.erase( it );
        foundIt = true;
      }
    }
    if ( !foundIt )
      recentConvergedAlleles.pop_back();
  }
  resizeRecentConvergedAlleles( c, x );

  LOG( Logger::LOGLEVEL_DEBUG, "divergeAllele(%u)\tc:%u,x:%u\tca.count:%u,rca.size:%u",
       locus, c, x,
       convergedAlleles.count(),
       recentConvergedAlleles.size() );
}

void Population::resizeRecentConvergedAlleles( unsigned int c, unsigned int x ) {
  recentConvergedAlleles.resize( c - std::min( c, x ) );
}

Population::Population() :
    data( new MatrixType( boost::extents
                          [Parameters::getInstance()->getN()]
                          [Parameters::getInstance()->getMaxA()] ) ),
    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;
    }
    if ( probability > cT )
      convergeAllele( locus );
  }
}

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

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::processConvergedAllele( unsigned int locus, AlleleType allele ) {
  const AlleleType a = Parameters::getInstance()->getA( locus );
  const FloatType cT = Parameters::getInstance()->getConvergenceThreshold();
  const FloatType mu = Parameters::getInstance()->getMu();
  const bool oldConverged = convergedAlleles.test( locus );
  bool newConverged = false;
  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
    newValue = std::max( limitProbabilityMin, newValue );
    newValue = std::min( limitProbabilityMax, newValue );
    // Check for convergence
    if ( newValue > cT )
      newConverged = true;
    // Assign the new value
    *cell = newValue;
  }
  // Update this gene's convergence status if needed
  if ( oldConverged && !newConverged )
    divergeAllele( locus );
  else if ( !oldConverged && newConverged )
    convergeAllele( locus );

  LOG( Logger::LOGLEVEL_DEBUG, "processConvergedAllele\tlocus:%u,allele:%u\toldConverged:%c,newConverged:%c",
       locus, allele, oldConverged ? 'T' : 'F', newConverged ? 'T' : 'F' );
}

unsigned int Population::prepareAllelesMessage( MessagePayloadType *buffer ) {
  const unsigned int y = Parameters::getInstance()->getEpochY();
  const unsigned int x = Parameters::getInstance()->getEpochX();
  const unsigned int c = Statistics::getInstance()->getNumConvergedAlleles();
  LOG( Logger::LOGLEVEL_DEBUG, "prepareAllelesMessage\ty:%u,x:%u,c:%u\tca.count:%u,rca.size:%u",
       y, x, c, convergedAlleles.count(), recentConvergedAlleles.size() );
  std::list<unsigned int>::iterator itRCA;
  boost::dynamic_bitset<>::size_type itCA;
  unsigned int itBuffer = 0;
  MessagePayloadType locus, allele;
  // Make sure the programmer did not make any mistakes...
  assert( convergedAlleles.count() == c );
  assert( recentConvergedAlleles.size() >= c - std::min( c, x ) );
  resizeRecentConvergedAlleles( c, x );
  // Remove c-x (or less) items from CA, leaving CA with exactly min(c,x) items
  for ( itRCA = recentConvergedAlleles.begin();
        itRCA != recentConvergedAlleles.end();
        itRCA++ )
    convergedAlleles.reset( *itRCA );
  assert( convergedAlleles.count() == std::min( c, x ) );
  // Iterate through CA and add all set indices to buffer
  for ( itCA = convergedAlleles.find_first();
        itCA != convergedAlleles.npos;
        itCA = convergedAlleles.find_next( itCA ) ) {
    locus = static_cast<MessagePayloadType>( itCA );
    allele = static_cast<MessagePayloadType>( getConvergedAllele( static_cast<unsigned int>( itCA ) ) );
    buffer[++itBuffer] = locus;
    buffer[++itBuffer] = allele;
  }
  // Put the removed items back in CA
  for ( itRCA = recentConvergedAlleles.begin();
        itRCA != recentConvergedAlleles.end();
        itRCA++ )
    convergedAlleles.set( *itRCA );
  // CA should be back where it was
  assert( convergedAlleles.count() == c );
  // By now we should have added 2*min(c,x) items to the buffer
  assert( itBuffer == 2 * std::min( c, x ) );
  // Set the buffer's header
  buffer[0] = static_cast<MessagePayloadType>( std::min( c, x ) );
  return std::min( c, x );
}

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() );
  }
}
