Sim/Face.cs
using System;
using System.Collections.Generic;

namespace SquishyTerry.Sim;

/// <summary>
/// The squishy face soft body, reconstructed from the original C++ (wasm) build.
///
/// A unit icosphere (4 subdivisions) is carved into a face: a mouth dent, flattened eye/nose/ear
/// patches and a hair shell. Each step runs position-based dynamics with a BDF2 predictor:
/// pre-tensioned distance constraints, dihedral bending, shape matching on the feature patches
/// and a weak pull back to the rest shape. One step = one 60 Hz frame; there is no dt.
/// </summary>
public sealed class Face
{
	// Per-vertex state (original offsets in the 4496-byte Face object in comments).
	public Vec3[] Pos;          // +0
	public Vec3[] PrevPos;      // +12
	public Vec3[] Normals;      // +24
	public int[] Indices;       // +36
	public float[] InvMass;     // +48
	public float[] Wrinkle;     // +60
	public Vec3[] Vel;          // +72
	public Vec3[] PrevVel;      // +84
	public Vec3[] Pred;         // +96
	public Vec3[] RestPos;      // +108
	public int[] HiddenTriangles; // +120
	public float[] Mouthiness;  // +132

	public readonly FaceStyle Style;

	/// <summary>Hair shell, or null for a bald style.</summary>
	public Shell Hair;
	/// <summary>Eyelid shells over each eyeball, or nulls.</summary>
	public readonly Shell[] Lids = new Shell[2];
	public readonly Part[] Eyes = new Part[2];
	public Part Nose;
	public readonly Part[] Ears = new Part[2];

	public readonly int[] EyeVertex = new int[2];
	public int MouthVertex;
	public int NoseVertex;
	public readonly int[] EarVertex = new int[2];
	/// <summary>Rings of vertices around the mouth vertex that form the mouth (2 in the original).</summary>
	public int MouthRings => Style.MouthRings;

	List<Edge> _edges;
	List<Bend> _bends;
	// Flattened copies of _edges/_bends for the solver (rebuilt whenever the lists change).
	DistanceConstraints _distance;
	BendConstraints _bending;

	// Per-step scratch: restRotation applied to RestPos (cached while the rotation is unchanged)
	// and the per-vertex rest-pull weights k / (1 - k) while dragging.
	Vec3[] _rotatedRest;
	Quat _rotatedRestFor;
	bool _rotatedRestValid;
	float[] _pullK;
	float[] _pullKeep;

	readonly ShapeRegion[] _eyeRegion = { new(), new() };
	readonly ShapeRegion _mouthRegion = new();
	readonly ShapeRegion _noseRegion = new();
	readonly ShapeRegion[] _earRegion = { new(), new() };

	public int VertexCount => Pos.Length;
	public int TriangleCount => Indices.Length / 3;

	public Face() : this( FaceStyle.Original ) { }

	public Face( FaceStyle style )
	{
		Style = style;
		var posList = new List<Vec3>();
		var idxList = new List<int>();
		Geometry.Icosphere( 4, posList, idxList );
		Pos = posList.ToArray();
		Indices = idxList.ToArray();
		int n = Pos.Length;

		EyeVertex[0] = Nearest( style.LeftEye );
		EyeVertex[1] = Nearest( style.RightEye );
		NoseVertex = Nearest( style.Nose );
		EarVertex[0] = Nearest( style.LeftEar );
		EarVertex[1] = Nearest( style.RightEar );
		MouthVertex = Nearest( style.Mouth );

		if ( style.HeadShape is not null )
		{
			for ( int v = 0; v < n; v++ )
				Pos[v] = style.HeadShape( Pos[v] );
		}

		Wrinkle = new float[n];
		Array.Fill( Wrinkle, 1.0f );

		var hidden = new List<int>();
		if ( style.Hair )
		{
			Hair = Shell.CreateHair( Pos, Indices );
			for ( int t = 0; t < TriangleCount; t++ )
			{
				if ( Hair.Inside[Indices[t * 3]] && Hair.Inside[Indices[t * 3 + 1]] && Hair.Inside[Indices[t * 3 + 2]] )
					hidden.Add( t );
			}
			for ( int v = 0; v < n; v++ )
				if ( Hair.Touched[v] ) Wrinkle[v] = 0.0f;
		}
		HiddenTriangles = hidden.ToArray();

		// Carve the mouth: pull the two rings around the mouth vertex in towards it.
		var mouthRing = Geometry.RingDistances( n, Indices, MouthVertex, MouthRings );
		{
			var c = Pos[MouthVertex];
			var ms = style.MouthScale;
			for ( int v = 0; v < n; v++ )
			{
				int k = mouthRing[v];
				if ( k <= 0 || k > MouthRings ) continue;
				var d = Pos[v] - c;
				float s = (k / (float)MouthRings) * (1.0f / d.Length);
				Pos[v] = new Vec3( d.X * s * ms.X + c.X, d.Y * s * ms.Y + c.Y, d.Z * s * ms.Z + c.Z );
			}

			if ( style.MouthCurve != 0 )
			{
				for ( int v = 0; v < n; v++ )
				{
					if ( mouthRing[v] < 0 || mouthRing[v] > MouthRings ) continue;
					float dx = Pos[v].X - c.X;
					Pos[v] = new Vec3( Pos[v].X, Pos[v].Y - style.MouthCurve * dx * dx, Pos[v].Z );
				}
			}
		}

		PrevPos = Geometry.Copy( Pos );
		RestPos = Geometry.Copy( Pos );
		Normals = Geometry.Copy( Pos );
		InvMass = new float[n];
		Vel = new Vec3[n];
		PrevVel = new Vec3[n];
		Pred = new Vec3[n];

		_edges = Geometry.BuildEdges( Pos, Indices, 0.75f );
		_bends = Geometry.BuildBends( Pos, Indices );
		_distance = new DistanceConstraints( _edges );
		_bending = new BendConstraints( _bends );
		_rotatedRest = new Vec3[n];
		_pullK = new float[n];
		_pullKeep = new float[n];

		for ( int i = 0; i < 30; i++ ) WarmupStep();

		CarvePatch( NoseVertex, style.NoseRings, style.NosePatch, _noseRegion );
		CarvePatch( EyeVertex[0], style.EyeRings, style.EyePatch, _eyeRegion[0] );
		CarvePatch( EyeVertex[1], style.EyeRings, style.EyePatch, _eyeRegion[1] );
		CarvePatch( EarVertex[0], style.EarRings, style.EarPatch, _earRegion[0] );
		CarvePatch( EarVertex[1], style.EarRings, style.EarPatch, _earRegion[1] );
		Geometry.ComputeNormals( Pos, Indices, Normals );

		Nose = new Part( Pos[NoseVertex], Normals[NoseVertex], style.NoseSize, style.NoseCut, 0.0f, 0.0f, style.NoseShape, style.NoseDetail, style.NoseMask );
		Eyes[0] = new Part( Pos[EyeVertex[0]], Normals[EyeVertex[0]], style.EyeSize, style.EyeCut, 0.0f, 0.0f );
		Eyes[1] = new Part( Pos[EyeVertex[1]], Normals[EyeVertex[1]], style.EyeSize, style.EyeCut, 0.0f, 0.0f );
		var earShape = style.EarShape;
		Func<Vec3, Vec3> leftEarShape = earShape is null ? null : p => { var q = earShape( new Vec3( -p.X, p.Y, p.Z ) ); return new Vec3( -q.X, q.Y, q.Z ); };
		Ears[0] = new Part( Pos[EarVertex[0]], Normals[EarVertex[0]], style.EarSize, style.EarCut, -style.EarShear, -style.EarRotation, leftEarShape, style.EarDetail );
		Ears[1] = new Part( Pos[EarVertex[1]], Normals[EarVertex[1]], style.EarSize, style.EarCut, style.EarShear, style.EarRotation, earShape, style.EarDetail );

		if ( style.Eyelids )
		{
			var lidNormal = new Vec3( 0, 1, style.LidTilt ).Normalized;
			float thickness = style.LidThickness;
			for ( int i = 0; i < 2; i++ )
			{
				Lids[i] = new Shell( Eyes[i].RestUnaligned, Eyes[i].Indices, new Vec3( 0, style.LidHeight, 0 ), lidNormal, _ => thickness );
				Lids[i].Update( Eyes[i].Pos, Eyes[i].Normals );
			}
		}

		var mouthIdx = new List<int>();
		for ( int v = 0; v < n; v++ )
			if ( mouthRing[v] != -1 ) mouthIdx.Add( v );
		_mouthRegion.Init( true, mouthIdx, Pos );

		Mouthiness = new float[n];
		for ( int v = 0; v < n; v++ )
			Mouthiness[v] = mouthRing[v] == -1 ? 0.0f : 1.0f - mouthRing[v] / (float)MouthRings;
		for ( int v = 0; v < n; v++ )
			if ( (uint)mouthRing[v] < MouthRings + 1 ) Wrinkle[v] = 0.0f;

		// The mouth centre and first ring lose their springs; no bending anywhere in the mouth.
		bool InInner( int v ) => mouthRing[v] > -1 && mouthRing[v] < MouthRings;
		_edges.RemoveAll( e => InInner( e.A ) || InInner( e.B ) );
		_bends.RemoveAll( b => mouthRing[b.P1] != -1 || mouthRing[b.P2] != -1 || mouthRing[b.P3] != -1 || mouthRing[b.P4] != -1 );
		_distance = new DistanceConstraints( _edges );
		_bending = new BendConstraints( _bends );

		for ( int i = 0; i < 15; i++ ) WarmupStep();
	}

	void WarmupStep() =>
		Step( Quat.Identity, false, 0, Vec3.Zero, 0.8f );

	int Nearest( Vec3 target )
	{
		int best = 0;
		float bestDist = 100000000.0f;
		for ( int i = 0; i < Pos.Length; i++ )
		{
			float d = (target - Pos[i]).Length;
			if ( d < bestDist ) { bestDist = d; best = i; }
		}
		return best;
	}

	/// <summary>Scale a feature patch towards its centre, flatten it and make it shape-matched (f_ic).</summary>
	void CarvePatch( int center, int rings, Vec3 scale, ShapeRegion region )
	{
		int n = Pos.Length;
		var ring = Geometry.RingDistances( n, Indices, center, rings );
		var c = Pos[center];
		float invRings = 1.0f / rings;
		for ( int v = 0; v < n; v++ )
		{
			int k = ring[v];
			if ( k < 1 || !(k <= (float)rings) ) continue;
			var d = Pos[v] - c;
			float x = invRings * k * (1.0f / d.Length);
			Pos[v] = new Vec3( d.X * scale.X * x + c.X, d.Y * scale.Y * x + c.Y, d.Z * scale.Z * x + c.Z );
		}

		// Outer boundary loop: edges of triangles with two verts on the last ring and one inside it.
		var next = new Dictionary<int, int>();
		int inner = rings - 1;
		for ( int t = 0; t + 2 < Indices.Length; t += 3 )
		{
			int a = Indices[t], b = Indices[t + 1], cc = Indices[t + 2];
			if ( ring[a] == rings && ring[b] == rings && ring[cc] == inner ) next[a] = b;
			if ( ring[b] == rings && ring[cc] == rings && ring[a] == inner ) next[b] = cc;
			if ( ring[cc] == rings && ring[a] == rings && ring[b] == inner ) next[cc] = a;
		}
		int start = Array.IndexOf( ring, rings );
		var loop = new List<int> { start };
		while ( true )
		{
			next.TryGetValue( loop[^1], out int nx );
			if ( nx == loop[0] || loop.Count > n ) break;
			loop.Add( nx );
		}

		var p0 = Pos[loop[0]];
		var normal = Vec3.Cross( Pos[loop[1]] - p0, Pos[loop[2]] - p0 ).Normalized;
		if ( Style.FitPatchPlanes )
		{
			// Newell's method: best-fit normal of the whole loop, through its centroid.
			var sum = Vec3.Zero;
			var centroid = Vec3.Zero;
			for ( int i = 0; i < loop.Count; i++ )
			{
				var a = Pos[loop[i]];
				var b = Pos[loop[(i + 1) % loop.Count]];
				sum += new Vec3( (a.Y - b.Y) * (a.Z + b.Z), (a.Z - b.Z) * (a.X + b.X), (a.X - b.X) * (a.Y + b.Y) );
				centroid += a;
			}
			normal = sum.Normalized;
			p0 = centroid / loop.Count;
		}
		for ( int v = 0; v < n; v++ )
		{
			if ( ring[v] == -1 ) continue;
			float d = Vec3.Dot( Pos[v] - p0, normal );
			Pos[v] -= Style.PatchFlatten == 1.0f ? normal * d : normal * (d * Style.PatchFlatten);
		}

		var idx = new List<int>();
		for ( int v = 0; v < n; v++ )
			if ( ring[v] != -1 ) idx.Add( v );
		region.Init( false, idx, Pos );

		for ( int v = 0; v < n; v++ )
			if ( ring[v] <= rings && ring[v] > -1 ) Wrinkle[v] = 0.0f;
	}

	/// <summary>
	/// Advance one frame (f_zd). The JS always passed an identity <paramref name="restRotation"/>.
	/// </summary>
	/// <param name="dragging">Whether a triangle is grabbed.</param>
	/// <param name="triangle">Grabbed triangle index.</param>
	/// <param name="target">Where the grabbed triangle's first vertex should go (face-local).</param>
	/// <param name="damping">Velocity damping: 0.8 while dragging, 0.97 otherwise.</param>
	public void Step( Quat restRotation, bool dragging, int triangle, Vec3 target, float damping )
	{
		// The per-vertex loops below spell out the Vec3 math by hand: the s&box editor compiles game
		// code unoptimized, where every operator/property is a real call. Each expression keeps the
		// original evaluation order, so the results are bit-identical to the operator version.
		Vec3[] pos = Pos, prevPos = PrevPos, vel = Vel, prevVel = PrevVel, pred = Pred, rest = RestPos;
		float[] invMass = InvMass;
		int n = pos.Length;
		int t0 = Indices[triangle * 3], t1 = Indices[triangle * 3 + 1], t2 = Indices[triangle * 3 + 2];

		Array.Fill( invMass, 1.0f );
		if ( dragging )
		{
			invMass[t0] = 0.0f;
			invMass[t1] = 0.0f;
			invMass[t2] = 0.0f;
		}

		// BDF2 predictor with h = 1 frame.
		for ( int v = 0; v < n; v++ )
		{
			ref Vec3 vl = ref vel[v];
			float vx = vl.X * damping, vy = vl.Y * damping, vz = vl.Z * damping;
			vl.X = vx; vl.Y = vy; vl.Z = vz;
			ref Vec3 p = ref pos[v];
			ref Vec3 pp = ref prevPos[v];
			ref Vec3 pv = ref prevVel[v];
			ref Vec3 pr = ref pred[v];
			pr.X = p.X * 1.333333f + vx * 0.888889f + pp.X * -0.333333f + pv.X * -0.222222f;
			pr.Y = p.Y * 1.333333f + vy * 0.888889f + pp.Y * -0.333333f + pv.Y * -0.222222f;
			pr.Z = p.Z * 1.333333f + vz * 0.888889f + pp.Z * -0.333333f + pv.Z * -0.222222f;
		}

		if ( dragging )
		{
			// Carry the grabbed triangle rigidly, rotated so it stays tangent to the surface.
			var a = rest[t0].Normalized;
			var b = target.Normalized;
			var axis = Vec3.Cross( a, b );
			float s = axis.Length;
			var r = Quat.Identity;
			if ( s > 0 )
				r = Quat.FromAxisAngle( axis / s, MathF.Asin( MathF.Min( s, 1.0f ) ) );
			pred[t0] = target;
			pred[t1] = target + r.Rotate( rest[t1] - rest[t0] );
			pred[t2] = target + r.Rotate( rest[t2] - rest[t0] );
		}

		// Rest shape under restRotation; constant for the step (and across steps while it is unchanged).
		Vec3[] rotRest = _rotatedRest;
		if ( !_rotatedRestValid || restRotation.X != _rotatedRestFor.X || restRotation.Y != _rotatedRestFor.Y
			|| restRotation.Z != _rotatedRestFor.Z || restRotation.W != _rotatedRestFor.W )
		{
			float qx = restRotation.X, qy = restRotation.Y, qz = restRotation.Z, qw = restRotation.W;
			for ( int v = 0; v < n; v++ )
			{
				ref Vec3 rp = ref rest[v];
				// Quat.Rotate
				float tx = 2f * (qy * rp.Z - qz * rp.Y);
				float ty = 2f * (qz * rp.X - qx * rp.Z);
				float tz = 2f * (qx * rp.Y - qy * rp.X);
				ref Vec3 o = ref rotRest[v];
				o.X = rp.X + qw * tx + (qy * tz - qz * ty);
				o.Y = rp.Y + qw * ty + (qz * tx - qx * tz);
				o.Z = rp.Z + qw * tz + (qx * ty - qy * tx);
			}
			_rotatedRestFor = restRotation;
			_rotatedRestValid = true;
		}

		// Rest-pull weight k, constant over the solver iterations.
		float[] pullK = _pullK, pullKeep = _pullKeep;
		if ( dragging )
		{
			var grab = rest[t0];
			for ( int v = 0; v < n; v++ )
			{
				ref Vec3 rp = ref rest[v];
				float k = invMass[v] * 0.01f * (0.5f - (grab.X * rp.X + grab.Y * rp.Y + grab.Z * rp.Z) * 0.5f);
				pullK[v] = k;
				pullKeep[v] = 1.0f - k;
			}
		}

		for ( int iter = 0; iter < 3; iter++ )
		{
			Geometry.SolveDistance( _distance, pred, invMass );
			Geometry.SolveBending( _bending, pred, invMass, 0.25f );
			if ( _mouthRegion.Enabled ) _mouthRegion.Solve( pred );

			if ( dragging )
			{
				for ( int v = 0; v < n; v++ )
				{
					float k = pullK[v], keep = pullKeep[v];
					ref Vec3 rr = ref rotRest[v];
					ref Vec3 pr = ref pred[v];
					pr.X = rr.X * k + pr.X * keep;
					pr.Y = rr.Y * k + pr.Y * keep;
					pr.Z = rr.Z * k + pr.Z * keep;
				}
			}
			else
			{
				// InvMass is all 1 here, so k = InvMass * 0.001 is the same for every vertex.
				float k = 1.0f * 0.001f, keep = 1.0f - k;
				for ( int v = 0; v < n; v++ )
				{
					ref Vec3 rr = ref rotRest[v];
					ref Vec3 pr = ref pred[v];
					pr.X = rr.X * k + pr.X * keep;
					pr.Y = rr.Y * k + pr.Y * keep;
					pr.Z = rr.Z * k + pr.Z * keep;
				}
			}

			if ( _eyeRegion[0].Enabled ) _eyeRegion[0].Solve( pred );
			if ( _eyeRegion[1].Enabled ) _eyeRegion[1].Solve( pred );
			if ( _noseRegion.Enabled ) _noseRegion.Solve( pred );
			if ( _earRegion[0].Enabled ) _earRegion[0].Solve( pred );
			if ( _earRegion[1].Enabled ) _earRegion[1].Solve( pred );
		}

		for ( int v = 0; v < n; v++ )
		{
			ref Vec3 p = ref pos[v];
			ref Vec3 pr = ref pred[v];
			ref Vec3 pp = ref prevPos[v];
			ref Vec3 vl = ref vel[v];
			prevVel[v] = vl;
			vl.X = p.X * -2.0f + pr.X * 1.5f + pp.X * 0.5f;
			vl.Y = p.Y * -2.0f + pr.Y * 1.5f + pp.Y * 0.5f;
			vl.Z = p.Z * -2.0f + pr.Z * 1.5f + pp.Z * 0.5f;
			pp = p;
			p = pr;
		}

		Geometry.ComputeNormals( pos, Indices, Normals );
		Hair?.Update( pos, Normals );

		Eyes[0]?.Update( pos[EyeVertex[0]], Normals[EyeVertex[0]], _eyeRegion[0].Rotation, vel[EyeVertex[0]] );
		Eyes[1]?.Update( pos[EyeVertex[1]], Normals[EyeVertex[1]], _eyeRegion[1].Rotation, vel[EyeVertex[1]] );
		Nose?.Update( pos[NoseVertex], Normals[NoseVertex], _noseRegion.Rotation, vel[NoseVertex] );
		Ears[0]?.Update( pos[EarVertex[0]], Normals[EarVertex[0]], _earRegion[0].Rotation, vel[EarVertex[0]] );
		Ears[1]?.Update( pos[EarVertex[1]], Normals[EarVertex[1]], _earRegion[1].Rotation, vel[EarVertex[1]] );
		for ( int i = 0; i < 2; i++ )
			if ( Lids[i] is not null ) Lids[i].Update( Eyes[i].Pos, Eyes[i].Normals );
	}

	/// <summary>
	/// Triangle under a face-local ray, or -1 (f_oxa). Hair blocks grabbing, and a hit on the mouth
	/// is redirected to the nearest non-mouth triangle.
	/// </summary>
	public int RayIntersect( Vec3 origin, Vec3 dir )
	{
		int faceTri = ClosestHit( origin, dir, Pos, Indices, out float faceT );
		if ( faceTri == -1 ) return -1;
		if ( Hair is not null )
		{
			int hairTri = ClosestHit( origin, dir, Hair.Pos, Hair.Indices, out float hairT );
			if ( hairTri != -1 && hairT < faceT ) return -1;
		}

		int a = Indices[faceTri * 3], b = Indices[faceTri * 3 + 1], c = Indices[faceTri * 3 + 2];
		if ( !(Mouthiness[a] > 0) && !(Mouthiness[b] > 0) && !(Mouthiness[c] > 0) )
			return faceTri;
		return NearestNonMouthTriangle( faceTri );
	}

	static int ClosestHit( Vec3 origin, Vec3 dir, Vec3[] pos, int[] idx, out float bestT )
	{
		bestT = float.MaxValue;
		int best = -1;
		for ( int t = 0; t + 2 < idx.Length; t += 3 )
		{
			if ( !Geometry.RayTriangle( origin, dir, pos[idx[t]], pos[idx[t + 1]], pos[idx[t + 2]], out float d ) ) continue;
			if ( d < bestT ) { bestT = d; best = t / 3; }
		}
		if ( best == -1 ) bestT = -1;
		return best;
	}

	int NearestNonMouthTriangle( int tri )
	{
		var r = RestPos;
		var sum = r[Indices[tri * 3]] + r[Indices[tri * 3 + 1]] + r[Indices[tri * 3 + 2]];
		int best = -1;
		float bestD = float.MaxValue;
		for ( int t = 0; t < TriangleCount; t++ )
		{
			int a = Indices[t * 3], b = Indices[t * 3 + 1], c = Indices[t * 3 + 2];
			var d = (sum - r[a] - r[b] - r[c]) * 0.333333f;
			float len = d.Length;
			if ( len < bestD && !(Mouthiness[a] > 0) && !(Mouthiness[b] > 0) && !(Mouthiness[c] > 0) )
			{
				bestD = len;
				best = t;
			}
		}
		return best;
	}

	/// <summary>Current rotation of an eye's shape-matched patch (what the eyeball rides on).</summary>
	public Quat EyeRotation( int eye ) => _eyeRegion[eye].Rotation;

	/// <summary>Centre of an eyeball's sphere, in face space.</summary>
	public Vec3 EyeCentre( int eye )
	{
		var e = Eyes[eye];
		var local = new Vec3( 0, 0, -Style.EyeCut * Style.EyeSize.Z );
		return Pos[EyeVertex[eye]] + _eyeRegion[eye].Rotation.Rotate( e.Align.Rotate( local ) );
	}

	public float ComputeVolume() => Metrics.Volume( Pos, Indices );
	public float ComputeArea() => Metrics.Area( Pos, Indices );
	public float ComputeAverageSpeed() => Metrics.AverageSpeed( Vel );
}