/*
* 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 "NeuralTanksGameMode.h"
#include "NeuralTanksGameLevel.h"
#include "NeuralTanksCheatManager.h"
#include "Utils/NeuralTanksUtils.h"
#include "Pawns/TankPawn.h"
#include "Pawns/AITankPawn.h"
#include "Pawns/PlayerTankPawn.h"
#include "NeuralTanksHud.h"
#include <Kismet/GameplayStatics.h>

/// <summary>
/// Default constructor.
/// </summary>
ANeuralTanksGameMode::ANeuralTanksGameMode()
{
	myGeneticAlgorithm = NULL;
	myCurrSecondsInGeneration = 0;
	myGenerationCounter = 0;
	myCrossoverRate = 0.7f;
	myMutationRate = 0.1f;
	myNumSecondsPerGeneration = 12;
	myMaxAITankTurnRate = 0.3f;
	myPlayerProjectileVelocity = 6000.0f;
	myMinWorldBoundsPosition = FVector(-30000.0f, -30000.0f, 0.0f);
	myMaxWorldBoundsPosition = FVector(30000.0f, 30000.0f, 0.0f);

	static ConstructorHelpers::FClassFinder<APlayerTankPawn> PlayerPawnClassFinder(TEXT("/Game/Blueprints/Tanks/BP_PlayerTankPawn"));
	DefaultPawnClass = PlayerPawnClassFinder.Class;

	HUDClass = ANeuralTanksHud::StaticClass();

	PrimaryActorTick.bCanEverTick = true;

	myIsInitialized = false;
	myHasRoundEnded = false;
}

/// <summary>
/// Called when the level first loads. This triggers the initialize function after 5 seconds.
/// We want to give the player a few seconds to get their bearings. 
/// </summary>
void ANeuralTanksGameMode::BeginPlay()
{
	FTimerHandle timerHandle;

	GetWorldTimerManager().SetTimer(timerHandle, this, &ANeuralTanksGameMode::Initialize, 5.0f, false);
}

/// <summary>
/// Initializes all of the enemy tanks in the level and their corresponding Neural Networks.
/// Also records the total number of tank ammunitions and the min/max level bounds.
/// </summary>
void ANeuralTanksGameMode::Initialize()
{
	int index;
	FText localizedTankAmmoDescription;
	ANeuralTanksGameLevel* theGameLevel;
	TArray<AActor*> outFoundActors;

	//Get references to all of the AI Tanks in the world.
	UGameplayStatics::GetAllActorsOfClass(GetWorld(), AAITankPawn::StaticClass(), outFoundActors);
	for (index = 0; index < outFoundActors.Num(); index++)
	{
		myTanks.Add(Cast<AAITankPawn>(outFoundActors[index]));
	}
	myNumTanks = myTanks.Num();
	outFoundActors.Empty();

	if (myNumTanks > 0)
	{
		//Get references to all of the TankAmmunitions in the world.
		UGameplayStatics::GetAllActorsOfClass(GetWorld(), ATankAmmunitionItem::StaticClass(), outFoundActors);
		for (index = 0; index < outFoundActors.Num(); index++)
		{
			myTankAmmunitions.Add(Cast<ATankAmmunitionItem>(outFoundActors[index]));
		}
		myNumTankAmmunitions = myTankAmmunitions.Num();

		//Get the total number of weights used in the tank's NN so we can initialize the Genetic Algorithm.
		myTotalNumWeightsInNeuralNetwork = myTanks[0]->GetNumberOfWeights();

		//Initialize the Genetic Algorithm class
		myGeneticAlgorithm = NewObject<UGeneticAlgorithm>();
		myGeneticAlgorithm->Initialize(myNumTanks, myCrossoverRate, myMutationRate, myTotalNumWeightsInNeuralNetwork, myTanks);

		//Get the AI bounds for the AIs from the current level's blueprint.
		theGameLevel = Cast<ANeuralTanksGameLevel>(GetWorld()->GetLevelScriptActor());
		if (theGameLevel != NULL)
		{
			myMinWorldBoundsPosition = theGameLevel->GetMinAIBoundsPosition();
			myMaxWorldBoundsPosition = theGameLevel->GetMaxAIBoundsPosition();
		}
		else
		{
			UNeuralTanksUtils::WriteToLog(TEXT("ANeuralTanksGameMode::Initialize(): Failed to find Neural Tanks Game Level Blueprint."));
		}

		//Get the localized description for the tank ammo.
		localizedTankAmmoDescription = FText::FromStringTable(TEXT("/Game/Text/NeuralTanksStringTable"), TEXT("Ammunition_NeuralTankAmmunitionDescription"));

		//Initialize the TankAmmunitions.
		for (index = 0; index < myNumTankAmmunitions; index++)
		{
			myTankAmmunitions[index]->Initialize(1, EItemType::VE_AMMO, myTankAmmunitions[index]->GetFName().ToString(), localizedTankAmmoDescription, myMinWorldBoundsPosition.X, myMinWorldBoundsPosition.Y, myMaxWorldBoundsPosition.X, myMaxWorldBoundsPosition.Y);
		}

		//Initialize the tanks and set the weights for their neural networks.
		for (index = 0; index < myNumTanks; index++)
		{
			myTanks[index]->Initialize(myMaxAITankTurnRate, myMinWorldBoundsPosition.X, myMinWorldBoundsPosition.Y, myMaxWorldBoundsPosition.X, myMaxWorldBoundsPosition.Y, myTankAmmunitions);
			myTanks[index]->SetWeights(myTanks[index]->GetChromosome()->GetWeights());
		}

		myPlayerPawn = Cast<APlayerTankPawn>(UGameplayStatics::GetPlayerPawn(GetWorld(), 0));

		myIsInitialized = true;
		myHasRoundEnded = false;
	}
	else
	{
		myTotalNumWeightsInNeuralNetwork = 0;
	}
}

/// <summary>
/// Called once per frame. This checks to see if the current generation is done and
/// the tanks are ready for reproduction. When this happens, the generation is incremented
/// and the genetic algorithm calculates the new/updated weights for the Neural Networks.
/// </summary>
/// <param name="deltaTime">Time in milliseconds between frames.</param>
void ANeuralTanksGameMode::Tick(float deltaTime)
{
	Super::Tick(deltaTime);

	if (myIsInitialized == true)
	{
		//Check if the player has defeated all the tanks.
		CheckForEndOfRound();

		//The updating for the individual Enemy Tank's AI/Neural Network is done through the UE4 Behavior Tree.

		//Checks to see if this generation has finished. When the generation is finished,
		//the reproduction process begins for creating/modifying the Neural Network Data.
		myCurrSecondsInGeneration += deltaTime;
		if(myCurrSecondsInGeneration >= myNumSecondsPerGeneration && myNumTanks >= 2)
		{
			//Another generation has been completed. Time to run the GA and update the tanks with their new NNs
			PrintStatistics();

			//Update the stats to be used in our stat window
			myAverageFitnesses.Add(myGeneticAlgorithm->GetAverageFitness());
			myBestFitnesses.Add(myGeneticAlgorithm->GetBestFitness());

			//Increment the generation counter and prevent integer overflow.
			myGenerationCounter = FMath::Clamp(myGenerationCounter + 1, 0, MAX_int32);

			//Reset cycles
			myCurrSecondsInGeneration = 0;

			//Run the Genetic Algorithm to create a new population
			myGeneticAlgorithm->Update(deltaTime, myTanks);

			//Insert the new (hopefully)improved brains back into the tanks and reset their positions etc
			for (int index = 0; index < myNumTanks; index++)
			{
				myTanks[index]->SetWeights(myTanks[index]->GetChromosome()->GetWeights());

				myTanks[index]->ResetTank();
			}
		}
	}
}

/// <summary>
/// Checks to see if all of the enemy tanks are destroyed or
/// the player is killed.
/// </summary>
void ANeuralTanksGameMode::CheckForEndOfRound()
{
	//Removes the dead tanks from the collection.
	RemoveDeadTanks();

	if (myHasRoundEnded == false)
	{
		if (myNumTanks == 0)
		{
			RoundEnded(true);
		}
		else if (myPlayerPawn->GetCurrentHealth() <= 0.0)
		{
			RoundEnded(false);
		}
	}
}

/// <summary>
/// Called when this round has ended for some reason.
/// </summary>
/// <param name="hasPlayerWon">true if the player destroyed all the enemy tanks, false if the player died.</param>
void ANeuralTanksGameMode::RoundEnded(bool hasPlayerWon)
{
	myIsInitialized = false;
	myHasRoundEnded = true;

	//Call blueprint facing event.
	OnRoundEndedBlueprint(hasPlayerWon);

	//Call C++ registered delegates.
	if (OnRoundEnded_Event.IsBound() == true)
	{
		OnRoundEnded_Event.Broadcast(hasPlayerWon);
	}
}

/// <summary>
/// Prints the Neural Network statistics to the log.
/// (This is only used for debugging and doesn't run in shipped builds).
/// </summary>
void ANeuralTanksGameMode::PrintStatistics()
{
	UNeuralTanksUtils::WriteToLog(FString::Printf(TEXT("Generation: %i		Best Fitness:%f		Average Fitness: :%f"), myGenerationCounter, GetBestFitness(), GetAverageFitness()));
}

/// <summary>
/// Checks the fitness for all the tanks in the collection and if the fitness
/// of a particular tank is below zero (negative) it is removed from the collection.
/// </summary>
void ANeuralTanksGameMode::RemoveDeadTanks()
{
	for (int index = 0; index < myTanks.Num(); index++)
	{
		if (myTanks[index]->GetFitness() < 0.0f)
		{
			myTanks.RemoveAt(index);
		}
	}

	myNumTanks = myTanks.Num();
}

/// <summary>
/// Returns true if the enemy tanks and level is initialized, false otherwise.
/// </summary>
/// <returns>true if the enemy tanks and level is initialized, false otherwise.</returns>
bool ANeuralTanksGameMode::GetIsInitialized()
{
	return myIsInitialized;
}

/// <summary>
/// Returns the current generation number of the Neural Network.
/// </summary>
/// <returns>The current generation number of the Neural Network.</returns>
int ANeuralTanksGameMode::GetCurrentGeneration()
{
	return myGenerationCounter;
}

/// <summary>
/// Returns the average calculated fitness of the Neural Network.
/// </summary>
/// <returns>The average calculated fitness of the Neural Network.</returns>
float ANeuralTanksGameMode::GetAverageFitness()
{
	return myGeneticAlgorithm->GetAverageFitness();
}

/// <summary>
/// Returns the best calculated fitness in the Neural Network.
/// </summary>
/// <returns>The best calculated fitness in the Neural Network.</returns>
float ANeuralTanksGameMode::GetBestFitness()
{
	return myGeneticAlgorithm->GetBestFitness();
}