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

namespace
{
	// --- small helpers ---------------------------------------------------------

	bool RelClose(double A, double B, double RelTol)
	{
		const double Scale = FMath::Max3(FMath::Abs(A), FMath::Abs(B), 1.0);
		return FMath::Abs(A - B) <= RelTol * Scale;
	}

	struct FRot3
	{
		double M[3][3] = { {1,0,0},{0,1,0},{0,0,1} };
		bool bImproper = false;
	};

	FVector ApplyRot(const double R[3][3], const FVector& V)
	{
		return FVector(
			R[0][0] * V.X + R[0][1] * V.Y + R[0][2] * V.Z,
			R[1][0] * V.X + R[1][1] * V.Y + R[1][2] * V.Z,
			R[2][0] * V.X + R[2][1] * V.Y + R[2][2] * V.Z);
	}

	// All 48 signed permutation rotations split into proper (det+1) and improper (det-1).
	const TArray<FRot3>& AxisRotations()
	{
		static TArray<FRot3> Cached;
		if (Cached.Num() == 0)
		{
			static const int32 Perm[6][3] = { {0,1,2},{0,2,1},{1,0,2},{1,2,0},{2,0,1},{2,1,0} };
			static const int32 Parity[6] = { 1, -1, -1, 1, 1, -1 };
			for (int32 p = 0; p < 6; ++p)
			{
				for (int32 sgn = 0; sgn < 8; ++sgn)
				{
					const int32 e0 = (sgn & 1) ? -1 : 1;
					const int32 e1 = (sgn & 2) ? -1 : 1;
					const int32 e2 = (sgn & 4) ? -1 : 1;
					const int32 Det = Parity[p] * e0 * e1 * e2;
					FRot3 R;
					FMemory::Memzero(R.M, sizeof(R.M));
					R.M[0][Perm[p][0]] = e0;
					R.M[1][Perm[p][1]] = e1;
					R.M[2][Perm[p][2]] = e2;
					R.bImproper = (Det < 0);
					Cached.Add(R);
				}
			}
		}
		return Cached;
	}

	// Uniform spatial hash over a point cloud for order-independent nearest-neighbour queries.
	struct FPointGrid
	{
		double InvCell = 1.0;
		const TArray<FVector>* Pts = nullptr;
		TMap<FIntVector, TArray<int32>> Cells;

		void Build(const TArray<FVector>& InPts, double Cell)
		{
			Pts = &InPts;
			InvCell = 1.0 / FMath::Max(Cell, 1e-6);
			Cells.Reserve(InPts.Num());
			for (int32 i = 0; i < InPts.Num(); ++i)
			{
				Cells.FindOrAdd(KeyOf(InPts[i])).Add(i);
			}
		}

		FIntVector KeyOf(const FVector& P) const
		{
			return FIntVector(
				FMath::FloorToInt(P.X * InvCell),
				FMath::FloorToInt(P.Y * InvCell),
				FMath::FloorToInt(P.Z * InvCell));
		}

		// Nearest point within MaxDist. Returns index or INDEX_NONE; OutDist set when found.
		int32 Nearest(const FVector& Q, double MaxDist, double& OutDist) const
		{
			const FIntVector Base = KeyOf(Q);
			double Best = MaxDist * MaxDist;
			int32 BestIdx = INDEX_NONE;
			// Search enough neighbour cells to cover MaxDist (cell size = 1/InvCell), not just +/-1,
			// so reported deviations for over-tolerance vertices aren't truncated.
			const int32 Rad = FMath::Max(1, FMath::CeilToInt(MaxDist * InvCell));
			for (int32 dz = -Rad; dz <= Rad; ++dz)
			for (int32 dy = -Rad; dy <= Rad; ++dy)
			for (int32 dx = -Rad; dx <= Rad; ++dx)
			{
				const TArray<int32>* Bucket = Cells.Find(Base + FIntVector(dx, dy, dz));
				if (!Bucket)
				{
					continue;
				}
				for (int32 Idx : *Bucket)
				{
					const double D2 = FVector::DistSquared((*Pts)[Idx], Q);
					if (D2 < Best)
					{
						Best = D2;
						BestIdx = Idx;
					}
				}
			}
			if (BestIdx != INDEX_NONE)
			{
				OutDist = FMath::Sqrt(Best);
			}
			return BestIdx;
		}
	};

	// Transform canonical point p -> member space via q = s*R*(p - cA) + cB.
	FVector Map(const double R[3][3], double s, const FVector& cA, const FVector& cB, const FVector& P)
	{
		return s * ApplyRot(R, P - cA) + cB;
	}

	// Cheap score over a subsample: sum of NN distances, with a fixed penalty for misses (lower is better).
	double ScoreCandidate(const TArray<FVector>& A, const FPointGrid& BGrid,
		const double R[3][3], double s, const FVector& cA, const FVector& cB,
		double Tol, int32 Stride)
	{
		const double MissPenalty = Tol * 1000.0;
		double Sum = 0.0;
		for (int32 i = 0; i < A.Num(); i += Stride)
		{
			const FVector Q = Map(R, s, cA, cB, A[i]);
			double Dist;
			if (BGrid.Nearest(Q, Tol, Dist) != INDEX_NONE)
			{
				Sum += Dist;
			}
			else
			{
				Sum += MissPenalty;
			}
		}
		return Sum;
	}

	// Full verification of a winning candidate: per-vertex max/rms deviation + reverse coverage.
	void VerifyCandidate(const TArray<FVector>& A, const TArray<FVector>& B, const FPointGrid& BGrid,
		const double R[3][3], double s, const FVector& cA, const FVector& cB,
		double Tol, double& OutMaxDev, double& OutRms, bool& OutCoverageOK)
	{
		double MaxDev = 0.0;
		double SumSq = 0.0;
		int32 Count = 0;
		TBitArray<> Hit(false, B.Num());
		bool bAllWithin = true;
		for (const FVector& P : A)
		{
			const FVector Q = Map(R, s, cA, cB, P);
			double Dist;
			const int32 Idx = BGrid.Nearest(Q, Tol * 4.0, Dist); // widen search so over-tol matches still report
			if (Idx == INDEX_NONE || Dist > Tol)
			{
				bAllWithin = false;
				MaxDev = FMath::Max(MaxDev, Idx == INDEX_NONE ? Tol * 4.0 : Dist);
			}
			else
			{
				Hit[Idx] = true;
				MaxDev = FMath::Max(MaxDev, Dist);
			}
			SumSq += (Idx == INDEX_NONE) ? (Tol * 4.0) * (Tol * 4.0) : Dist * Dist;
			++Count;
		}
		OutMaxDev = MaxDev;
		OutRms = (Count > 0) ? FMath::Sqrt(SumSq / Count) : 0.0;
		OutCoverageOK = bAllWithin && (Hit.CountSetBits() >= B.Num());
	}
}

namespace OptimizerMatcher
{
	bool FingerprintCompatible(const FOptMeshGeom& A, const FOptMeshGeom& B, const FOptimizerScanSettings& S)
	{
		if (!A.bValid || !B.bValid)
		{
			return false;
		}
		if (A.bRenderDerived != B.bRenderDerived)
		{
			return false; // never cross-compare source vs render-derived geometry
		}
		// Hard gates: welded vertex + triangle counts are exactly rotation/mirror invariant.
		if (A.WeldedVertexCount != B.WeldedVertexCount || A.TriangleCount != B.TriangleCount)
		{
			return false;
		}
		if (!S.bMergeAcrossMaterials && A.MaterialHash != B.MaterialHash)
		{
			return false;
		}

		const double RelTol = FMath::Max((double)S.ScalarRelTolerance, 1e-6);

		if (S.bAllowUniformScale)
		{
			// Require area/volume/eigenvalue ratios to agree on a single scale factor s.
			if (A.SurfaceArea <= 0 || B.SurfaceArea <= 0)
			{
				return false;
			}
			const double sArea = FMath::Sqrt(B.SurfaceArea / A.SurfaceArea);
			const double sVol = (A.Volume > 0 && B.Volume > 0) ? FMath::Pow(B.Volume / A.Volume, 1.0 / 3.0) : sArea;
			const double sEig = (A.EigenValues[0] > 0 && B.EigenValues[0] > 0) ? FMath::Sqrt(B.EigenValues[0] / A.EigenValues[0]) : sArea;
			if (!RelClose(sArea, sVol, 0.02) || !RelClose(sArea, sEig, 0.02))
			{
				return false;
			}
		}
		else
		{
			if (!RelClose(A.SurfaceArea, B.SurfaceArea, RelTol))
			{
				return false;
			}
			if (A.Volume > 0 && B.Volume > 0 && !RelClose(A.Volume, B.Volume, RelTol))
			{
				return false;
			}
			for (int32 i = 0; i < 3; ++i)
			{
				if (!RelClose(A.EigenValues[i], B.EigenValues[i], FMath::Max(RelTol, 1e-3)))
				{
					return false;
				}
			}
		}

		// Radial histograms are normalized per-mesh (scale tolerant); compare by L1.
		double L1 = 0.0;
		for (int32 i = 0; i < FOptMeshGeom::NumRadialBins; ++i)
		{
			L1 += FMath::Abs(A.RadialHistogram[i] - B.RadialHistogram[i]);
		}
		return L1 <= 0.02;
	}

	bool RecoverDelta(const FOptMeshGeom& A, const FOptMeshGeom& B, const FOptimizerScanSettings& S, FOptDelta& Out)
	{
		Out = FOptDelta();
		if (A.Positions.Num() == 0 || B.Positions.Num() == 0)
		{
			return false;
		}

		const double Tol = FMath::Max((double)S.AcceptTolerance, 1e-5);
		FPointGrid BGrid;
		BGrid.Build(B.Positions, Tol);

		// Scoring stride keeps the 24/48 candidate scan cheap on dense meshes.
		const int32 Stride = FMath::Max(1, A.Positions.Num() / 1500);

		auto BuildMatrix = [&](const double R[3][3], double s, const FVector& cA, const FVector& cB) -> FMatrix
		{
			// Linear part L = s*R (column form); UE row-vector matrix is L^T with translation row tt.
			const FVector tt = cB - s * ApplyRot(R, cA);
			FMatrix Mtx = FMatrix::Identity;
			for (int32 i = 0; i < 3; ++i)
			{
				for (int32 j = 0; j < 3; ++j)
				{
					Mtx.M[i][j] = s * R[j][i];
				}
			}
			Mtx.M[3][0] = tt.X; Mtx.M[3][1] = tt.Y; Mtx.M[3][2] = tt.Z;
			return Mtx;
		};

		// --- Exact-only mode: just test true identity (same orientation AND pivot) ---
		if (S.Precision == EOptimizerMatchPrecision::ExactOnly)
		{
			static const double I3[3][3] = { {1,0,0},{0,1,0},{0,0,1} };
			double MaxDev, Rms; bool Cov;
			VerifyCandidate(A.Positions, B.Positions, BGrid, I3, 1.0, FVector::ZeroVector, FVector::ZeroVector, Tol, MaxDev, Rms, Cov);
			if (Cov)
			{
				Out.bValid = true;
				Out.bIdentity = true;
				Out.MaxDev = MaxDev;
				Out.Rms = Rms;
				Out.CanonToMember = FMatrix::Identity;
				return true;
			}
			return false;
		}

		// --- Rigid mode: identity-with-pivot fast path, then brute-force axis rotations ---
		const double sEst = (S.bAllowUniformScale && A.SurfaceArea > 0)
			? FMath::Sqrt(B.SurfaceArea / A.SurfaceArea) : 1.0;

		struct FBest { double Score = TNumericLimits<double>::Max(); FRot3 R; double s = 1.0; };
		FBest Best;

		// Candidate set.
		TArray<const FRot3*> Candidates;
		static const FRot3 Identity; // R = I, proper
		Candidates.Add(&Identity);
		for (const FRot3& R : AxisRotations())
		{
			if (R.bImproper && !S.bAllowMirrored)
			{
				continue;
			}
			Candidates.Add(&R);
		}

		for (const FRot3* R : Candidates)
		{
			const double Score = ScoreCandidate(A.Positions, BGrid, R->M, sEst, A.Centroid, B.Centroid, Tol, Stride);
			if (Score < Best.Score)
			{
				Best.Score = Score;
				Best.R = *R;
				Best.s = sEst;
			}
		}

		// Full verify the winner.
		double MaxDev, Rms; bool Cov;
		VerifyCandidate(A.Positions, B.Positions, BGrid, Best.R.M, Best.s, A.Centroid, B.Centroid, Tol, MaxDev, Rms, Cov);
		if (!Cov)
		{
			return false;
		}

		Out.bValid = true;
		Out.bMirrored = Best.R.bImproper;
		Out.bScaled = !FMath::IsNearlyEqual(Best.s, 1.0, 1e-4);
		Out.Scale = Best.s;
		Out.MaxDev = MaxDev;
		Out.Rms = Rms;
		Out.CanonToMember = BuildMatrix(Best.R.M, Best.s, A.Centroid, B.Centroid);
		Out.bIdentity = !Out.bMirrored && !Out.bScaled && Out.CanonToMember.Equals(FMatrix::Identity, 1e-4);
		return true;
	}
}