using System;

namespace CharacterEditor;
/// <summary>Fits a shared identity to observed surface landmarks using calibrated native controls.</summary>
public static class PhotoSurfaceFit
{
	public sealed record Diagnostics(float[][] Poses, float[] Perspective, float Error);
	// The deployed analyzer returns 25 landmarks. Keep that established path
	// unchanged while the denser analyzer and surface fitter are being evaluated.
	public static bool CanFit(IReadOnlyList<FaceMeasurement> views) => views.Any(v => v.FrontUsable && PhotoSurfaceBasis.Ids.Count(v.Landmarks.ContainsKey) >= 96);
	public static Vector3[] Geometry(PhotoFace face)
	{
		var points = new Vector3[PhotoSurfaceBasis.Ids.Length];
		for (int p = 0; p < points.Length; p++)
		{
			var overall = Vector3.Zero;
			for (int h = 0; h < 7; h++)
				overall += PhotoSurfaceBasis.Heads[h][p] * face.HeadWeights[h];
			points[p] = overall;
			for (int r = 0; r < FaceRegions.Ids.Length; r++)
			{
				if (!face.RegionWeights.TryGetValue(FaceRegions.Ids[r], out var weights))
					continue;
				var regional = Vector3.Zero;
				for (int h = 0; h < 7; h++)
					regional += PhotoSurfaceBasis.Heads[h][p] * weights[h];
				points[p] += (regional - overall) * PhotoSurfaceBasis.Masks[p][r];
			}
		}

		for (int c = 0; c < PhotoSurfaceBasis.Controls.Length; c++)
		{
			var control = PhotoSurfaceBasis.Controls[c];
			float offset = face.Offsets.GetValueOrDefault(control);
			if (MathF.Abs(offset) < .000001f)
				continue;
			var weights = face.WeightsFor(FaceRegions.ForSlider(control));
			for (int h = 0; h < 7; h++)
			{
				var delta = PhotoSurfaceBasis.Deltas[h * PhotoSurfaceBasis.Controls.Length * 2 + c * 2 + (offset >= 0 ? 1 : 0)];
				float amount = weights[h] * MathF.Abs(offset) * 2;
				for (int p = 0; p < points.Length; p++)
					points[p] += delta[p] * amount;
			}
		}

		for (int c = 0; c < PhotoFaceCorrections.Controls.Length; c++)
		{
			var control = PhotoFaceCorrections.Controls[c];
			float amount = face.Corrections.GetValueOrDefault(control.Id);
			if (control.Fine)
				continue;
			if (MathF.Abs(amount) < .000001f)
				continue;
			var weights = face.WeightsFor(control.Region);
			for (int h = 0; h < 7; h++)
			{
				var delta = PhotoCorrectionBasis.Deltas[h * PhotoFaceCorrections.Controls.Length * 2 + c * 2 + (amount >= 0 ? 1 : 0)];
				for (int p = 0; p < points.Length; p++)
					points[p] += delta[p] * weights[h] * MathF.Abs(amount);
			}
		}

		PhotoContourWarp.Apply(points, face);
		return points;
	}

	public static Vector2[] Project(Vector3[] points, float[] pose, float perspective = 0)
	{
		var m = pose ?? new float[]
		{
			1,
			0,
			0,
			0,
			0,
			1,
			0,
			0,
			0,
			0,
			1,
			0,
			0,
			0,
			0,
			1
		};
		var result = new Vector2[points.Length];
		// A camera in front of the human (+X, looking -X) has image-right +Y.
		// This convention must also be used when binding detector points to triangles.
		float meanDepth = points.Average(p => p.y * m[2] + (p.z - 63) * m[6] + p.x * m[10]);
		for (int i = 0; i < points.Length; i++)
		{
			var p = points[i];
			float depth = p.y * m[2] + (p.z - 63) * m[6] + p.x * m[10];
			float divisor = Math.Clamp(1 - perspective * (depth - meanDepth), .5f, 1.5f);
			result[i] = new Vector2(p.y * m[0] + (p.z - 63) * m[4] + p.x * m[8], -(p.y * m[1] + (p.z - 63) * m[5] + p.x * m[9])) / divisor;
		}

		return result;
	}

	sealed class View
	{
		public FaceMeasurement Measurement;
		public int[] Indices;
		public Vector2[] Observed;
		public float[] Weights;
		public float Scale;
		public float Perspective;
		public float YawCorrection;
		public float PitchCorrection;
		public int[][] Features;
	}

	static float CorrectionPenalty(PhotoFaceCorrections.Control control) => control.Fine ? .0000003f : .00002f;
	static View Prepare(FaceMeasurement measurement)
	{
		// At a strong turn the detector invents the hidden eye and cheek. Only the
		// visible central profile constrains depth; those invented points cannot
		// become a second set of widths or force the far eyelids to move.
		int[] profileIds =
		{
			10,
			152,
			168,
			6,
			197,
			195,
			5,
			4,
			1,
			19,
			94,
			2,
			0,
			13,
			14,
			17, 151, 9, 8,
			// Visible eyebrow contours constrain ridge depth at strong turns.
			// Hidden-side points are still downweighted below.
			70, 63, 105, 66, 107, 55, 65, 52, 53, 46,
			300, 293, 334, 296, 336, 285, 295, 282, 283, 276
		};
		var indices = Enumerable.Range(0, PhotoSurfaceBasis.Ids.Length).Where(i => measurement.Landmarks.ContainsKey(PhotoSurfaceBasis.Ids[i]) && (measurement.Yaw < 60 || profileIds.Contains(PhotoSurfaceBasis.Ids[i]))).ToArray();
		var observed = indices.Select(i => measurement.Landmarks[PhotoSurfaceBasis.Ids[i]]).ToArray();
		var groups = indices.GroupBy(i => PhotoSurfaceBasis.Regions[i]).ToDictionary(x => x.Key, x => x.Count());
		var weights = indices.Select(i =>
		{
			string group = PhotoSurfaceBasis.Regions[i];
			float weight = group == "outline" ? (measurement.Yaw < 20 ? .75f : .25f) : group == "cheeks" ? .75f : 1;
			float side = -PhotoSurfaceBasis.Heads[2][i].y * (measurement.Pose?[8] ?? 0);
			if (measurement.Yaw > 25 && side < -.4f)
				weight *= measurement.Yaw > 40 ? .05f : .15f;
			return weight / groups[group];
		}).ToArray();
		float scale = measurement.Landmarks.TryGetValue(234, out var l) && measurement.Landmarks.TryGetValue(454, out var r) ? (l - r).Length : 0;
		scale = Math.Max(scale, (observed.Max(x => x.y) - observed.Min(x => x.y)) * .55f);
		return new()
		{
			Measurement = measurement,
			Indices = indices,
			Observed = observed,
			Weights = weights,
			Features = measurement.Yaw > 30 ? Array.Empty<int[]>() : new[]
			{
				Enumerable.Range(0, indices.Length).Where(k => indices[k] < 16).ToArray(),
				Enumerable.Range(0, indices.Length).Where(k => indices[k] >= 16 && indices[k] < 32).ToArray(),
				Enumerable.Range(0, indices.Length).Where(k => PhotoSurfaceBasis.Regions[indices[k]] == "mouth").ToArray(),
				Enumerable.Range(0, indices.Length).Where(k => PhotoSurfaceBasis.Regions[indices[k]] == "nose").ToArray(),
				Enumerable.Range(0, indices.Length).Where(k => PhotoSurfaceBasis.Regions[indices[k]] == "jaw").ToArray(),
				Enumerable.Range(0, indices.Length).Where(k => PhotoSurfaceBasis.Regions[indices[k]] == "forehead").ToArray(),
				Enumerable.Range(0, indices.Length).Where(k => PhotoSurfaceBasis.Regions[indices[k]] == "cheeks").ToArray()
			}.Where(f => f.Length >= 6).ToArray(),
			Scale = Math.Max(scale, .0001f)
		};
	}

	static float[] CorrectedPose(View view)
	{
		var original = view.Measurement.Pose ?? new float[]
		{
			1,
			0,
			0,
			0,
			0,
			1,
			0,
			0,
			0,
			0,
			1,
			0,
			0,
			0,
			0,
			1
		};
		var result = original.ToArray();
		float y = view.YawCorrection * MathF.PI / 180, p = view.PitchCorrection * MathF.PI / 180, cy = MathF.Cos(y), sy = MathF.Sin(y), cp = MathF.Cos(p), sp = MathF.Sin(p);
		for (int row = 0; row < 3; row++)
		{
			float x = original[row] * cy - original[8 + row] * sy, z = original[row] * sy + original[8 + row] * cy;
			result[row] = x;
			result[4 + row] = original[4 + row] * cp + z * sp;
			result[8 + row] = -original[4 + row] * sp + z * cp;
		}

		return result;
	}

	static void DetectorBias(Vector2[] projection, PhotoFace face, float[] pose)
	{
		float yaw = MathF.Atan2(pose[8], pose[10]) * 180 / MathF.PI;
		int lo = yaw < 0 ? 0 : 1, hi = lo + 1;
		float t = Math.Clamp((yaw - PhotoDetectorBasis.Yaws[lo]) / (PhotoDetectorBasis.Yaws[hi] - PhotoDetectorBasis.Yaws[lo]), 0, 1);
		for (int p = 0; p < projection.Length; p++)
			for (int h = 0; h < 7; h++)
			{
				float weight = face.HeadWeights[h];
				for (int r = 0; r < FaceRegions.Ids.Length; r++)
					if (face.RegionWeights.TryGetValue(FaceRegions.Ids[r], out var w))
						weight += (w[h] - face.HeadWeights[h]) * PhotoSurfaceBasis.Masks[p][r];
				projection[p] += (PhotoDetectorBasis.Bias[h][lo][p] * (1 - t) + PhotoDetectorBasis.Bias[h][hi][p] * t) * weight;
			}

		PhotoContourWarp.DetectorBias(projection, face, yaw);
	}

	static float[] Residual(PhotoFace face, Vector3[] geometry, View view, bool localContours = false)
	{
		var pose = CorrectedPose(view);
		var projection = Project(geometry, pose, view.Perspective);
		if (view.Measurement.DetectorLandmarks)
			DetectorBias(projection, face, pose);
		float sum = view.Weights.Sum();
		var pc = Vector2.Zero;
		var oc = Vector2.Zero;
		for (int k = 0; k < view.Indices.Length; k++)
		{
			pc += projection[view.Indices[k]] * view.Weights[k];
			oc += view.Observed[k] * view.Weights[k];
		}

		pc /= sum;
		oc /= sum;
		float dot = 0, cross = 0, length = 0;
		for (int k = 0; k < view.Indices.Length; k++)
		{
			var p = projection[view.Indices[k]] - pc;
			var o = view.Observed[k] - oc;
			float w = view.Weights[k];
			dot += Vector2.Dot(p, o) * w;
			cross += (p.x * o.y - p.y * o.x) * w;
			length += p.LengthSquared * w;
		}

		float a = dot / Math.Max(length, .000001f), b = cross / Math.Max(length, .000001f);
		var residual = new float[view.Indices.Length * 2];
		var fitted = localContours ? new Vector2[view.Indices.Length] : null;
		for (int k = 0; k < view.Indices.Length; k++)
		{
			var p = projection[view.Indices[k]] - pc;
			var predicted = new Vector2(a * p.x - b * p.y, b * p.x + a * p.y) + oc;
			if (localContours)
				fitted[k] = predicted;
			var delta = (predicted - view.Observed[k]) / view.Scale;
			float d = delta.Length;
			// A mislabeled/occluded feature must not drag the whole head toward it.
			float robust = d < .035f ? 1 : (.07f * d - .001225f) / (d * d);
			float weight = MathF.Sqrt(view.Weights[k] / sum * robust);
			residual[k * 2] = delta.x * weight;
			residual[k * 2 + 1] = delta.y * weight;
		}

		if (!localContours)
			return residual;
		var values = new List<float>(residual);
		foreach (var feature in view.Features)
		{
			// Local centering separates contour shape from feature placement. Eye
			// aperture is measured relative to the eye, rather than the whole face,
			// so a meaningful eyelid mismatch cannot disappear in a low global RMS.
			var predictedCenter = feature.Aggregate(Vector2.Zero, (s, k) => s + fitted[k]) / feature.Length;
			var observedCenter = feature.Aggregate(Vector2.Zero, (s, k) => s + view.Observed[k]) / feature.Length;
			float extent = Math.Max(feature.Max(k => view.Observed[k].x) - feature.Min(k => view.Observed[k].x), feature.Max(k => view.Observed[k].y) - feature.Min(k => view.Observed[k].y));
			float weight = MathF.Sqrt(.04f / feature.Length);
			foreach (int k in feature)
			{
				var delta = (fitted[k] - predictedCenter - view.Observed[k] + observedCenter) / Math.Max(extent, .0001f);
				float distance = delta.Length;
				float robust = distance <= .12f ? 1 : MathF.Sqrt((.24f * distance - .0144f) / (distance * distance));
				values.Add(delta.x * weight * robust);
				values.Add(delta.y * weight * robust);
			}
		}

		return values.ToArray();
	}

	static float Error(PhotoFace face, Vector3[] geometry, View view, bool localContours = false) => Residual(face, geometry, view, localContours).Sum(r => r * r);
	public static float Error(PhotoFace face, IReadOnlyList<FaceMeasurement> views)
	{
		var usable = views.Where(v => PhotoSurfaceBasis.Ids.Count(v.Landmarks.ContainsKey) >= 12).Select(Prepare).ToArray();
		if (usable.Length == 0)
			return 0;
		var geometry = Geometry(face);
		return usable.Average(v => Error(face, geometry, v));
	}

	public static PhotoFace Generate(string name, IReadOnlyList<FaceMeasurement> measurements, float[] prior, System.Threading.CancellationToken cancellation = default, Action<Diagnostics> diagnostics = null, bool fitContours = true)
	{
		var views = measurements.Where(v => PhotoSurfaceBasis.Ids.Count(v.Landmarks.ContainsKey) >= 12).Select(Prepare).ToArray();
		var face = new PhotoFace
		{
			Name = name,
			HeadWeights = prior.ToArray()
		};
		bool depth = measurements.Any(v => v.ProfileUsable && v.Yaw >= 25);
		bool localContours = false;
		// The former whole-head volume reaches the lower neck. Identity proportions
		// use face-only width/height corrections so photo fitting cannot move the collar.
		var controls = PhotoSurfaceBasis.Controls.Where(c => c != "head" && (depth || c is not ("nose" or "chin" or "mouthprojection"))).ToArray();
		float Score()
		{
			cancellation.ThrowIfCancellationRequested();
			var geometry = Geometry(face);
			float score = views.Sum(v => Error(face, geometry, v, localContours) * v.Measurement.Quality) / views.Sum(v => v.Measurement.Quality);
			score += views.Sum(v => (v.YawCorrection * v.YawCorrection + v.PitchCorrection * v.PitchCorrection) * .00000001f);
			score += .000015f * face.Offsets.Values.Sum(x => x * x);
			score += PhotoFaceCorrections.Controls.Sum(c => CorrectionPenalty(c) * MathF.Pow(face.Corrections.GetValueOrDefault(c.Id), 2));
			score += .000008f * face.HeadWeights.Zip(prior).Sum(x => MathF.Pow(x.First - x.Second, 2));
			foreach (var w in face.RegionWeights.Values)
				score += .000006f * w.Zip(face.HeadWeights).Sum(x => MathF.Pow(x.First - x.Second, 2));
			return score;
		}

		float best = Score();
		void Blend(float[] weights, float step)
		{
			for (int from = 0; from < 7; from++)
				for (int to = 0; to < 7; to++)
				{
					if (from == to || !PhotoFaceFit.SameFamily(from, to) || weights[from] < step)
						continue;
					weights[from] -= step;
					weights[to] += step;
					float score = Score();
					if (score < best - 1e-10f)
						best = score;
					else
					{
						weights[from] += step;
						weights[to] -= step;
					}
				}
		}

		void Controls(float step)
		{
			foreach (var id in controls)
			{
				float original = face.Offsets.GetValueOrDefault(id), chosen = original;
				foreach (float sign in new[]
				{
					-1f,
					1f
				}

				)
				{
					float value = Math.Clamp(original + step * sign, -.35f, .35f);
					face.Offsets[id] = value;
					float score = Score();
					if (score < best - 1e-10f)
					{
						best = score;
						chosen = value;
					}
				}

				face.Offsets[id] = chosen;
			}
		}

		void Corrections(float step)
		{
			foreach (var control in PhotoFaceCorrections.Controls.Where(c => !c.Fine))
			{
				float original = face.Corrections.GetValueOrDefault(control.Id), chosen = original;
				foreach (float direction in new[]
				{
					-1f,
					1f
				}

				)
				{
					float value = Math.Clamp(original + direction * step, -1, 1);
					face.Corrections[control.Id] = value;
					float score = Score();
					if (score < best - 1e-10f)
					{
						best = score;
						chosen = value;
					}
				}

				face.Corrections[control.Id] = chosen;
			}
		}

		void Cameras(float step)
		{
			foreach (var view in views)
				foreach (bool yaw in new[]
				{
					true,
					false
				}

				)
				{
					float original = yaw ? view.YawCorrection : view.PitchCorrection, chosen = original;
					foreach (float direction in new[]
					{
						-1f,
						1f
					}

					)
					{
						float value = Math.Clamp(original + direction * step, -20, 20);
						if (yaw)
							view.YawCorrection = value;
						else
							view.PitchCorrection = value;
						float score = Score();
						if (score < best - 1e-10f)
						{
							best = score;
							chosen = value;
						}
					}

					if (yaw)
						view.YawCorrection = chosen;
					else
						view.PitchCorrection = chosen;
				}
		}

		foreach (float step in new[]
		{
			8f,
			4f,
			2f,
			.75f
		}

		)
			for (int pass = 0; pass < 2; pass++)
				Cameras(step);
		foreach (float step in new[]
		{
			.2f,
			.08f,
			.025f,
			.008f
		}

		)
			for (int pass = 0; pass < 3; pass++)
			{
				Blend(face.HeadWeights, step);
				Controls(step);
				Corrections(step * 2);
				Cameras(step * 15);
			}

		foreach (var region in new[]
		{
			"nose",
			"jaw",
			"eyes",
			"mouth",
			"forehead",
			"cheeks"
		}

		)
			face.RegionWeights[region] = face.HeadWeights.ToArray();
		foreach (float step in new[]
		{
			.15f,
			.05f,
			.015f,
			.005f
		}

		)
			for (int pass = 0; pass < 2; pass++)
			{
				foreach (var weights in face.RegionWeights.Values)
					Blend(weights, step);
				Controls(step);
				Corrections(step * 2);
				Cameras(step * 15);
			}

		// Camera scale, translation and roll are solved analytically for every view.
		// A small perspective search accounts for differing photo focal distances.
		foreach (var view in views)
		{
			float chosen = view.Perspective;
			foreach (float value in new[]
			{
				0f,
				.015f,
				.03f,
				.05f
			}

			)
			{
				view.Perspective = value;
				float score = Score();
				if (score < best)
				{
					best = score;
					chosen = value;
				}
			}

			view.Perspective = chosen;
		}

		foreach (float step in new[]
		{
			.015f,
			.005f
		}

		)
		{
			Blend(face.HeadWeights, step);
			foreach (var weights in face.RegionWeights.Values)
				Blend(weights, step);
			Controls(step);
			Corrections(step * 2);
			Cameras(step * 15);
		}

		// Single-parameter steps can stall when eye spacing, head proportions and
		// framing need to change together. Refine those variables jointly.
		var parameters = new List<PhotoLeastSquares.Parameter>();
		foreach (var id in controls)
			parameters.Add(new(() => face.Offsets.GetValueOrDefault(id), v => face.Offsets[id] = v, -.35f, .35f, .002f, .15f));
		foreach (var control in PhotoFaceCorrections.Controls.Where(c => !c.Fine))
			parameters.Add(new(() => face.Corrections.GetValueOrDefault(control.Id), v => face.Corrections[control.Id] = v, -1, 1, .004f, .3f));
		foreach (var view in views)
		{
			parameters.Add(new(() => view.YawCorrection, v => view.YawCorrection = v, -20, 20, .1f, 4));
			parameters.Add(new(() => view.PitchCorrection, v => view.PitchCorrection = v, -20, 20, .1f, 4));
			parameters.Add(new(() => view.Perspective, v => view.Perspective = v, 0, .05f, .001f, .01f));
		}

		float[] Residuals()
		{
			cancellation.ThrowIfCancellationRequested();
			var geometry = Geometry(face);
			var values = new List<float>();
			float quality = views.Sum(v => v.Measurement.Quality);
			foreach (var view in views)
			{
				float w = MathF.Sqrt(view.Measurement.Quality / quality);
				values.AddRange(Residual(face, geometry, view, localContours).Select(x => x * w));
			}

			values.AddRange(controls.Select(id => face.Offsets.GetValueOrDefault(id) * MathF.Sqrt(.000015f)));
			values.AddRange(PhotoFaceCorrections.Controls.Select(c => face.Corrections.GetValueOrDefault(c.Id) * MathF.Sqrt(CorrectionPenalty(c))));
			foreach (var view in views)
			{
				values.Add(view.YawCorrection * .0001f);
				values.Add(view.PitchCorrection * .0001f);
			}

			return values.ToArray();
		}

		PhotoLeastSquares.Refine(parameters, Residuals, Score, cancellation);
		// Local contours refine a settled identity and camera. Keeping camera pose
		// fixed here prevents an eyelid or jaw mismatch being hidden by a new angle.
		if (fitContours)
		{
			localContours = true;
			var contours = new List<PhotoLeastSquares.Parameter>();
			foreach (var control in PhotoFaceCorrections.Controls.Where(c => c.Fine && (!c.Depth || depth)))
				contours.Add(new(() => face.Corrections.GetValueOrDefault(control.Id), v => face.Corrections[control.Id] = v, -1, 1, .004f, .35f));
			PhotoLeastSquares.Refine(contours, Residuals, Score, cancellation, 24);
			// Existing sliders can relax as the finer fields take over their contour.
			foreach (var id in controls)
				contours.Add(new(() => face.Offsets.GetValueOrDefault(id), v => face.Offsets[id] = v, -.35f, .35f, .002f, .1f));
			PhotoLeastSquares.Refine(contours, Residuals, Score, cancellation, 12);
		}

		best = Score();
		foreach (var key in face.Offsets.Keys.ToArray())
			if (MathF.Abs(face.Offsets[key]) < .00001f)
				face.Offsets.Remove(key);
		// Remove floating-point drift so endpoints remain valid portable coefficients.
		foreach (var weights in face.RegionWeights.Values.Prepend(face.HeadWeights))
		{
			for (int h = 0; h < 7; h++)
				weights[h] = Math.Clamp(weights[h], 0, 1);
			float total = weights.Sum();
			for (int h = 0; h < 7; h++)
				weights[h] /= total;
		}

		diagnostics?.Invoke(new(views.Select(CorrectedPose).ToArray(), views.Select(v => v.Perspective).ToArray(), Score()));
		return face;
	}
}