#include "genetic/explicitsimulation.h"
#include "genetic/operator.h"
#include "genetic/fitnessfunction.h"
#include "genetic/phenotype.h"

#include <bu/random.h>
#include <bu/sio.h>
#include <bu/mutexlocker.h>

#include <stdlib.h>

using namespace Bu;

Genetic::ExplicitSimulation::ExplicitSimulation( Genetic::Operator *pOper,
		Genetic::FitnessFunction *pFunc, int iThreads, int iPopSize,
		float fKeep, float fRandom, bool bKeepBest, bool bRecalcSurvivors ) :
	pOper( pOper ),
	iPopSize( iPopSize ),
	fKeep( fKeep ),
	fRandom( fRandom ),
	bKeepBest( bKeepBest ),
	bRecalcSurvivors( bRecalcSurvivors ),
	bRunning( true )
{
	for( int j = 0; j < iPopSize; j++ )
	{
		xPop.addPhenotype( pOper->random() );
	}

	if( iThreads < 1 )
		iThreads = 1;

	int iId = 0;
	lProcessor.append( new Processor( *this, pFunc, qWork, iId++ ) );
	while( lProcessor.getSize() < iThreads )
		lProcessor.append(
			new Processor( *this, pFunc->clone(), qWork, iId++ )
			);
	for( ProcessorList::iterator i = lProcessor.begin(); i; i++ )
		(*i)->start();

	updateFitness();
}

Genetic::ExplicitSimulation::~ExplicitSimulation()
{
	delete pOper;

	setRunning( false );

	qWork.unblockAll();
	for( ProcessorList::iterator i = lProcessor.begin(); i; i++ )
	{
		(*i)->join();
		delete *i;
	}
}

void Genetic::ExplicitSimulation::timestep()
{
	PhenotypeList lNew;
	
	int iChildren = iPopSize*(1.0-fKeep-fRandom);

	// Create children
	for( int j = 0; j < iChildren; j++ )
	{
		PhenotypeList lParents;
		for( int k = 0; k < pOper->parentCount(); k++ )
			lParents.append( xPop.getPhenotype( selectWeighted() ) );
		lNew.append( pOper->mate( lParents ) );
	}

	// Select phenotypes for keeping
	int iKeep = iPopSize*fKeep;
	FitnessHash hTempFitness;
	for( int j = 0; j < iKeep; j++ )
	{
		Genetic::PhenotypeId id = selectWeighted();
		mFitness.lock();
		lNew.append( xPop.takePhenotype( id ) );
		hTempFitness.insert( id, hFitness.get( id ) );
		hFitness.erase( id );
		dTotalFitness -= hTempFitness.get( id );
		mFitness.unlock();
	}

	mFitness.lock();
	if( bKeepBest && hFitness.has( uMaxFitness ) )
	{
		lNew.append( xPop.takePhenotype( uMaxFitness ) );
		hTempFitness.insert( uMaxFitness, hFitness.get( uMaxFitness ) );
		hFitness.erase( uMaxFitness );
		dTotalFitness -= hTempFitness.get( uMaxFitness );
	}
	mFitness.unlock();

	// Fill in the remainder with random phenotypes
	while( lNew.getSize() < iPopSize )
	{
		lNew.append( pOper->random() );
	}

	// Refill the population
	mFitness.lock();
	if( bRecalcSurvivors )
		hFitness.clear();
	else
		hFitness = hTempFitness;
	mFitness.unlock();
	xPop.clear();
	xPop.timestep();
	for( PhenotypeList::iterator i = lNew.begin(); i; i++ )
	{
		xPop.addPhenotype( *i );
	}

	updateFitness();
}

Genetic::PhenotypeId Genetic::ExplicitSimulation::selectWeighted()
{
	Bu::MutexLocker ml( mFitness );

	double dSel = Bu::Random::randNorm()*dTotalFitness;
	double dRun = 0.0;

	for( FitnessHash::iterator i = hFitness.begin(); i; i++ )
	{
		dRun += *i;
		if( dSel < dRun )
			return i.getKey();
	}
/*
	sio << "Genetic::ExplicitSimulation::selectWeighted() - failed, picked random"
		<< sio.nl;
	sio << "  " << dMinFitness << " - " << dMaxFitness << " -- " << dTotalFitness << " > " << dRun << " > " << dSel << sio.nl;
//	abort();
*/
    int iSel = Bu::Random::rand( hFitness.getSize() );
    for( FitnessHash::iterator i = hFitness.begin(); i; i++ )
    {
        if( --iSel <= 0 )
            return i.getKey();
    }

	return uMaxFitness;
}

Genetic::PhenotypeId Genetic::ExplicitSimulation::getMaxFitnessId() const
{
    return uMaxFitness;
}

Genetic::Phenotype *Genetic::ExplicitSimulation::getMaxFitnessPhenotype() const
{
    return xPop.getPhenotype( uMaxFitness );
}

void Genetic::ExplicitSimulation::updateFitness()
{
	mFitness.lock();
	dMinFitness = -1.0;
	dTotalFitness = 0.0;
	mFitness.unlock();
	cWorkDone.lock();
	iWorkDone = xPop.getSize();
	cWorkDone.unlock();

	for( Population::iterator i = xPop.begin(); i; i++ )
	{
		mFitness.lock();
		bool bHas = hFitness.has( i.getKey() );
		mFitness.unlock();
		if( bHas )
		{
			setFitness( i.getKey(), hFitness.get( i.getKey() ) );
			cWorkDone.lock();
			iWorkDone--;
			cWorkDone.unlock();
		}
		else
		{
			qWork.enqueue( *i );
		}
	}
	
	cWorkDone.lock();
	if( iWorkDone == 0 )
	{
		cWorkDone.unlock();
		return;
	}

	while( iWorkDone > 0 )
		cWorkDone.wait();
	cWorkDone.unlock();
}

void Genetic::ExplicitSimulation::setFitness( Genetic::PhenotypeId id,
		double dFitness )
{
	Bu::MutexLocker ml( mFitness );
	if( dFitness < 0.0 )
		dFitness = 0.0;

	hFitness.insert( id, dFitness );

	dTotalFitness += dFitness;
	if( dMinFitness < 0.0 )
	{
		dMinFitness = dMaxFitness = dFitness;
		uMaxFitness = id;
	}
	else if( dMinFitness > dFitness )
	{
		dMinFitness = dFitness;
	}
	else if( dMaxFitness < dFitness )
	{
		dMaxFitness = dFitness;
		uMaxFitness = id;
	}
}

void Genetic::ExplicitSimulation::workDone()
{
	cWorkDone.lock();
	iWorkDone--;

	if( iWorkDone == 0 )
		cWorkDone.signal();
	cWorkDone.unlock();
}

void Genetic::ExplicitSimulation::setRunning( bool b )
{
	Bu::MutexLocker ml( mRunning );
	bRunning = b;
}

bool Genetic::ExplicitSimulation::isRunning()
{
	Bu::MutexLocker ml( mRunning );
	return bRunning;
}

Genetic::ExplicitSimulation::Processor::Processor(
		Genetic::ExplicitSimulation &rSim,
		Genetic::FitnessFunction *pFunc,
		Genetic::ExplicitSimulation::WorkQueue &rqWork,
		int iId ) :
	rSim( rSim ),
	pFunc( pFunc ),
	rqWork( rqWork ),
	iId( iId )
{
}

Genetic::ExplicitSimulation::Processor::~Processor()
{
	delete pFunc;
}

void Genetic::ExplicitSimulation::Processor::run()
{
	while( rSim.isRunning() )
	{
		Genetic::Phenotype *pPhen = rqWork.dequeue( true );
		if( pPhen == NULL )
			continue;

		double dFitness = (*pFunc)( pPhen );
		rSim.setFitness( pPhen->getId(), dFitness );
		rSim.workDone();
	}
}