/*
* Neural Tanks is a game created by Eddie O'Hagan that leverages Neural Networking
* to use for AI. This was created in Unreal Engine 4 which is developed by Epic Games.
* (Copyright Epic Games, Inc. All Rights Reserved.)
*
* Smart Tanks was originally developed by Mat Buckland in 2002 as an example
* on how Neural Networking can be used to train a set of minesweepers
* to collect mines. I (Eddie O'Hagan) updated this project in 2024 to use more
* Object Oriented Programming and modern practices. To see the original tutorial;
* visit <see href="http://www.ai-junkie.com/ann/evolved/nnt1.html">Neural Networking Tutorial</see>
* and <see href="http://www.ai-junkie.com/ga/intro/gat1.html">Genetic Algorithm Tutorial</see>
*/
#include "GeneticAlgorithm.h"
#include <algorithm>
#include <Kismet/GameplayStatics.h>

/// <summary>
/// Default constructor.
/// </summary>
UGeneticAlgorithm::UGeneticAlgorithm()
{
	myCurrentGeneration = 0U;
	myPopulationSize = 0U;
	myChromosomeLength = 0U;
	myCrossoverRate = 0.0;
	myMutationRate = 0.0;
	myBestChromosomeIndex = 0;
	myTotalFitness = 0;
	myBestFitness = 0;
	myAverageFitness = 0;
	myWorstFitness = 9999999;
}

/// <summary>
/// Parameterized constructor. Initializes the chromosomes for the provided Tanks.
/// </summary>
/// <param name="populationSize">The number of Tanks/Chromosomes in the population.</param>
/// <param name="crossoverRate">The crossover rate assigned to the Chromosomes.</param>
/// <param name="mutationRate">The mutation rate assigned to the Chromosomes.</param>
/// <param name="chromosomeLength">The number of weights in a single Chromosomes.</param>
/// <param name="tanks">The collection of Tanks that each have their own Chromosome.</param>
void UGeneticAlgorithm::Initialize(const int populationSize, const float crossoverRate, const float mutationRate, const int chromosomeLength, TArray<AAITankPawn*>& tanks)
{
	UChromosome* currChromosome;
	myCurrentGeneration = 0U;
	myPopulationSize = populationSize;
	myChromosomeLength = chromosomeLength;
	myCrossoverRate = crossoverRate;
	myMutationRate = mutationRate;
	myBestChromosomeIndex = 0;
	myTotalFitness = 0;
	myBestFitness = 0;
	myAverageFitness = 0;
	myWorstFitness = 9999999;

	for (int index = 0U; index < myPopulationSize; index++)
	{
		currChromosome = NewObject<UChromosome>();
		currChromosome->Initialize(myChromosomeLength, myCrossoverRate, myMutationRate);
		tanks[index]->SetChromosome(currChromosome);
	}
}

/// <summary>
/// Calculates the best fitness, the worst fitness, and the average fitness for this generation
/// of Chromosomes.
/// </summary>
void UGeneticAlgorithm::CalculateStatistics(TArray<AAITankPawn*>& tanks)
{
	float currFitness;

	myPopulationSize = tanks.Num();
	if (myPopulationSize >= 2)
	{
		for (int index = 0U; index < myPopulationSize; index++)
		{
			currFitness = tanks[index]->GetChromosomeFitness();
			myTotalFitness += currFitness;

			if (currFitness > myBestFitness)
			{
				myBestFitness = currFitness;
				myBestChromosomeIndex = index;
			}

			if (currFitness < myWorstFitness)
			{
				myWorstFitness = currFitness;
			}
		}

		if (myPopulationSize != 0)
		{
			myAverageFitness = myTotalFitness / myPopulationSize;
		}
	}
}

/// <summary>
/// Resets the total fitness, best fitness, and average fitness to 0.
/// Also resets the worst fitness to max (9999999).
/// </summary>
void UGeneticAlgorithm::ResetAlgorithm()
{
	myTotalFitness = 0;
	myBestFitness = 0;
	myAverageFitness = 0;
	myWorstFitness = 9999999;
}

/// <summary>
/// Returns the index of a randomly selected Tank/Chromosome based on its fitness. 
/// The way this function selects a chromosome is similar to a roulette wheel. Depending 
/// on the fitness, (the higher the better) the Tank/Chromosome has a better chance 
/// of being picked. This allows for Tanks/Chromosomes that have performed well to 
/// have a better chance of reproducing and being picked.
/// </summary>
/// <returns>The index of the randomly selected Tank/Chromosome.</returns>
int UGeneticAlgorithm::RoulettePickChromosome(TArray<AAITankPawn*>& tanks)
{
	int index;
	int numberOfTanks = tanks.Num();
	float scoreSlice;
	float scoreTotal = 0.0;
	float currTotal = 0.0;

	//Add up all of the scores.
	for (index = 0U; index < numberOfTanks; index++)
	{
		scoreTotal += tanks[index]->GetChromosomeFitness();
	}

	//Take a (random) percentage of the total scores (a slice) and make that our random range.
	scoreSlice = FMath::FRandRange(0.0f, 1.0f) * scoreTotal;

	//Add up all the fitness until it reaches the scoreSlice. 
	for (index = 0U; index < numberOfTanks; index++)
	{
		currTotal += tanks[index]->GetChromosomeFitness();
		if (currTotal >= scoreSlice)
		{
			//This chromosome falls in the range.
			return index;
		}
	}

	throw new std::exception("GeneticAlgorithm::RoulettePickChromosome() Failed, did not find a chromosome, this should never happen.");
	return 0;
}

/// <summary>
/// Called once a frame to update/run the Genetic Algorithm.
/// </summary>
void UGeneticAlgorithm::Update(float deltaTime, TArray<AAITankPawn*>& tanks)
{
	int indexOne;
	int indexTwo;
	UChromosome* parentOne;
	UChromosome* parentTwo;
	UChromosome* babyOne = NULL;
	UChromosome* babyTwo = NULL;
	TArray<UChromosome*> newPopulation;

	myPopulationSize = tanks.Num();

	//We need at least two tanks.
	if (myPopulationSize >= 2)
	{
		//reset the appropriate variables
		ResetAlgorithm();

		//sort the population (for scaling and elitism)
		tanks.Sort();

		CalculateStatistics(tanks);

		while (newPopulation.Num() < myPopulationSize)
		{
			indexOne = RoulettePickChromosome(tanks);
			indexTwo = RoulettePickChromosome(tanks);

			parentOne = tanks[indexOne]->GetChromosome();
			parentTwo = tanks[indexTwo]->GetChromosome();

			babyOne = NewObject<UChromosome>();
			babyTwo = NewObject<UChromosome>();

			if (parentOne->Crossover(parentOne, parentTwo, babyOne, babyTwo) == true)
			{
				babyOne->Mutate();
				babyTwo->Mutate();

				babyOne->SetFitness(0);
				babyTwo->SetFitness(0);
				newPopulation.Add(babyOne);
				newPopulation.Add(babyTwo);
			}
			else
			{
				parentOne->Mutate();
				parentTwo->Mutate();

				parentOne->SetFitness(0);
				parentTwo->SetFitness(0);
				newPopulation.Add(parentOne);
				newPopulation.Add(parentTwo);
			}
		}

		//Overwrite the old population with the new one.
		for (int index = 0; index < myPopulationSize; index++) 
		{
			tanks[index]->CopyChromosomeData(newPopulation[index]);
		}
	}
}

/// <summary>
/// Returns the average fitness in the current generation.
/// </summary>
/// <returns>The average fitness in the current generation.</returns>
float UGeneticAlgorithm::GetAverageFitness()
{
	return myAverageFitness;
}

/// <summary>
/// Returns the best fitness in the current generation.
/// </summary>
/// <returns>The best fitness in the current generation.</returns>
float UGeneticAlgorithm::GetBestFitness()
{
	return myBestFitness;
}

/// <summary>
/// Returns the Chromosome reference at the specified index.
/// </summary>
/// <param name="tankIndex">The index in the Tank/Chromosome collection.</param>
/// <param name="tanks">The collection of Tanks to search.</param>
/// <returns>The corresponding Chromosome from the collection based on the provided tankIndex.</returns>
UChromosome* UGeneticAlgorithm::GetChromosomeAtIndex(const int tankIndex, TArray<AAITankPawn*>& tanks)
{
	if (tankIndex < myPopulationSize)
	{
		return tanks[tankIndex]->GetChromosome();
	}
	
	throw new std::exception("GeneticAlgorithm::GetChromosomeAtIndex(const int) failed, chromosomeIndex is out of bounds");
	return NULL;
}

/// <summary>
/// Updates the fitness value for a specified Chromosome in the collection.
/// </summary>
/// <param name="chromosomeIndex">The index of the Tank/Chromosome in the collection.</param>
/// <param name="fitnessForChromosome">The new fitness of the Chromosome.</param>
void UGeneticAlgorithm::SetFitnessForChromosome(const int chromosomeIndex, const float fitnessForChromosome, TArray<AAITankPawn*>& tanks)
{
	UChromosome* currChromosome = GetChromosomeAtIndex(chromosomeIndex, tanks);
	currChromosome->SetFitness(fitnessForChromosome);
}