#ifndef GENALG_H
#define GENALG_H

#include <climits>
#include <cmath>
//#include <cstdlib> // for rand()
#include <cstring>

#include <bitset>
#include <random>
#include <vector>

namespace genalg {

template<typename T>
struct Phenotype {
	double fitness;
	T chromosome;
};

//typedef double (*FitnessFunc)(Phenotype* p);

template<typename T, typename F, typename U, typename R>
std::vector<Phenotype<T>> Iterate(
		std::vector<Phenotype<T>>* population,
		const F &FitnessFunc,
		U user_data,
		R* rng,
		double crossover_rate,
		double mutation_rate );

template<typename T, typename R>
std::vector<Phenotype<T>> RandomGarbage(
		size_t members,
		R* rng );

#ifdef GENALG_IMPLEMENTATION

/** \brief Breeds a new generation of phenotypes.
 *
 * This is the main evolution function. Feed it your population and get
 * a slightly evolved population out the other end. Do this many times
 * to optimize for the fitness function.
 *
 * The new population's size will be a multiple of 2 due to how the
 * recombination works. Aside from that, the input and output
 * populations will always be the same size.
 *
 * When experimenting with crossover rates and mutation rates, a value
 * of 0.7 has been suggested for the crossover and 0.01 or 0.001 for the
 * mutation.
 *
 * Note that the human generational mutation rate is roughly on the
 * order of 0.00000002, or 1 out of every 500 million base pairs. For
 * most contemporary programs, this would translate to negligibly few
 * mutations.
 *
 * Example:
 * `population = genalg::Iterate( &population, MyFunc, my_data, &mt, 0.7, 0.01 );`
 *
 * @param  population     Pointer to `std::vector` of phenotypes.
 * @param  FitnessFunc    Function that ingests a
 *                        `Phenotype<MyChromosome>*` and user data and
 *                        returns a floating-point fitness value.
 * @param  user_data      A variable of any type to be passed as the
 *                        second parameter to `FitnessFunc`.
 * @param  rng            Pointer to instance of an RNG from `<random>`.
 * @param  crossover_rate The probability that 2 chosen chromosomes
 *                        will swap bits.
 * @param  mutation_rate  A bit's probability of flipping.
 * @return `std::vector` of the new generation of phenotypes.
 */
template<typename T, typename F, typename U, typename R>
std::vector<Phenotype<T>> Iterate(
		std::vector<Phenotype<T>>* population,
		const F &FitnessFunc,
		U user_data,
		R* rng,
		double crossover_rate,
		double mutation_rate ){

	const size_t chromosome_bytes = sizeof(T);
	const size_t chromosome_bits = chromosome_bytes * CHAR_BIT;
	// TODO(fluffrabbit): Investigate the contiguousness of std::bitset.
	std::bitset<chromosome_bits> bits[3];

	double total_fitness = 0.0;

	// Decide the fitness of each member of the population.
	for( auto &p : *population ){
		p.fitness = FitnessFunc( &p, user_data );
		// Yes, negative fitness is possible.
		total_fitness += std::max( 0.0, static_cast<double>( p.fitness ) );
	}

	// If the total fitness is infinity or NaN, just try not to cause a crash.
	if( std::isinf( total_fitness ) || std::isnan( total_fitness ) ){
		total_fitness = 0.0;
	}

	// Initialize the random number distributions.
	std::uniform_real_distribution<double> roulette_dist( 0.0, total_fitness );
	std::uniform_real_distribution<double> unit_dist( 0.0, 1.0 );
	std::uniform_int_distribution<size_t> gene_dist( 0, chromosome_bits );

	std::vector<Phenotype<T>> new_population;

	// Iterate through the population, producing 2 new members for every 2 old members.
	for( size_t half_n = 0; half_n < population->size() / 2; half_n++ ){
		double
			slot = 0.0,
			ball1 = roulette_dist( *rng ),
			ball2 = roulette_dist( *rng );
		bool
			got_bits1 = false,
			got_bits2 = false;

		// Pick 2 mates using roulette wheel selection.
		for( auto &p : *population ){
			int target = 0;
			if( ball1 >= slot && ball1 < slot + p.fitness ){
				target = 1;
				got_bits1 = true;
			}
			if( ball2 >= slot && ball2 < slot + p.fitness ){
				target = 2;
				got_bits2 = true;
			}
			if( target > 0 ){
				std::memcpy(
					static_cast<void*>( target == 1 ? &bits[1] : &bits[2] ),
					static_cast<void*>( &p.chromosome ),
					chromosome_bytes
				);
			}
			if( got_bits1 && got_bits2 ){
				break;
			}
			slot += p.fitness;
		}

		if( unit_dist( *rng ) < crossover_rate ){
			// Swap genes.
			size_t gene_offset_bits = gene_dist( *rng );
			size_t gene_offset_bytes = gene_offset_bits / CHAR_BIT;

			// Use bits[0] as a buffer for the first chromosome's data.
			std::memcpy(
				static_cast<void*>( &bits[0] ),
				static_cast<void*>( &bits[1] ),
				chromosome_bytes
			);

			if( gene_offset_bytes > 0 ){
				// Swap bytes from the starts of the chromosomes to the offset.
				// This is faster than iterating over every bit.
				std::memcpy(
					static_cast<void*>( &bits[1] ),
					static_cast<void*>( &bits[2] ),
					gene_offset_bytes
				);
				std::memcpy(
					static_cast<void*>( &bits[2] ),
					static_cast<void*>( &bits[0] ),
					gene_offset_bytes
				);
			}

			// Swap the bits in the byte of the crossover.
			for( size_t i = gene_offset_bytes * CHAR_BIT; i < gene_offset_bits; i++ ){
				bits[1][i] = bits[2][i];
				bits[2][i] = bits[0][i];
			}
		}

		// Mutate the chromosomes.
		for( size_t i = 0; i < chromosome_bits; i++ ){
			if( unit_dist( *rng ) < mutation_rate ){
				bits[1].flip( i );
			}
			if( unit_dist( *rng ) < mutation_rate ){
				bits[2].flip( i );
			}
		}

		// Add the phenotypes to the new population.
		for( int i = 0; i < 2; i++ ){
			new_population.push_back( { 0.0, {} } );
			std::memcpy(
				static_cast<void*>( &new_population.back().chromosome ),
				static_cast<void*>( i == 0 ? &bits[1] : &bits[2] ),
				chromosome_bytes
			);
		}
	}

	return new_population;
}

/** \brief Returns a randomized `std::vector` of `Phenotype`s.
 *
 * Returns a vector containing `Phenotype`s with randomized chromosomes.
 * The randomization happens per byte without knowledge of what the
 * chromosome affects. This is a worst-case initialization for a quick
 * start that <b>should not be used in production</b>.
 *
 * Example:
 * `auto population = genalg::RandomGarbage<MyChromosome>( 20, &mt );`
 *
 * You should initialize with plausible random values instead.
 *
 * @param  <T>     struct type to use for chromosomes.
 * @param  members Number of `Phenotype`s to create.
 * @param  rng     Pointer to an instance of an RNG from `<random>`.
 * @return `std::vector` of randomized `Phenotype`s.
 */
template<typename T, typename R>
std::vector<Phenotype<T>> RandomGarbage(
		size_t members,
		R* rng ){

	std::vector<Phenotype<T>> population;

	unsigned char bytes[sizeof(T)];

	std::uniform_int_distribution<int> byte_dist( 0, UCHAR_MAX );

	// Create a random initial population.
	for( size_t i = 0; i < members; i++ ){
		for( unsigned char &b : bytes ){
			//b = std::rand() % ( UCHAR_MAX + 1 );
			b = byte_dist( *rng );
		}
		population.push_back( {
			0.0,
			*static_cast<T*>( static_cast<void*>( bytes ) )
		} );
	}
	return population;
}

#endif // GENALG_IMPLEMENTATION

} // namespace genalg

#endif // GENALG_H