// All rights reserved Dominik Pavlicek 2023


#include "Layout/ComboActionForceDirectedSolveLayoutStrategy.h"

#include "ComboActionGraph.h"
#include "Ed/EdComboActionGraph.h"
#include "Ed/EdComboActionGraphNode.h"
#include "Nodes/ComboActionGraphNode.h"
#include "Settings/ComboActionGraphEditorSettings.h"

UComboActionForceDirectedSolveLayoutStrategy::UComboActionForceDirectedSolveLayoutStrategy()
{
	this->bRandomInit = false;
	this->CoolDownRate = 10;
	this->InitTemperature = 10.f;
}

static inline float CoolDown(float Temp, float CoolDownRate)
{
	if (Temp < .01) return .01;
	return Temp - (Temp / CoolDownRate);
}

static inline float GetAttractForce(float X, float K)
{
	return (X * X) / K;
}

static inline float GetRepulseForce(float X, float k)
{
	return X != 0 ? k * k / X : TNumericLimits<float>::Max();
}

void UComboActionForceDirectedSolveLayoutStrategy::Layout(UEdGraph *InEdGraph)
{
	this->EdGraph = Cast<UEdComboActionGraph>(InEdGraph);
	check(this->EdGraph != nullptr);

	this->EdGraph->RebuildComboActionGraph();
	this->Graph = this->EdGraph->GetComboActionGraph();
	check(this->Graph != nullptr);

	if (this->Settings != nullptr)
	{
		this->OptimalDistance = Settings->GetOptimalDistance();
		this->MaxIteration = Settings->GetMaxIteration();
		this->bRandomInit = Settings->IsRandomInit();
	}

	FBox2D PreTreeBound(ForceInitToZero);
	for (int32 i = 0; i < Graph->RootNodes.Num(); ++i)
	{
		PreTreeBound = this->LayoutOneTree(Graph->RootNodes[i], PreTreeBound);
	}
}

FBox2D UComboActionForceDirectedSolveLayoutStrategy::LayoutOneTree(UComboActionGraphNode *RootNode, const FBox2D &PreTreeBound)
{
	float Temp = this->InitTemperature;
	FBox2D TreeBound = GetActualBounds(RootNode);
	TreeBound.Min.X += PreTreeBound.Max.X + this->OptimalDistance;
	TreeBound.Max.X += PreTreeBound.Max.X + this->OptimalDistance;

	if (this->bRandomInit)
	{
		this->RandomLayoutOneTree(RootNode, TreeBound);
	}

	float RepulseForce, AttractForce, Distance;
	FVector2D Diff;

	TMap<UEdGraphNode*, FVector2D> NodeToDisplacement;

	for (const TObjectPtr<UEdGraphNode> &EdNode : this->EdGraph->Nodes)
	{
		NodeToDisplacement.Add(EdNode, FVector2D(0.0f, 0.0f));
	}

	for (int32 IterationNum = 0; IterationNum < this->MaxIteration; IterationNum++)
	{
		// Calculate the repulsive forces.
		for (int32 i = 0; i < this->EdGraph->Nodes.Num(); i++)
		{
			for (int32 j = 0; j < this->EdGraph->Nodes.Num(); j++)
			{
				if (i == j)
				{
					continue;
				}
				Diff.X = this->EdGraph->Nodes[i]->NodePosX - this->EdGraph->Nodes[j]->NodePosX;
				Diff.Y = this->EdGraph->Nodes[i]->NodePosY - this->EdGraph->Nodes[j]->NodePosY;
				Distance = Diff.Size();
				Diff.Normalize();

				RepulseForce = Distance > 2 * this->OptimalDistance ? 0 : GetRepulseForce(Distance, this->OptimalDistance);

				NodeToDisplacement[this->EdGraph->Nodes[i]] += Diff * RepulseForce;
			}
		}

		// Calculate the attractive forces.
		int Level = 0;
		TArray<UComboActionGraphNode*> CurrLevelNodes = { RootNode };
		TArray<UComboActionGraphNode*> NextLevelNodes;

		while (CurrLevelNodes.Num() != 0)
		{
			for (int32 i = 0; i < CurrLevelNodes.Num(); i++)
			{
				UComboActionGraphNode *Node = CurrLevelNodes[i];
				check(Node != nullptr);

				UEdComboActionGraphNode *EdNode_ParentNode = this->EdGraph->NodeMap[Node];

				for (int32 j = 0; j < Node->ChildrenNodes.Num(); j++)
				{
					NextLevelNodes.Add(Node->ChildrenNodes[j]);

					UEdComboActionGraphNode *EdNode_ChildNode = this->EdGraph->NodeMap[Node->ChildrenNodes[j]];

					Diff.X = EdNode_ChildNode->NodePosX - EdNode_ParentNode->NodePosY;
					Diff.Y = EdNode_ChildNode->NodePosY - EdNode_ParentNode->NodePosY;
					Distance = Diff.Size();
					Diff.Normalize();

					AttractForce = GetAttractForce(Distance, this->OptimalDistance);

					NodeToDisplacement[EdNode_ParentNode] += Distance * Diff;
					NodeToDisplacement[EdNode_ChildNode] -= Distance * Diff;
				}
			}

			CurrLevelNodes = NextLevelNodes;
			NextLevelNodes.Reset();
			Level++;
		}

		for (UEdGraphNode *EdNode : this->EdGraph->Nodes)
		{
			Distance = NodeToDisplacement[EdNode].Size();
			NodeToDisplacement[EdNode].Normalize();

			float Minimum = Distance < Temp ? Distance : Temp;
			EdNode->NodePosX += NodeToDisplacement[EdNode].X * Minimum;
			EdNode->NodePosY += NodeToDisplacement[EdNode].Y * Minimum;
		}

		Temp = CoolDown(Temp, this->CoolDownRate);
	}

	FBox2D ActualBound = GetActualBounds(RootNode);

	FVector2D Center = ActualBound.GetCenter();
	FVector2D TreeCenter = TreeBound.GetCenter();

	FVector2D Scale = (TreeBound.Max - TreeBound.Min) / (ActualBound.Max - ActualBound.Min);

	for (UEdGraphNode *EdNode : EdGraph->Nodes)
	{
		EdNode->NodePosX = TreeCenter.X + Scale.X * (EdNode->NodePosX - Center.X);
		EdNode->NodePosY = TreeCenter.Y + Scale.Y * (EdNode->NodePosY - Center.Y);
	}

	return TreeBound;
}