// Copyright IHY.
#include "OptimizerGeometry.h"

#include "Engine/StaticMesh.h"
#include "MeshDescription.h"
#include "StaticMeshAttributes.h"
#include "StaticMeshResources.h"
#include "Rendering/PositionVertexBuffer.h"
#include "Materials/MaterialInterface.h"

namespace
{
	// Read LOD0 source geometry into compact arrays. Returns false if no source model.
	bool ReadFromMeshDescription(const UStaticMesh* Mesh, TArray<FVector>& OutRaw, TArray<FIntVector>& OutTris)
	{
		if (!Mesh->IsMeshDescriptionValid(0))
		{
			return false;
		}
		const FMeshDescription* MD = Mesh->GetMeshDescription(0);
		if (!MD || MD->Vertices().Num() == 0)
		{
			return false;
		}

		FStaticMeshConstAttributes Attributes(*MD);
		TVertexAttributesConstRef<FVector3f> Positions = Attributes.GetVertexPositions();

		OutRaw.Reset(MD->Vertices().Num());
		TMap<FVertexID, int32> Compact;
		Compact.Reserve(MD->Vertices().Num());
		for (const FVertexID VertexID : MD->Vertices().GetElementIDs())
		{
			Compact.Add(VertexID, OutRaw.Num());
			OutRaw.Add(FVector(Positions[VertexID]));
		}

		OutTris.Reset(MD->Triangles().Num());
		for (const FTriangleID TriangleID : MD->Triangles().GetElementIDs())
		{
			TArrayView<const FVertexID> Tri = MD->GetTriangleVertices(TriangleID);
			if (Tri.Num() == 3)
			{
				OutTris.Add(FIntVector(Compact[Tri[0]], Compact[Tri[1]], Compact[Tri[2]]));
			}
		}
		return OutRaw.Num() > 0;
	}

	// Fallback: built render-data LOD0 (welded/optimized). Tag the result as render-derived.
	bool ReadFromRenderData(const UStaticMesh* Mesh, TArray<FVector>& OutRaw, TArray<FIntVector>& OutTris)
	{
		const FStaticMeshRenderData* RD = Mesh->GetRenderData();
		if (!RD || RD->LODResources.Num() == 0)
		{
			return false;
		}
		const FStaticMeshLODResources& LOD = RD->LODResources[0];
		const FPositionVertexBuffer& Pos = LOD.VertexBuffers.PositionVertexBuffer;
		const uint32 NumVerts = Pos.GetNumVertices();
		if (NumVerts == 0)
		{
			return false;
		}

		OutRaw.Reset(NumVerts);
		for (uint32 i = 0; i < NumVerts; ++i)
		{
			OutRaw.Add(FVector(Pos.VertexPosition(i)));
		}

		TArray<uint32> Indices;
		LOD.IndexBuffer.GetCopy(Indices);
		OutTris.Reset(Indices.Num() / 3);
		for (int32 i = 0; i + 2 < Indices.Num(); i += 3)
		{
			OutTris.Add(FIntVector((int32)Indices[i], (int32)Indices[i + 1], (int32)Indices[i + 2]));
		}
		return true;
	}
}

namespace OptimizerGeometry
{
	void SymmetricEigen3x3(const double In[3][3], double OutValues[3], double OutVecs[3][3])
	{
		double a[3][3];
		double v[3][3] = { {1,0,0}, {0,1,0}, {0,0,1} };
		for (int32 i = 0; i < 3; ++i)
		{
			for (int32 j = 0; j < 3; ++j)
			{
				a[i][j] = In[i][j];
			}
		}

		// Cyclic Jacobi rotations on the three off-diagonal entries.
		for (int32 Sweep = 0; Sweep < 64; ++Sweep)
		{
			const double Off = FMath::Abs(a[0][1]) + FMath::Abs(a[0][2]) + FMath::Abs(a[1][2]);
			if (Off < 1e-18)
			{
				break;
			}
			static const int32 P[3] = { 0, 0, 1 };
			static const int32 Q[3] = { 1, 2, 2 };
			for (int32 k = 0; k < 3; ++k)
			{
				const int32 p = P[k];
				const int32 q = Q[k];
				if (FMath::Abs(a[p][q]) < 1e-300)
				{
					continue;
				}
				const double Theta = (a[q][q] - a[p][p]) / (2.0 * a[p][q]);
				double t = (Theta >= 0.0 ? 1.0 : -1.0) / (FMath::Abs(Theta) + FMath::Sqrt(Theta * Theta + 1.0));
				const double c = 1.0 / FMath::Sqrt(t * t + 1.0);
				const double s = t * c;

				// Rotate a: a = Jᵀ a J
				const double app = a[p][p];
				const double aqq = a[q][q];
				const double apq = a[p][q];
				a[p][p] = c * c * app - 2.0 * s * c * apq + s * s * aqq;
				a[q][q] = s * s * app + 2.0 * s * c * apq + c * c * aqq;
				a[p][q] = 0.0;
				a[q][p] = 0.0;
				const int32 r = 3 - p - q; // the third index
				const double arp = a[r][p];
				const double arq = a[r][q];
				a[r][p] = c * arp - s * arq;
				a[p][r] = a[r][p];
				a[r][q] = s * arp + c * arq;
				a[q][r] = a[r][q];

				// Accumulate eigenvectors: v = v J
				for (int32 i = 0; i < 3; ++i)
				{
					const double vip = v[i][p];
					const double viq = v[i][q];
					v[i][p] = c * vip - s * viq;
					v[i][q] = s * vip + c * viq;
				}
			}
		}

		int32 Order[3] = { 0, 1, 2 };
		const double Diag[3] = { a[0][0], a[1][1], a[2][2] };
		// Sort indices by eigenvalue DESC.
		if (Diag[Order[0]] < Diag[Order[1]]) { Swap(Order[0], Order[1]); }
		if (Diag[Order[0]] < Diag[Order[2]]) { Swap(Order[0], Order[2]); }
		if (Diag[Order[1]] < Diag[Order[2]]) { Swap(Order[1], Order[2]); }

		for (int32 i = 0; i < 3; ++i)
		{
			OutValues[i] = Diag[Order[i]];
			OutVecs[0][i] = v[0][Order[i]];
			OutVecs[1][i] = v[1][Order[i]];
			OutVecs[2][i] = v[2][Order[i]];
		}
	}

	bool ExtractGeom(const UStaticMesh* Mesh, float WeldEps, bool bWantPositions, FOptMeshGeom& Out)
	{
		Out = FOptMeshGeom();
		if (!Mesh)
		{
			return false;
		}

		TArray<FVector> Raw;
		TArray<FIntVector> Tris;
		if (ReadFromMeshDescription(Mesh, Raw, Tris))
		{
			Out.bRenderDerived = false;
		}
		else if (ReadFromRenderData(Mesh, Raw, Tris))
		{
			Out.bRenderDerived = true;
		}
		else
		{
			return false;
		}

		Out.RawVertexCount = Raw.Num();
		Out.TriangleCount = Tris.Num();

		// Surface area + signed volume from the raw triangle soup (welding doesn't change these).
		double Area = 0.0;
		double Vol6 = 0.0;
		for (const FIntVector& T : Tris)
		{
			const FVector& A = Raw[T.X];
			const FVector& B = Raw[T.Y];
			const FVector& C = Raw[T.Z];
			Area += 0.5 * FVector::CrossProduct(B - A, C - A).Size();
			Vol6 += FVector::DotProduct(A, FVector::CrossProduct(B, C));
		}
		Out.SurfaceArea = Area;
		Out.Volume = FMath::Abs(Vol6) / 6.0;

		// Weld positions onto a grid so seam-split duplicates collapse to one unique position.
		const double InvEps = (WeldEps > UE_KINDA_SMALL_NUMBER) ? (1.0 / (double)WeldEps) : 1.0;
		TMap<FIntVector, int32> Grid;
		Grid.Reserve(Raw.Num());
		TArray<FVector> Welded;
		Welded.Reserve(Raw.Num());
		for (const FVector& P : Raw)
		{
			const FIntVector Key(
				FMath::RoundToInt(P.X * InvEps),
				FMath::RoundToInt(P.Y * InvEps),
				FMath::RoundToInt(P.Z * InvEps));
			if (!Grid.Contains(Key))
			{
				Grid.Add(Key, Welded.Num());
				Welded.Add(P);
			}
		}
		Out.WeldedVertexCount = Welded.Num();
		if (Welded.Num() == 0)
		{
			return false;
		}

		// Centroid (mean of welded positions) + bounds.
		FVector Sum(0.0);
		FBox Box(ForceInit);
		for (const FVector& P : Welded)
		{
			Sum += P;
			Box += P;
		}
		Out.Centroid = Sum / (double)Welded.Num();
		Out.LocalBounds = Box;

		// Covariance of centered welded cloud -> eigenvalues (rotation invariant).
		double Cov[3][3] = { {0,0,0}, {0,0,0}, {0,0,0} };
		double MaxR = 0.0;
		for (const FVector& P : Welded)
		{
			const FVector d = P - Out.Centroid;
			Cov[0][0] += d.X * d.X; Cov[0][1] += d.X * d.Y; Cov[0][2] += d.X * d.Z;
			Cov[1][1] += d.Y * d.Y; Cov[1][2] += d.Y * d.Z;
			Cov[2][2] += d.Z * d.Z;
			MaxR = FMath::Max(MaxR, d.Size());
		}
		const double InvN = 1.0 / (double)Welded.Num();
		Cov[0][0] *= InvN; Cov[0][1] *= InvN; Cov[0][2] *= InvN;
		Cov[1][1] *= InvN; Cov[1][2] *= InvN; Cov[2][2] *= InvN;
		Cov[1][0] = Cov[0][1]; Cov[2][0] = Cov[0][2]; Cov[2][1] = Cov[1][2];
		double EVecs[3][3];
		SymmetricEigen3x3(Cov, Out.EigenValues, EVecs);

		// Radial histogram of |v - centroid|, normalized by the max radius then to sum 1.
		if (MaxR > UE_KINDA_SMALL_NUMBER)
		{
			const double InvMax = (double)FOptMeshGeom::NumRadialBins / MaxR;
			for (const FVector& P : Welded)
			{
				const double r = (P - Out.Centroid).Size();
				int32 Bin = (int32)FMath::FloorToDouble(r * InvMax);
				Bin = FMath::Clamp(Bin, 0, FOptMeshGeom::NumRadialBins - 1);
				Out.RadialHistogram[Bin] += 1.0;
			}
			for (int32 i = 0; i < FOptMeshGeom::NumRadialBins; ++i)
			{
				Out.RadialHistogram[i] *= InvN;
			}
		}

		// Material / section signature (kept separate from geometry grouping).
		uint32 H = 0;
		const TArray<FStaticMaterial>& Mats = Mesh->GetStaticMaterials();
		for (const FStaticMaterial& M : Mats)
		{
			const FString N = M.MaterialInterface ? M.MaterialInterface->GetPathName() : TEXT("None");
			H = HashCombine(H, GetTypeHash(N));
		}
		Out.MaterialHash = H;
		Out.SectionCount = Mats.Num();

		if (bWantPositions)
		{
			Out.Positions = MoveTemp(Welded);
		}

		Out.bValid = true;
		return true;
	}
}