using System;
using System.Collections.Generic;

namespace SquishyTerry.Sim;

/// <summary>
/// Shape-matching cluster (Müller et al. 2005). Linear regions snap rigidly to their best-fit
/// rotation; the quadratic mode (used by the mouth) allows some deformation with beta = 0.8.
/// Original: 712-byte struct, init f_bj, solve f_yb.
/// </summary>
public sealed class ShapeRegion
{
	public bool Enabled;
	public Quat Rotation = Quat.Identity;
	public bool Quadratic;
	public int[] Indices = Array.Empty<int>();
	public Vec3[] RestPositions = Array.Empty<Vec3>();
	public Vec3 RestCentroid;

	// Precomputed at Init (all constant): rest offsets q = RestPositions - RestCentroid, their
	// quadratic terms q~ (9 per vertex) and Aqq^-1 (row-major 9x9). The solve loops below spell out
	// the math by hand (the editor build is unoptimized) in the original evaluation order.
	Vec3[] _restOffset = Array.Empty<Vec3>();
	double[] _restQuad = Array.Empty<double>();
	readonly double[] _aqqInv = new double[81];
	readonly float[] _apq = new float[9];
	readonly float[] _rm = new float[9];
	readonly double[] _apqt = new double[27];
	readonly double[] _a = new double[27];
	static readonly float[] s_extractScratch = new float[9];

	/// <param name="quadratic">Allow quadratic deformation (the mouth) instead of a rigid fit.</param>
	/// <param name="indices">Vertices in the region.</param>
	/// <param name="positions">Rest positions (all vertices).</param>
	public void Init( bool quadratic, IReadOnlyList<int> indices, IReadOnlyList<Vec3> positions )
	{
		Enabled = true;
		Quadratic = quadratic;
		Indices = new int[indices.Count];
		RestPositions = new Vec3[indices.Count];

		float sx = 0, sy = 0, sz = 0;
		for ( int i = 0; i < indices.Count; i++ )
		{
			Indices[i] = indices[i];
			var p = positions[indices[i]];
			RestPositions[i] = p;
			sx += p.X; sy += p.Y; sz += p.Z;
		}
		float n = (uint)indices.Count;
		RestCentroid = new Vec3( sx / n, sy / n, sz / n );
		Rotation = Quat.Identity;

		_restOffset = new Vec3[RestPositions.Length];
		for ( int i = 0; i < RestPositions.Length; i++ )
			_restOffset[i] = RestPositions[i] - RestCentroid;

		if ( !quadratic ) return;

		var q = new double[9];
		_restQuad = new double[RestPositions.Length * 9];
		for ( int i = 0; i < RestPositions.Length; i++ )
		{
			Quadratic9( _restOffset[i], q );
			for ( int k = 0; k < 9; k++ ) _restQuad[i * 9 + k] = q[k];
		}

		var aqq = new double[9, 9];
		foreach ( var rp in RestPositions )
		{
			Quadratic9( rp - RestCentroid, q );
			for ( int r = 0; r < 9; r++ )
				for ( int c = 0; c < 9; c++ )
					aqq[r, c] += q[r] * q[c];
		}
		var inv = new double[9, 9];
		Invert9( aqq, inv );
		for ( int r = 0; r < 9; r++ )
			for ( int c = 0; c < 9; c++ )
				_aqqInv[r * 9 + c] = inv[r, c];
	}

	static void Quadratic9( Vec3 v, Span<double> q )
	{
		q[0] = v.X; q[1] = v.Y; q[2] = v.Z;
		q[3] = v.X * v.X; q[4] = v.Y * v.Y; q[5] = v.Z * v.Z;
		q[6] = v.Y * v.X; q[7] = v.Z * v.Y; q[8] = v.Z * v.X;
	}

	/// <summary>Project the region's vertices in <paramref name="p"/> onto the matched shape.</summary>
	public void Solve( Vec3[] p )
	{
		int[] idx = Indices;
		Vec3[] rel = _restOffset;
		int count = idx.Length;
		float cx = 0, cy = 0, cz = 0;
		for ( int i = 0; i < count; i++ )
		{
			ref Vec3 v = ref p[idx[i]];
			cx += v.X; cy += v.Y; cz += v.Z;
		}
		float n = (uint)count;
		cx /= n; cy /= n; cz /= n;

		// Apq = sum (p - c)(q - q0)^T, column-major.
		float a0 = 0, a1 = 0, a2 = 0, a3 = 0, a4 = 0, a5 = 0, a6 = 0, a7 = 0, a8 = 0;
		for ( int i = 0; i < count; i++ )
		{
			ref Vec3 q = ref rel[i];
			ref Vec3 v = ref p[idx[i]];
			float dx = v.X - cx, dy = v.Y - cy, dz = v.Z - cz;
			a0 += q.X * dx; a1 += q.X * dy; a2 += q.X * dz;
			a3 += q.Y * dx; a4 += q.Y * dy; a5 += q.Y * dz;
			a6 += q.Z * dx; a7 += q.Z * dy; a8 += q.Z * dz;
		}
		var apq = _apq;
		apq[0] = a0; apq[1] = a1; apq[2] = a2;
		apq[3] = a3; apq[4] = a4; apq[5] = a5;
		apq[6] = a6; apq[7] = a7; apq[8] = a8;
		ExtractRotation( apq, ref Rotation, 30 );

		if ( !Quadratic )
		{
			// p = c + Rotation.Rotate( q )
			float rx = Rotation.X, ry = Rotation.Y, rz = Rotation.Z, rw = Rotation.W;
			for ( int i = 0; i < count; i++ )
			{
				ref Vec3 q = ref rel[i];
				float tx = 2f * (ry * q.Z - rz * q.Y);
				float ty = 2f * (rz * q.X - rx * q.Z);
				float tz = 2f * (rx * q.Y - ry * q.X);
				ref Vec3 o = ref p[idx[i]];
				o.X = cx + (q.X + rw * tx + (ry * tz - rz * ty));
				o.Y = cy + (q.Y + rw * ty + (rz * tx - rx * tz));
				o.Z = cz + (q.Z + rw * tz + (rx * ty - ry * tx));
			}
			return;
		}

		// A~ = (sum (p - c) q~^T) * Aqq^-1 ; A = 0.8 A~ + 0.2 [R | 0]   (row-major 3x9)
		var apqt = _apqt;
		Array.Clear( apqt );
		var qt = _restQuad;
		for ( int i = 0; i < count; i++ )
		{
			ref Vec3 v = ref p[idx[i]];
			double dx = v.X - cx, dy = v.Y - cy, dz = v.Z - cz;
			int o = i * 9;
			for ( int k = 0; k < 9; k++ )
			{
				double qk = qt[o + k];
				apqt[k] += dx * qk;
				apqt[9 + k] += dy * qk;
				apqt[18 + k] += dz * qk;
			}
		}

		var rm = _rm;
		Rotation.ToMatrix( rm );

		var a = _a;
		var aqqInv = _aqqInv;
		for ( int r = 0; r < 3; r++ )
		{
			int ro = r * 9;
			for ( int k = 0; k < 9; k++ )
			{
				double s = 0;
				for ( int m = 0; m < 9; m++ ) s += apqt[ro + m] * aqqInv[m * 9 + k];
				double rot = k < 3 ? rm[k * 3 + r] : 0.0;
				a[ro + k] = 0.8 * s + 0.2 * rot;
			}
		}

		for ( int i = 0; i < count; i++ )
		{
			int o = i * 9;
			double gx = 0, gy = 0, gz = 0;
			for ( int k = 0; k < 9; k++ )
			{
				double qk = qt[o + k];
				gx += a[k] * qk;
				gy += a[9 + k] * qk;
				gz += a[18 + k] * qk;
			}
			ref Vec3 v = ref p[idx[i]];
			v.X = cx + (float)gx;
			v.Y = cy + (float)gy;
			v.Z = cz + (float)gz;
		}
	}

	/// <summary>
	/// Müller et al. 2016, "A Robust Method to Extract the Rotational Part of Deformations" (f_tsc).
	/// </summary>
	public static void ExtractRotation( ReadOnlySpan<float> a, ref Quat q, int maxIter )
	{
		var s = s_extractScratch;
		for ( int i = 0; i < 9; i++ ) s[i] = a[i];
		ExtractRotation( s, ref q, maxIter );
	}

	/// <summary>
	/// <see cref="ExtractRotation(ReadOnlySpan{float}, ref Quat, int)"/> on a column-major float[9],
	/// with the quaternion/matrix math written out (same arithmetic order as Quat.ToMatrix,
	/// Quat.FromAxisAngle, Quat operator * and Quat.Normalized).
	/// </summary>
	public static void ExtractRotation( float[] a, ref Quat q, int maxIter )
	{
		float a0 = a[0], a1 = a[1], a2 = a[2], a3 = a[3], a4 = a[4], a5 = a[5], a6 = a[6], a7 = a[7], a8 = a[8];
		float qx = q.X, qy = q.Y, qz = q.Z, qw = q.W;
		for ( int iter = 0; iter < maxIter; iter++ )
		{
			// Quat.ToMatrix (column-major)
			double x = qx, y = qy, z = qz, w = qw;
			float r0 = (float)(1 - 2 * y * y - 2 * z * z);
			float r1 = (float)(2 * x * y + 2 * w * z);
			float r2 = (float)(2 * x * z - 2 * w * y);
			float r3 = (float)(2 * x * y - 2 * w * z);
			float r4 = (float)(1 - 2 * x * x - 2 * z * z);
			float r5 = (float)(2 * y * z + 2 * w * x);
			float r6 = (float)(2 * x * z + 2 * w * y);
			float r7 = (float)(2 * y * z - 2 * w * x);
			float r8 = (float)(1 - 2 * x * x - 2 * y * y);

			// omega = sum_col cross( R col, A col ), dot = sum_col dot( R col, A col )
			float ox = 0, oy = 0, oz = 0, dot = 0;
			ox += r1 * a2 - r2 * a1; oy += r2 * a0 - r0 * a2; oz += r0 * a1 - r1 * a0;
			dot += r0 * a0 + r1 * a1 + r2 * a2;
			ox += r4 * a5 - r5 * a4; oy += r5 * a3 - r3 * a5; oz += r3 * a4 - r4 * a3;
			dot += r3 * a3 + r4 * a4 + r5 * a5;
			ox += r7 * a8 - r8 * a7; oy += r8 * a6 - r6 * a8; oz += r6 * a7 - r7 * a6;
			dot += r6 * a6 + r7 * a7 + r8 * a8;

			float denom = (float)(MathF.Abs( dot ) + 1e-9);
			ox /= denom; oy /= denom; oz /= denom;
			float len = MathF.Sqrt( ox * ox + oy * oy + oz * oz );
			if ( len < 1e-9 ) break;

			// rq = Quat.FromAxisAngle( omega / len, len )
			float half = len * 0.5f;
			float sn = MathF.Sin( half );
			float bx = ox / len * sn, by = oy / len * sn, bz = oz / len * sn, bw = MathF.Cos( half );

			// q = (rq * q).Normalized
			float nx = bw * qx + bx * qw + by * qz - bz * qy;
			float ny = bw * qy - bx * qz + by * qw + bz * qx;
			float nz = bw * qz + bx * qy - by * qx + bz * qw;
			float nw = bw * qw - bx * qx - by * qy - bz * qz;
			float inv = 1.0f / MathF.Sqrt( nx * nx + ny * ny + nz * nz + nw * nw );
			qx = nx * inv; qy = ny * inv; qz = nz * inv; qw = nw * inv;
		}
		q = new Quat( qx, qy, qz, qw );
	}

	/// <summary>9x9 inverse with partial pivoting (the original used an Eigen PartialPivLU).</summary>
	static void Invert9( double[,] m, double[,] inv )
	{
		const int n = 9;
		var a = new double[n, n];
		for ( int r = 0; r < n; r++ )
			for ( int c = 0; c < n; c++ )
				a[r, c] = m[r, c];
		for ( int r = 0; r < n; r++ )
			for ( int c = 0; c < n; c++ )
				inv[r, c] = r == c ? 1 : 0;

		for ( int col = 0; col < n; col++ )
		{
			int pivot = col;
			double best = Math.Abs( a[col, col] );
			for ( int r = col + 1; r < n; r++ )
			{
				double v = Math.Abs( a[r, col] );
				if ( v > best ) { best = v; pivot = r; }
			}
			if ( pivot != col )
			{
				for ( int c = 0; c < n; c++ )
				{
					(a[col, c], a[pivot, c]) = (a[pivot, c], a[col, c]);
					(inv[col, c], inv[pivot, c]) = (inv[pivot, c], inv[col, c]);
				}
			}
			double d = a[col, col];
			for ( int c = 0; c < n; c++ ) { a[col, c] /= d; inv[col, c] /= d; }
			for ( int r = 0; r < n; r++ )
			{
				if ( r == col ) continue;
				double f = a[r, col];
				if ( f == 0 ) continue;
				for ( int c = 0; c < n; c++ )
				{
					a[r, c] -= f * a[col, c];
					inv[r, c] -= f * inv[col, c];
				}
			}
		}
	}
}