#include "Debug/DescentImageStats.h"

namespace Descent::Image
{
	double Luma01(const FColor& C)
	{
		return (0.2126 * C.R + 0.7152 * C.G + 0.0722 * C.B) / 255.0;
	}

	FDescentLumaStats ComputeStats(TConstArrayView64<FColor> Pixels, int32 Width, int32 Height)
	{
		FDescentLumaStats S;
		S.Width = Width;
		S.Height = Height;
		const int64 N = static_cast<int64>(Width) * Height;
		if (Width <= 0 || Height <= 0 || Pixels.Num() != N)
		{
			return S;
		}

		constexpr int32 Bins = 1021; // luma * 255 * 4, rounded: 0 .. 1020
		TArray<uint32> Hist;
		Hist.SetNumZeroed(Bins);

		const int32 CX0 = FMath::RoundToInt(Width * 0.3), CX1 = FMath::RoundToInt(Width * 0.7);
		const int32 CY0 = FMath::RoundToInt(Height * 0.3), CY1 = FMath::RoundToInt(Height * 0.7);

		double LumaSum = 0.0, CentreSum = 0.0, R = 0.0, G = 0.0, B = 0.0;
		int64 CentreCount = 0, NonBlack = 0;
		for (int32 Y = 0; Y < Height; ++Y)
		{
			const FColor* Row = Pixels.GetData() + static_cast<int64>(Y) * Width;
			const bool bCentreRow = Y >= CY0 && Y < CY1;
			for (int32 X = 0; X < Width; ++X)
			{
				const FColor& C = Row[X];
				const double L = Luma01(C);
				LumaSum += L;
				R += C.R;
				G += C.G;
				B += C.B;
				if (FMath::Max3(C.R, C.G, C.B) > 8)
				{
					++NonBlack;
				}
				if (bCentreRow && X >= CX0 && X < CX1)
				{
					CentreSum += L;
					++CentreCount;
				}
				++Hist[FMath::Clamp(FMath::RoundToInt(L * 255.0 * 4.0), 0, Bins - 1)];
			}
		}

		S.MeanLuma = LumaSum / N;
		S.CentreLuma = CentreCount > 0 ? CentreSum / CentreCount : S.MeanLuma;
		S.NonBlack = static_cast<double>(NonBlack) / N;
		S.MeanR = R / (255.0 * N);
		S.MeanG = G / (255.0 * N);
		S.MeanB = B / (255.0 * N);

		const double Target = 0.99 * N;
		double Cumulative = 0.0;
		for (int32 I = 0; I < Bins; ++I)
		{
			Cumulative += Hist[I];
			if (Cumulative >= Target)
			{
				S.P99Luma = I / (255.0 * 4.0);
				break;
			}
		}
		BlockContrastRect(Pixels, Width, Height, 0.3, 0.3, 0.7, 0.7, S.CentreDetail, S.CentreRms);
		return S;
	}

	void BlockContrastRect(TConstArrayView64<FColor> Pixels, int32 Width, int32 Height, double X0, double Y0, double X1, double Y1, double& OutDetail, double& OutRms)
	{
		OutDetail = 0.0;
		OutRms = 0.0;
		if (Width <= 0 || Height <= 0 || Pixels.Num() != static_cast<int64>(Width) * Height)
		{
			return;
		}
		const int32 Block = FMath::Max(4, Width / 160);
		const int32 PX0 = FMath::Clamp(FMath::FloorToInt(FMath::Min(X0, X1) * Width), 0, Width - 1);
		const int32 PY0 = FMath::Clamp(FMath::FloorToInt(FMath::Min(Y0, Y1) * Height), 0, Height - 1);
		const int32 Nx = (FMath::Clamp(FMath::CeilToInt(FMath::Max(X0, X1) * Width), PX0 + 1, Width) - PX0) / Block;
		const int32 Ny = (FMath::Clamp(FMath::CeilToInt(FMath::Max(Y0, Y1) * Height), PY0 + 1, Height) - PY0) / Block;
		if (Nx < 2 || Ny < 2)
		{
			return;
		}
		TArray<double> Mean;
		Mean.SetNumZeroed(Nx * Ny);
		for (int32 By = 0; By < Ny; ++By)
		{
			for (int32 Bx = 0; Bx < Nx; ++Bx)
			{
				double Sum = 0.0;
				for (int32 Y = 0; Y < Block; ++Y)
				{
					const FColor* Row = Pixels.GetData() + static_cast<int64>(PY0 + By * Block + Y) * Width + PX0 + Bx * Block;
					for (int32 X = 0; X < Block; ++X) { Sum += Luma01(Row[X]); }
				}
				Mean[By * Nx + Bx] = Sum / (Block * Block);
			}
		}
		double All = 0.0, Sq = 0.0, Pairs = 0.0;
		int64 NumPairs = 0;
		for (int32 By = 0; By < Ny; ++By)
		{
			for (int32 Bx = 0; Bx < Nx; ++Bx)
			{
				const double A = Mean[By * Nx + Bx];
				All += A;
				Sq += A * A;
				if (Bx + 1 < Nx) { const double B = Mean[By * Nx + Bx + 1]; Pairs += FMath::Abs(A - B) / (A + B + 0.02); ++NumPairs; }
				if (By + 1 < Ny) { const double B = Mean[(By + 1) * Nx + Bx]; Pairs += FMath::Abs(A - B) / (A + B + 0.02); ++NumPairs; }
			}
		}
		const double N = static_cast<double>(Nx) * Ny;
		const double M = All / N;
		OutDetail = NumPairs > 0 ? Pairs / NumPairs : 0.0;
		OutRms = M > 1e-6 ? FMath::Sqrt(FMath::Max(0.0, Sq / N - M * M)) / M : 0.0;
	}

	double MeanLumaRect(TConstArrayView64<FColor> Pixels, int32 Width, int32 Height, double X0, double Y0, double X1, double Y1)
	{
		if (Width <= 0 || Height <= 0 || Pixels.Num() != static_cast<int64>(Width) * Height)
		{
			return 0.0;
		}
		const int32 PX0 = FMath::Clamp(FMath::FloorToInt(FMath::Min(X0, X1) * Width), 0, Width - 1);
		const int32 PX1 = FMath::Clamp(FMath::CeilToInt(FMath::Max(X0, X1) * Width), PX0 + 1, Width);
		const int32 PY0 = FMath::Clamp(FMath::FloorToInt(FMath::Min(Y0, Y1) * Height), 0, Height - 1);
		const int32 PY1 = FMath::Clamp(FMath::CeilToInt(FMath::Max(Y0, Y1) * Height), PY0 + 1, Height);
		double Sum = 0.0;
		int64 Count = 0;
		for (int32 Y = PY0; Y < PY1; ++Y)
		{
			const FColor* Row = Pixels.GetData() + static_cast<int64>(Y) * Width;
			for (int32 X = PX0; X < PX1; ++X)
			{
				Sum += Luma01(Row[X]);
				++Count;
			}
		}
		return Count > 0 ? Sum / Count : 0.0;
	}

	double MeanLumaRing(TConstArrayView64<FColor> Pixels, int32 Width, int32 Height, double Cx, double Cy, double R0, double R1)
	{
		if (Width <= 0 || Height <= 0 || Pixels.Num() != static_cast<int64>(Width) * Height)
		{
			return 0.0;
		}
		const double Lo = FMath::Min(R0, R1), Hi = FMath::Max(R0, R1);
		const double Lo2 = Lo * Lo, Hi2 = Hi * Hi;
		double Sum = 0.0;
		int64 Count = 0;
		for (int32 Y = 0; Y < Height; ++Y)
		{
			const double Dy = (Y + 0.5) - Cy;
			const FColor* Row = Pixels.GetData() + static_cast<int64>(Y) * Width;
			for (int32 X = 0; X < Width; ++X)
			{
				const double Dx = (X + 0.5) - Cx;
				const double D2 = Dx * Dx + Dy * Dy;
				if (D2 >= Lo2 && D2 < Hi2)
				{
					Sum += Luma01(Row[X]);
					++Count;
				}
			}
		}
		return Count > 0 ? Sum / Count : 0.0;
	}
}