// Fill out your copyright notice in the Description page of Project Settings.

#include "DataComponent.h"
#include "Runtime/Core/Public/Misc/Paths.h"
#include "Runtime/Core/Public/HAL/FileManagerGeneric.h"
#include "HDF5Library.h"
#include "Plot3D.h"
#include "PawnBase.h"
#include "Plot2D.h"
#include "PlotCanvas.h"
#include "Runtime/Engine/Classes/Engine/World.h"
#include "Runtime/Engine/Classes/Kismet/KismetMathLibrary.h"
#include "Runtime/Core/Public/Async/ParallelFor.h"

UDataComponent::UDataComponent()
{
	D3_BrushAlpha = 0;
	D3_SpriteAlpha = 255;
	D3_BrushOverridesColor = false;
	D3_BrushColor = FColor(0, 0, 0, 0);
	D3_UseTransparency = true;
	D3_SpriteSizeFactor = 50;
	UseFileBounds = false;
}


// Called when the game starts
void UDataComponent::BeginPlay()
{
	Super::BeginPlay();

	if (!TransferFunction)
	{
		UE_LOG(LogTemp, Error, TEXT("No transfer function set!"));
		FGenericPlatformMisc::RequestExit(false);
	}

	// Find Files
	FString Directory = FPaths::ProjectDir();
	Directory += FString("/Resources/");
	FString Path = Directory + "*.h5*";
	FFileManagerGeneric::Get().FindFiles(FileNames, *Path, true, false);

	if (!FileNameToLoadOnStart.IsEmpty())
	{
		LoadFile(FileNameToLoadOnStart);
	}
}

void UDataComponent::LoadFile(FString FileName)
{
	// Load File
	TArray<FString> TempStringArray = TArray<FString>();
	LoadedFile = UHDF5Library::LoadFile(FileName, TempStringArray);
	CurrentFileName = FileName;

	if (!Cast<APawnBase>(GetOwner())->Plot3D)
	{
		FVector Location(-1610.0f, -341.0f, 100.0f);
		FRotator Rotation(0.0f, 0.0f, 0.0f);
		FActorSpawnParameters SpawnInfo;
		APlot3D* Plot3D = GetWorld()->SpawnActor<APlot3D>(Location, Rotation, SpawnInfo);
		Cast<APawnBase>(GetOwner())->Plot3D = Plot3D;
		Plot3D->Init(this);
	}

	/*
	The order would be "1,10,11,12,13,14,15,16,17,18,19,2,20,21,22,23,24" without sorting padded strings!
	*/
	for (int i = 0; i < TempStringArray.Num(); ++i)
	{
		int32 Length = TempStringArray[i].Len();
		for (int j = 0; j < (12 - Length); ++j) // pad the string to 10 characters
		{
			TempStringArray[i].InsertAt(5, " "); // 5 is the position after "step#"
		}
	}
	TempStringArray.Sort([](const FString& One, const FString& Two) {
		return One < Two;
	});
	for (int i = 0; i < TempStringArray.Num(); ++i)
	{
		TempStringArray[i] = TempStringArray[i].Replace(TEXT(" "), TEXT(""));
	}

	GroupNames.Empty();
	for (FString String : TempStringArray)
	{
		GroupNames.Emplace(String);
	}

	if (UseFileBounds)
	{
		SetFileBounds();
	}

	LoadStep(GroupNames[0], true);
}

void UDataComponent::LoadStep(FString GroupName, bool ResetBrush)
{
	UHDF5Library::LoadGroup(LoadedFile, GroupName, X, Y, Z, U, V, W);

	CurrentGroupName = GroupName;

	// Prepare Data
	Num = X.Num();
	if (OverrideMaxPointCount > 0 && OverrideMaxPointCount < Num)
	{
		Num = OverrideMaxPointCount;
	}

	if (ResetBrush) // TODO
	{
		IsBrushed = TArray<bool>();
		IsBrushed.Init(false, Num);
	}
	
	Length.Empty();
	for (int32 i = 0; i < X.Num(); i++)
	{
		Length.Emplace(FVector(U[i], V[i], W[i]).Size());
	}

	if (!UseFileBounds)
	{
		int32 MaxIndex, MinIndex;
		UKismetMathLibrary::MaxOfFloatArray(Length, MaxIndex, MaxLength);
		UKismetMathLibrary::MinOfFloatArray(Length, MinIndex, MinLength);
		UKismetMathLibrary::MaxOfFloatArray(X, MaxIndex, MaxX);
		UKismetMathLibrary::MinOfFloatArray(X, MinIndex, MinX);
		UKismetMathLibrary::MaxOfFloatArray(Y, MaxIndex, MaxY);
		UKismetMathLibrary::MinOfFloatArray(Y, MinIndex, MinY);
		UKismetMathLibrary::MaxOfFloatArray(Z, MaxIndex, MaxZ);
		UKismetMathLibrary::MinOfFloatArray(Z, MinIndex, MinZ);
		UKismetMathLibrary::MaxOfFloatArray(U, MaxIndex, MaxU);
		UKismetMathLibrary::MinOfFloatArray(U, MinIndex, MinU);
		UKismetMathLibrary::MaxOfFloatArray(V, MaxIndex, MaxV);
		UKismetMathLibrary::MinOfFloatArray(V, MinIndex, MinV);
		UKismetMathLibrary::MaxOfFloatArray(W, MaxIndex, MaxW);
		UKismetMathLibrary::MinOfFloatArray(W, MinIndex, MinW);

		UpdateDataScaleFactor();
	}
	
	PColor.Empty();
	for (int32 i = 0; i < Num; i++)
	{
		FColor TempColor = CalculateColorOfPoint(i, MaxLength, MinLength);

		PColor.Emplace(FVector(TempColor.R, TempColor.G, TempColor.B));
	}

	Cast<APawnBase>(GetOwner())->Plot3D->OnLoadStep(!ResetBrush);

	APawnBase* BasePawn = Cast<APawnBase>(GetOwner());
	for (APlot2D* Plot2D : BasePawn->Plots2D)
	{
		Plot2D->PlotCanvas->SetXAxisType(Plot2D->PlotCanvas->XAxisType); // this refreshes the min and max values
		Plot2D->PlotCanvas->SetYAxisType(Plot2D->PlotCanvas->YAxisType); // this refreshes the min and max values
		Plot2D->Redraw();
	}
}

FColor UDataComponent::CalculateColorOfPoint(int32 Index, float Min, float Max)
{
	float NormalizedValue = UKismetMathLibrary::NormalizeToRange(Length[Index], Min, Max);
	FVector ColorVec = TransferFunction->GetVectorValue(NormalizedValue);
	FColor Color = FColor((uint8)(ColorVec.X * 255), (uint8)(ColorVec.Y * 255), (uint8)(ColorVec.Z * 255));
	return Color;
}

void UDataComponent::SetFileBounds()
{
	TArray<float> XMinTemp = TArray<float>();
	TArray<float> XMaxTemp = TArray<float>();
	TArray<float> YMinTemp = TArray<float>();
	TArray<float> YMaxTemp = TArray<float>();
	TArray<float> ZMinTemp = TArray<float>();
	TArray<float> ZMaxTemp = TArray<float>();
	TArray<float> UMinTemp = TArray<float>();
	TArray<float> UMaxTemp = TArray<float>();
	TArray<float> VMinTemp = TArray<float>();
	TArray<float> VMaxTemp = TArray<float>();
	TArray<float> WMinTemp = TArray<float>();
	TArray<float> WMaxTemp = TArray<float>();
	TArray<float> LengthTemp = TArray<float>();

	for (FString Group : GroupNames)
	{
		TArray<float> XTemp = TArray<float>();
		TArray<float> YTemp = TArray<float>();
		TArray<float> ZTemp = TArray<float>();
		TArray<float> UTemp = TArray<float>();
		TArray<float> VTemp = TArray<float>();
		TArray<float> WTemp = TArray<float>();
		UHDF5Library::LoadGroup(LoadedFile, Group, XTemp, YTemp, ZTemp, UTemp, VTemp, WTemp);

		for (int32 i = 0; i < XTemp.Num(); i++)
		{
			LengthTemp.Emplace(FVector(UTemp[i], VTemp[i], WTemp[i]).Size());
		}

		float TempMaxX, TempMaxY, TempMaxZ, TempMaxU, TempMaxV, TempMaxW, TempMinX, TempMinY, TempMinZ, TempMinU, TempMinV, TempMinW;
		int32 MaxIndex, MinIndex;
		UKismetMathLibrary::MaxOfFloatArray(XTemp, MaxIndex, TempMaxX); XMaxTemp.Emplace(TempMaxX);
		UKismetMathLibrary::MinOfFloatArray(XTemp, MinIndex, TempMinX); XMinTemp.Emplace(TempMinX);
		UKismetMathLibrary::MaxOfFloatArray(YTemp, MaxIndex, TempMaxY); YMaxTemp.Emplace(TempMaxY);
		UKismetMathLibrary::MinOfFloatArray(YTemp, MinIndex, TempMinY); YMinTemp.Emplace(TempMinY);
		UKismetMathLibrary::MaxOfFloatArray(ZTemp, MaxIndex, TempMaxZ); ZMaxTemp.Emplace(TempMaxZ);
		UKismetMathLibrary::MinOfFloatArray(ZTemp, MinIndex, TempMinZ); ZMinTemp.Emplace(TempMinZ);
		UKismetMathLibrary::MaxOfFloatArray(UTemp, MaxIndex, TempMaxU); UMaxTemp.Emplace(TempMaxU);
		UKismetMathLibrary::MinOfFloatArray(UTemp, MinIndex, TempMinU); UMinTemp.Emplace(TempMinU);
		UKismetMathLibrary::MaxOfFloatArray(VTemp, MaxIndex, TempMaxV); VMaxTemp.Emplace(TempMaxV);
		UKismetMathLibrary::MinOfFloatArray(VTemp, MinIndex, TempMinV); VMinTemp.Emplace(TempMinV);
		UKismetMathLibrary::MaxOfFloatArray(WTemp, MaxIndex, TempMaxW); WMaxTemp.Emplace(TempMaxW);
		UKismetMathLibrary::MinOfFloatArray(WTemp, MinIndex, TempMinW); WMinTemp.Emplace(TempMinW);
	}

	int32 MaxIndex, MinIndex;
	UKismetMathLibrary::MaxOfFloatArray(LengthTemp, MaxIndex, MaxLength);
	UKismetMathLibrary::MinOfFloatArray(LengthTemp, MinIndex, MinLength);
	UKismetMathLibrary::MaxOfFloatArray(XMaxTemp, MaxIndex, MaxX);
	UKismetMathLibrary::MinOfFloatArray(XMinTemp, MinIndex, MinX);
	UKismetMathLibrary::MaxOfFloatArray(YMaxTemp, MaxIndex, MaxY);
	UKismetMathLibrary::MinOfFloatArray(YMinTemp, MinIndex, MinY);
	UKismetMathLibrary::MaxOfFloatArray(ZMaxTemp, MaxIndex, MaxZ);
	UKismetMathLibrary::MinOfFloatArray(ZMinTemp, MinIndex, MinZ);
	UKismetMathLibrary::MaxOfFloatArray(UMaxTemp, MaxIndex, MaxU);
	UKismetMathLibrary::MinOfFloatArray(UMinTemp, MinIndex, MinU);
	UKismetMathLibrary::MaxOfFloatArray(VMaxTemp, MaxIndex, MaxV);
	UKismetMathLibrary::MinOfFloatArray(VMinTemp, MinIndex, MinV);
	UKismetMathLibrary::MaxOfFloatArray(WMaxTemp, MaxIndex, MaxW);
	UKismetMathLibrary::MinOfFloatArray(WMinTemp, MinIndex, MinW);

	UpdateDataScaleFactor();
}

void UDataComponent::UpdateDataScaleFactor()
{
	float Range = FMath::Max3(MaxX - MinX, MaxY - MinY, MaxZ - MinZ);
	DataScaleFactor = 2 / Range;
}

void UDataComponent::RebuildAll()
{
	APawnBase* BasePawn = Cast<APawnBase>(GetOwner());
	BasePawn->Plot3D->RebuildCloud();
	for (APlot2D* Plot2D : BasePawn->Plots2D)
	{
		Plot2D->Redraw();
	}
}

void UDataComponent::AddFilter_Planes(FVector NormalRight, FVector LocationRight, bool UseTwoPlanes, FVector NormalLeft, FVector LocationLeft, bool Add, bool Rebuild, bool Invert)
{
	TArray<int32> Indexes = TArray<int32>();
	Indexes.Init(-1, Num);

	APlot3D* Plot3D = Cast<APawnBase>(GetOwner())->Plot3D;
	ParallelFor(Num, [this, Invert, &Indexes, LocationRight, LocationLeft, NormalRight, NormalLeft, UseTwoPlanes, Add, Plot3D](int32 CurrIdx)
	{
		FVector WorldLoc = Plot3D->CalculateWorldLocOfPoint(CurrIdx);
		FVector DifferenceRight = WorldLoc - LocationRight;
		FVector DifferenceLeft = WorldLoc - LocationLeft;
		float DotRight = UKismetMathLibrary::Dot_VectorVector(DifferenceRight, NormalRight);
		float DotLeft = UKismetMathLibrary::Dot_VectorVector(DifferenceLeft, NormalLeft);

		if (Invert)
		{
			if (!UseTwoPlanes)
			{
				if (DotRight >= 0.f)
				{
					Indexes[CurrIdx] = CurrIdx;
				}
			}
			else
			{
				if (!(DotRight <= 0.f && DotLeft >= 0.f))
				{
					Indexes[CurrIdx] = CurrIdx;
				}
			}
		}
		else
		{
			if (DotRight <= 0.f && (DotLeft >= 0.f || !UseTwoPlanes))
			{
				Indexes[CurrIdx] = CurrIdx;
			}
		}
	});

	BrushPoints(Indexes, Add, Rebuild);
}

void UDataComponent::AddFilter_ValueRange(EAxisType XType, float XMin, float XMax, EAxisType YType, float YMin, float YMax, bool Add, bool Rebuild, bool UseTwoValues)
{
	TArray<float> * XArray = nullptr;
	TArray<float> * YArray = nullptr;
	TArray<int32> Indexes = TArray<int32>();

	switch (XType)
	{
	case EAxisType::X:
		XArray = &X;
		break;
	case EAxisType::Y:
		XArray = &Y;
		break;
	case EAxisType::Z:
		XArray = &Z;
		break;
	case EAxisType::U:
		XArray = &U;
		break;
	case EAxisType::V:
		XArray = &V;
		break;
	case EAxisType::W:
		XArray = &W;
		break;
	case EAxisType::LengthUVW:
		XArray = &Length;
		break;
	default:
		break;
	}

	switch (YType)
	{
	case EAxisType::X:
		YArray = &X;
		break;
	case EAxisType::Y:
		YArray = &Y;
		break;
	case EAxisType::Z:
		YArray = &Z;
		break;
	case EAxisType::U:
		YArray = &U;
		break;
	case EAxisType::V:
		YArray = &V;
		break;
	case EAxisType::W:
		YArray = &W;
		break;
	case EAxisType::LengthUVW:
		YArray = &Length;
		break;
	default:
		break;
	}

	/*
	Indexes.Init(-1, Points.Num());
	ParallelFor(Points.Num(), [this, XArray, YArray, XMin, XMax, YMin, YMax, Add, &Indexes](int32 CurrIdx)
	{
		if ((*XArray)[CurrIdx] >= XMin && (*XArray)[CurrIdx] <= XMax && (*YArray)[CurrIdx] >= YMin && (*YArray)[CurrIdx] <= YMax)
		{
			Indexes[CurrIdx] = CurrIdx;
		}
	});*/

	for (int i = 0; i < Num; ++i)
	{
		if ((*XArray)[i] >= XMin && (*XArray)[i] <= XMax && (!UseTwoValues || ((*YArray)[i] >= YMin && (*YArray)[i] <= YMax)))
		{
			Indexes.Emplace(i);
		}
	}

	BrushPoints(Indexes, Add, Rebuild);
}

void UDataComponent::AddFilter_Box(FTransform Transform, FVector Extent, bool Add, bool Rebuild)
{
	TArray<int32> Indexes = TArray<int32>();

	/*for (int i = 0; i < Points.Num(); ++i)
	{
		FVector WorldLoc = CalculateWorldLocOfPoint(&Points[i]);
		if (UKismetMathLibrary::IsPointInBoxWithTransform(WorldLoc, Transform, Extent))
		{
			Indexes.Emplace(i);
		}
	}*/

	APlot3D* Plot3D = Cast<APawnBase>(GetOwner())->Plot3D;
	Indexes.Init(-1, Num);
	ParallelFor(Num, [this, Transform, Extent, Add, &Indexes, Plot3D](int32 CurrIdx)
	{
		FVector WorldLoc = Plot3D->CalculateWorldLocOfPoint(CurrIdx);
		if (UKismetMathLibrary::IsPointInBoxWithTransform(WorldLoc, Transform, Extent))
		{
			Indexes[CurrIdx] = CurrIdx;
		}
	});

	BrushPoints(Indexes, Add, Rebuild);
}

void UDataComponent::AddFilter_Sphere(FVector Location, float Radius, bool Add, bool Rebuild)
{
	TArray<int32> Indexes = TArray<int32>();

	/*for (int i = 0; i < Points.Num(); ++i)
	{
		FVector WorldLoc = CalculateWorldLocOfPoint(&Points[i]);
		if ((Location - WorldLoc).Size() <= Radius)
		{
			Indexes.Emplace(i);
		}
	}*/

	APlot3D* Plot3D = Cast<APawnBase>(GetOwner())->Plot3D;
	Indexes.Init(-1, Num);
	ParallelFor(Num, [this, Location, Radius, Add, &Indexes, Plot3D](int32 CurrIdx)
	{
		FVector WorldLoc = Plot3D->CalculateWorldLocOfPoint(CurrIdx);
		if ((Location - WorldLoc).Size() <= Radius)
		{
			Indexes[CurrIdx] = CurrIdx;
		}
	});

	BrushPoints(Indexes, Add, Rebuild);
}

void UDataComponent::BrushPoints(TArray<int32> BrushPoints, bool Add, bool Rebuild)
{
	for (int i = 0; i < BrushPoints.Num(); ++i)
	{
		BrushPoint(BrushPoints[i], Add);
	}

	APawnBase* BasePawn = Cast<APawnBase>(GetOwner());
	BasePawn->Plot3D->BrushPoints(BrushPoints, Rebuild);
	if (Rebuild)
	{
		for (APlot2D* Plot2D : BasePawn->Plots2D)
		{
			Plot2D->Redraw();
		}
	}
}

void UDataComponent::BrushPoint(int32 PointIndex, bool Add)
{
	if (PointIndex == -1)
	{
		return;
	}

	IsBrushed[PointIndex] = Add;
}

// does not use the standard BrushPoints, for performance reasons
void UDataComponent::ResetPoints(bool Rebuild)
{
	for (int i = 0; i < Num; ++i)
	{
		BrushPoint(i, false);
	}

	APawnBase* BasePawn = Cast<APawnBase>(GetOwner());
	BasePawn->Plot3D->ResetPoints(Rebuild);
	if (Rebuild)
	{
		for (APlot2D* Plot2D : BasePawn->Plots2D)
		{
			Plot2D->Redraw();
		}
	}
}

// does not use the standard BrushPoints, for performance reasons
void UDataComponent::InvertPoints(bool Rebuild)
{
	for (int i = 0; i < Num; ++i)
	{
		BrushPoint(i, !IsBrushed[i]);
	}

	/*ParallelFor(Points.Num(), [this](int32 CurrIdx)
	{
		BrushPoint(CurrIdx, Points[CurrIdx].IsEnabled());
	});*/

	/*ParallelFor(4, [this](int32 CurrIdx)
	{
		int32 Start = Points.Num() / 4 * CurrIdx;
		int32 End = Points.Num() / 4 * (CurrIdx + 1);
		for (int i = Start; i < End; ++i)
		{
			BrushPoint(i, Points[i].IsEnabled());
		}
	});*/

	APawnBase* BasePawn = Cast<APawnBase>(GetOwner());
	BasePawn->Plot3D->InvertPoints(Rebuild);
	if (Rebuild)
	{
		for (APlot2D* Plot2D : BasePawn->Plots2D)
		{
			Plot2D->Redraw();
		}
	}
}