using System;

namespace CharacterEditor;
/// <summary>Small damped joint refinement for coupled face proportions and camera pose.</summary>
internal static class PhotoLeastSquares
{
	internal sealed record Parameter(Func<float> Read, Action<float> Write, float Min, float Max, float Difference, float MaxStep);
	internal static void Refine(IReadOnlyList<Parameter> parameters, Func<float[]> residuals, Func<float> score, System.Threading.CancellationToken cancellation, int iterations = 12)
	{
		int count = parameters.Count;
		double damping = .001;
		float best = score();
		for (int iteration = 0; iteration < iterations; iteration++)
		{
			cancellation.ThrowIfCancellationRequested();
			var original = parameters.Select(p => p.Read()).ToArray();
			var residual = residuals();
			var jacobian = new double[count][];
			for (int p = 0; p < count; p++)
			{
				cancellation.ThrowIfCancellationRequested();
				var parameter = parameters[p];
				float value = Math.Clamp(original[p] + parameter.Difference, parameter.Min, parameter.Max);
				if (value == original[p])
					value = Math.Clamp(original[p] - parameter.Difference, parameter.Min, parameter.Max);
				parameter.Write(value);
				var changed = residuals();
				parameter.Write(original[p]);
				jacobian[p] = new double[residual.Length];
				double step = value - original[p];
				if (Math.Abs(step) > 1e-9)
					for (int r = 0; r < residual.Length; r++)
						jacobian[p][r] = (changed[r] - residual[r]) / step;
			}

			var matrix = new double[count, count];
			var rhs = new double[count];
			for (int p = 0; p < count; p++)
			{
				for (int r = 0; r < residual.Length; r++)
					rhs[p] -= jacobian[p][r] * residual[r];
				for (int q = 0; q <= p; q++)
				{
					double dot = 0;
					for (int r = 0; r < residual.Length; r++)
						dot += jacobian[p][r] * jacobian[q][r];
					matrix[p, q] = matrix[q, p] = dot;
				}

				matrix[p, p] += damping * Math.Max(matrix[p, p], 1e-7);
			}

			if (!Solve(matrix, rhs))
				break;
			bool accepted = false;
			for (int trial = 0; trial < 5; trial++)
			{
				float scale = MathF.Pow(.5f, trial);
				for (int p = 0; p < count; p++)
				{
					var parameter = parameters[p];
					float step = (float)Math.Clamp(rhs[p], -parameter.MaxStep, parameter.MaxStep) * scale;
					parameter.Write(Math.Clamp(original[p] + step, parameter.Min, parameter.Max));
				}

				float candidate = score();
				if (float.IsFinite(candidate) && candidate < best - 1e-10f)
				{
					best = candidate;
					accepted = true;
					break;
				}
			}

			if (accepted)
				damping = Math.Max(1e-6, damping * .4);
			else
			{
				for (int p = 0; p < count; p++)
					parameters[p].Write(original[p]);
				damping *= 10;
				if (damping > 1000)
					break;
			}
		}
	}

	static bool Solve(double[, ] matrix, double[] rhs)
	{
		int count = rhs.Length;
		for (int column = 0; column < count; column++)
		{
			int pivot = column;
			for (int row = column + 1; row < count; row++)
				if (Math.Abs(matrix[row, column]) > Math.Abs(matrix[pivot, column]))
					pivot = row;
			if (Math.Abs(matrix[pivot, column]) < 1e-12)
				return false;
			if (pivot != column)
			{
				for (int j = column; j < count; j++)
					(matrix[column, j], matrix[pivot, j]) = (matrix[pivot, j], matrix[column, j]);
				(rhs[column], rhs[pivot]) = (rhs[pivot], rhs[column]);
			}

			for (int row = column + 1; row < count; row++)
			{
				double factor = matrix[row, column] / matrix[column, column];
				for (int j = column; j < count; j++)
					matrix[row, j] -= factor * matrix[column, j];
				rhs[row] -= factor * rhs[column];
			}
		}

		for (int row = count - 1; row >= 0; row--)
		{
			for (int j = row + 1; j < count; j++)
				rhs[row] -= matrix[row, j] * rhs[j];
			rhs[row] /= matrix[row, row];
		}

		return rhs.All(double.IsFinite);
	}
}