/*
* 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 "AITankPawn.h"
#include <Kismet/GameplayStatics.h>
#include "Particles/ParticleSystemComponent.h"
#include "Landscape.h"

/// <summary>
/// The Default constructor for the AI Tank Pawn. Enables tick and sets
/// default values.
/// </summary>
AAITankPawn::AAITankPawn() : Super()
{
	myIsDestroyed = false;
	myIsInitialized = false;
	mySpeed = 0;
	mySpeedMultiplier = 10;
	myNumOfReadyTankAmmo = 0;
	myNumShotsSinceLastHit = 0;
	myNumShotsSinceLastHitThreshold = 3;
	myFitnessScore = 0;
	myLeftTrackInput = 0;
	myRightTrackInput = 0;
	myRotation = 0;
	myMaxTurnRate = 0.3f;
	myClosestTankAmmunitionIndex = 0;
	myMinXPosition = 0.0f;
	myMinYPosition = 0.0f;
	myMaxXPosition = 0.0f;
	myMaxYPosition = 0.0f;
	myPositionBuffer = 100.0f;

	PrimaryActorTick.bCanEverTick = true;

	//For standard AI, one hit = destroyed.
	myMaximumHealthPoints = 1.0f;
	myMaximumArmorPoints = 0.0f;
	myCurrentHealthPoints = myMaximumHealthPoints;
	myCurrentArmorPoints = myMaximumArmorPoints;
}

/// <summary>
/// Called before Begin Play. Finds the Static Mesh component and Destroyed static mesh component.
/// </summary>
/// <param name="transform">The transform the actor was constructed at.</param>
void AAITankPawn::OnConstruction(const FTransform& transform)
{
	Super::OnConstruction(transform);

	myAIStaticMeshComponent = Cast<UStaticMeshComponent>(myComponents["Tank_Body"]);
	myAIDestroyedStaticMeshComponent = Cast<UStaticMeshComponent>(myComponents["DestroyedTank"]);
}

/// <summary>
/// Called when this AI Tank Pawn first loads into the level. Creates the Neural Network 
/// and sets this AI Tank to an "alive" state.
/// </summary>
void AAITankPawn::BeginPlay()
{
	FQuat actorRotation;

	Super::BeginPlay();

	myNeuralNetwork = NewObject<UNeuralNetwork>();

	myRotation = FMath::FRandRange(0.0f, 1.0f) * (2.0f * PI);

	SetTankRotation(myRotation);
	myLookAtVector = GetActorRotation().Vector();

	ResetTank();
}

/// <summary>
/// Starts the tank and sets its world/level bounds and its turn rate. In addition to this, ammunition for the tank 
/// can be provided.
/// 
/// If the tank goes past the bounds, it is simply moved to the other side (it wraps around like in pac-man).
/// </summary>
/// <param name="maxTurnRate">The fastest rate that the AI Tank can turn.</param>
/// <param name="minXLocation">The lower X bounds of where this AI Tank can exist in the world.</param>
/// <param name="minYLocation">The lower Y bounds of where this AI Tank can exist in the world.</param>
/// <param name="maxXLocation">The upper X bounds of where this AI Tank can exist in the world.</param>
/// <param name="maxYLocation">The upper Y bounds of where this AI Tank can exist in the world.</param>
/// <param name="tankAmmunitions">The ammunition for the tank to use.</param>
void AAITankPawn::Initialize(const float maxTurnRate, const float minXLocation, const float minYLocation, const float maxXLocation, const float maxYLocation, TArray<ATankAmmunitionItem*>& tankAmmunitions)
{
	myIsInitialized = true;
	myHasStarted = true;

	myMaxTurnRate = maxTurnRate;
	myMinXPosition = minXLocation + myPositionBuffer;
	myMinYPosition = minYLocation + myPositionBuffer;
	myMaxXPosition = maxXLocation - myPositionBuffer;
	myMaxYPosition = maxYLocation - myPositionBuffer;
	myAllTankAmmunitions = tankAmmunitions;
}

/// <summary>
/// Called when this AI Tank collides with an item in the world. In the AI's case specifically,
/// we increase the fitness of this AI's neural network every time it finds a Tank Ammunition.
/// </summary>
/// <param name="theItem">The item that this AI Tank collided/overlapped with.</param>
void AAITankPawn::OnTankOverlappedItem(ATankItem* theItem)
{
	Super::OnTankOverlappedItem(theItem);

	if (theItem->GetItemType() == EItemType::VE_AMMO && theItem->GetOwner() != this)
	{
		//We have discovered a TankAmmunition so increase fitness.
		if (myChromosome != NULL)
		{
			IncrementFitness();
		}

		myNumOfReadyTankAmmo++;
	}
}

/// <summary>
/// Called when this AI Tank is hit by projectile. This includes player
/// and AI projectiles (Friendly Fire is possible).
/// </summary>
/// <param name="theAttackingPawn">The Tank that landed the hit on this AI Tank.</param>
/// <param name="damage">The amount of damage applied to this AI Tank.</param>
void AAITankPawn::OnTankHit(ATankPawn* theAttackingPawn, float damage)
{
	Super::OnTankHit(theAttackingPawn, damage);

	DamageTank(damage);
	
	if (myCurrentHealthPoints <= 0.0f)
	{
		DestroyTank();
	}
}

/// <summary>
/// Called when this AI Tank successfully hits another Tank. This is used in 
/// the advanced boss AI where if the boss (Sassy Sam) does not hit the player 
/// after shooting a certain number times. This allows the boss to know it needs
/// to relocate in-order to get a better shot.
/// </summary>
/// <param name="theHitTankPawn">The Tank this AI Tank successfully hit.</param>
void AAITankPawn::OnSuccessfullyHitTank(ATankPawn* theHitTankPawn)
{
	myNumShotsSinceLastHit = 0;
}

/// <summary>
/// Spawns and fires a projectile from the cannon. Also decrements the amount of 
/// ammunition this AI Tank has.
/// </summary>
/// <param name="cannonReloadTime">The amount of time in seconds before this AI Tank can fire again.</param>
/// <returns>true if a projectile was spawned and fired, false otherwise.</returns>
bool AAITankPawn::FireCannon(const float cannonReloadTime)
{
	if (Super::FireCannon(cannonReloadTime) == true)
	{
		myNumOfReadyTankAmmo--;
		if (myNumOfReadyTankAmmo < 0)
		{
			myNumOfReadyTankAmmo = 0;
		}

		myNumShotsSinceLastHit = myNumShotsSinceLastHit++;

		//Check if the tank has missed enough to warrant repositioning for a better angle.
		if (myNumShotsSinceLastHit >= myNumShotsSinceLastHitThreshold)
		{
			//Attempt to reposition tank through the blueprint/behavior tree.
			AttemptToReposition();
		}

		return true;
	}

	return false;
}

/// <summary>
/// Copies the provided chromosome's data (and weights) to this tank's chromosome.
/// </summary>
/// <param name="chromosomeToCopy">The chromosome to copy data from.</param>
void AAITankPawn::CopyChromosomeData(UChromosome* chromosomeToCopy)
{
	myChromosome->CopyChromosomeData(chromosomeToCopy);
}

/// <summary>
/// Calculates the distance to all the TankAmmunitions and returns the location of the closest one.
/// This is used for AI navigation.
/// </summary>
/// <param name="tankAmmunitions">A collection of all the tank ammunition in the level/world.</param>
/// <returns>The smallest calculated distance found.</returns>
FVector AAITankPawn::FindClosestTankAmmunition(TArray<ATankAmmunitionItem*>& tankAmmunitions)
{
	float distanceToTankAmmunition;
	float closestTankAmmunitionSoFar = MAX_flt;
	FVector vClosestObject(0, 0, 0);
	FVector vectorToTankAmmunition;

	if (myIsDestroyed == false)
	{
		//Cycle through TankAmmunitions to find closest.
		for (int index = 0; index < tankAmmunitions.Num(); index++)
		{
			//Calculate the vector distance from this AI Tank to the current tank ammo.
			vectorToTankAmmunition = tankAmmunitions[index]->GetActorLocation() - GetActorLocation();

			//Get the magnitude of the vector (which is how far this tank would have to travel).
			distanceToTankAmmunition = vectorToTankAmmunition.Size();

			//Check if the current distance is smaller then the one we found so far.
			if (distanceToTankAmmunition < closestTankAmmunitionSoFar)
			{
				closestTankAmmunitionSoFar = distanceToTankAmmunition;

				vClosestObject = GetActorLocation() - tankAmmunitions[index]->GetActorLocation();

				myClosestTankAmmunitionIndex = index;
			}
		}
	}

	return vClosestObject;
}

/// <summary>
/// Increments the fitness score of this AAITankPawn.
/// </summary>
void AAITankPawn::IncrementFitness()
{
	if (myIsDestroyed == false)
	{
		myFitnessScore++;

		//Copy over the updated fitness score to the chromosome.
		myChromosome->SetFitness(myFitnessScore);
	}
}

/// <summary>
/// Zeros out the current fitness score and resets the static mesh.
/// </summary>
void AAITankPawn::ResetTank()
{
	myIsDestroyed = false;
	myFitnessScore = 0;

	myAIStaticMeshComponent->SetVisibility(true);
	myAIDestroyedStaticMeshComponent->SetVisibility(false);

	myDestroyedTankParticles->Deactivate();
}

/// <summary>
/// Destroys the tank rendering it immobile and no longer used in the Neural Network.
/// </summary>
void AAITankPawn::DestroyTank()
{
	if (myIsDestroyed == false)
	{
		myIsDestroyed = true;
		myFitnessScore = -1;

		myAIStaticMeshComponent->SetVisibility(false);
		myAIDestroyedStaticMeshComponent->SetVisibility(true);

		myDestroyedTankParticles->Activate();

		//Stop the AI from controlling the pawn when it is destroyed.
		if (Controller != NULL)
		{
			Controller->UnPossess();
		}
	}
}

/// <summary>
///	First we take sensor readings and feed these into the AI Tank's brain
/// (which is the Neural Network).
///
///	The inputs are:
///		1. A vector to the closest TankAmmunition (X)
///		2. A vector to the closest TankAmmunition (Y)
///		3. The AI Tank's current 'look at' vector (X)
///		4. The AI Tank's current 'look at' vector (Y)
/// 
///	We receive two outputs from the brain, the lTrack and the rTrack values.
///	So given a force for each track we calculate the resultant rotation 
///	and acceleration and apply to current velocity vector. In other words;
/// we use the lTrack output for how much force to turn the tank left, 
/// and the rTrack output for how much force to turn the tank right.
/// </summary>
/// <returns>true if no errors occurred, false otherwise.</returns>
bool AAITankPawn::ProcessNeuralNetwork()
{
	float rotForce;
	FVector currActorLocation;
	TArray<float> inputs;
	TArray<float> output;

	//Ensure this tank is running.
	if (myIsInitialized == true && myIsDestroyed == false)
	{
		//Calculate the distance to closest TankAmmunition.
		FVector closestDistance = FindClosestTankAmmunition(myAllTankAmmunitions);

		//Normalize it because at this point we are interested in only the direction.
		closestDistance.Normalize();

		//Add in the vector to closest TankAmmunition as the first and second input to the Neural Network.
		inputs.Add(closestDistance.X);
		inputs.Add(closestDistance.Y);

		//Calculate the look at vector (the direction this tank is facing).
		myLookAtVector = GetActorRotation().Vector();

		//Add in Tank's look at vector as the third and fourth input to the Neural Network.
		inputs.Add(myLookAtVector.X);
		inputs.Add(myLookAtVector.Y);

		//Use the Neural Network to process the output.
		output = myNeuralNetwork->CalculateOutput(inputs);

		//Ensure there were no errors in calculating the output.
		if (output.Num() < myNeuralNetwork->GetNumberOfOutputs())
		{
			return false;
		}

		//Assign the outputs to the Tank's left & right tracks.
		myLeftTrackInput = output[0];
		myRightTrackInput = output[1];

		//Calculate steering forces
		rotForce = myLeftTrackInput - myRightTrackInput;

		//Clamp the Tank's rotation between its max turn rate.
		rotForce = FMath::Clamp<float>(rotForce, -myMaxTurnRate, myMaxTurnRate);

		//Calculate the final rotation.
		myRotation += rotForce;

		//Calculate the forward speed based on how much force was given to the left and right tracks.
		mySpeed = (myLeftTrackInput + myRightTrackInput) * mySpeedMultiplier;

		//Update the world location and rotation of the tank based on the calculated speed and rotation force.
		SetTankLocationAndRotation(GetActorLocation(), rotForce, mySpeed);
	}

	return true;
}

/// <summary>
/// Used for sorting the Tanks based on its fitness score.
/// </summary>
/// <param name="otherTank">The other Tank to compare against this Tank.</param>
/// <returns>true if this Tank's fitness is less than the otherTank, false otherwise.</returns>
bool AAITankPawn::operator<(const AAITankPawn& otherTank) const
{
	return (myChromosome->GetFitness() < otherTank.GetChromosomeFitness());
}

/// <summary>
/// Returns the total number of weights in the entire Neural Network.
/// </summary>
/// <returns>The total number of weights in the entire Neural Network.</returns>
int AAITankPawn::GetNumberOfWeights() const
{
	return myNeuralNetwork->GetTotalNumberOfNeuronInputs();
}

/// <summary>
/// Returns the fitness score of this Tank's chromosome. This is usually equal to 
/// the fitness score of the Tank itself. 
/// </summary>
/// <returns>The fitness score of this Tank's chromosome</returns>
float AAITankPawn::GetChromosomeFitness() const
{
	return myChromosome->GetFitness();
}

/// <summary>
/// Returns the fitness score of this Tank.
/// </summary>
/// <returns>The fitness score of this Tank.</returns>
float AAITankPawn::GetFitness() const
{
	return myFitnessScore;
}

/// <summary>
/// Returns the x,y position of this Tank.
/// </summary>
/// <returns>The x,y position of this Tank.</returns>
FVector AAITankPawn::GetPosition() const
{
	return GetActorLocation();
}

/// <summary>
/// Returns the chromosome of this Tank.
/// </summary>
/// <returns>The chromosome of this Tank.</returns>
UChromosome* AAITankPawn::GetChromosome()
{
	return myChromosome;
}

/// <summary>
/// Returns the behavior tree for this Tank.
/// </summary>
/// <returns>The behavior tree for this Tank.</returns>
UBehaviorTree* AAITankPawn::GetBehaviorTree()
{
	return myBehaviorTree;
}

/// <summary>
/// Sets the fitness for this Tank's Neural Network.
/// </summary>
/// <param name="fitness">The new fitness for this Tank's Neural Network.</param>
void AAITankPawn::SetFitness(const float fitness)
{
	myFitnessScore = fitness;
}

/// <summary>
/// Sets the rotation angle (in radians) for this Tank.
/// </summary>
/// <param name="angleInRadians">The new rotation angle (in radians) for this Tank.</param>
void AAITankPawn::SetTankRotation(const float angleInRadians)
{
	SetActorRotation(FQuat(FVector::UpVector, angleInRadians));
}

/// <summary>
/// Sets the weights for this Tank's Neural Network.
/// </summary>
/// <param name="weights">The new weights for this Tank's Neural Network.</param>
void AAITankPawn::SetWeights(TArray<float>& weights)
{
	myNeuralNetwork->SetWeights(weights);
}

/// <summary>
/// Sets the chromosome for this Tank.
/// </summary>
/// <param name="chromosome">The new chromosome for this Tank.</param>
void AAITankPawn::SetChromosome(UChromosome* chromosome)
{
	myChromosome = chromosome;
}

/// <summary>
/// Directly sets the 3D world/level position of this AI Tank with respect to the terrain and
/// the Neural Network. Since the AI Tank doesn't use gravity, we set its Z location component
/// to the height of where this Tank is located (its X,Y location) on the landscape. 
/// </summary>
/// <param name="theLocation">The X,Y world/level location to set the AI tank to.</param>
/// <param name="theRotationAngle">The angle in radians to rotate this tank to.</param>
/// <param name="theSpeed">The speed to use for moving the tank.</param>
void AAITankPawn::SetTankLocationAndRotation(const FVector& theLocation, const float theRotationAngle, const float theSpeed)
{
	//Adjust the height of the tank if there is a landscape. 
	FVector currActorLocation = theLocation;
	currActorLocation.Z = GetLandscapeHeightAtTankLocation();

	//update position and rotation
	SetActorRotation(GetActorRotation() + FQuat(FVector::UpVector, theRotationAngle).Rotator());
	SetActorLocation(currActorLocation + (GetActorRotation().Vector() * theSpeed));

	//Reset tank position if it goes out of bounds.
	if (currActorLocation.X > myMaxXPosition)
	{
		SetActorLocation(FVector(myMinXPosition, currActorLocation.Y, currActorLocation.Z));
	}
	if (currActorLocation.X < myMinXPosition)
	{
		SetActorLocation(FVector(myMaxXPosition, currActorLocation.Y, currActorLocation.Z));
	}
	if (currActorLocation.Y > myMaxYPosition)
	{
		SetActorLocation(FVector(currActorLocation.X, myMinYPosition, currActorLocation.Z));
	}
	if (currActorLocation.Y < myMinYPosition)
	{
		SetActorLocation(FVector(currActorLocation.X, myMaxYPosition, currActorLocation.Z));
	}
}