using AutoRig.Dl.RigNet;
using AutoRig.Mesh;

namespace AutoRig.Dl.UniRig;

using AutoRig.Dl;
using Vector3 = System.Numerics.Vector3;

/// <summary>
/// UniRig's input pipeline (verified in unirig-port-notes.md): bbox-normalize to
/// [-1,1]³, SamplerMix a 65536-point cloud (8192 vertex picks with vertex normals
/// + 57344 area-weighted surface samples with face normals — deterministic
/// equivalents of the reference's unseeded sampling), then the perceiver's
/// eval-time reduction: the constant seed-0 4096-subsample and fps down to 1024.
/// </summary>
public static class UniRigInput
{
    public const int NumSamples = 65536;
    public const int VertexSamples = 8192;
    public const int PreCount = 4096;     // token_num * 4
    public const int LatentCount = 1024;  // fps ratio 1/4

    public sealed class Prepared
    {
        /// <summary>The 4096-point pre-cloud (perceiver "data").</summary>
        public required Vector3[] PrePoints;
        public required Vector3[] PreNormals;

        /// <summary>fps picks (1024) into PrePoints (perceiver queries).</summary>
        public required int[] SampledIndices;

        /// <summary>Undo the normalization: world = p * Scale + Center.</summary>
        public required Vector3 Center;
        public required float Scale;
    }

    public static Prepared Prepare( RigMesh mesh )
    {
        ArgumentNullException.ThrowIfNull( mesh );

        var (points, normals, center, scale) = SampleCloud( mesh, NumSamples, VertexSamples );

        // ---- perceiver eval reduction ----
        var prePoints = new Vector3[PreCount];
        var preNormals = new Vector3[PreCount];
        for ( var i = 0; i < PreCount; i++ )
        {
            var index = PerceiverSubsample.Indices[i];
            prePoints[i] = points[index];
            preNormals[i] = normals[index];
        }

        var preTensor = Tensor.From(
            prePoints.SelectMany( p => new[] { p.X, p.Y, p.Z } ).ToArray(), PreCount, 3 );
        var sampled = PointNet.FarthestPointSample( preTensor, ratio: 0.25f );

        return new Prepared
        {
            PrePoints = prePoints,
            PreNormals = preNormals,
            SampledIndices = sampled,
            Center = center,
            Scale = scale,
        };
    }

    /// <summary>SamplerMix (shared with SkinTokens, whose counts differ): bbox
    /// normalize into [-1,1]³, then vertex picks followed by surface samples.</summary>
    internal static (Vector3[] Points, Vector3[] Normals, Vector3 Center, float Scale)
        SampleCloud( RigMesh mesh, int numSamples, int vertexSamples )
    {
        var bounds = mesh.ComputeBounds();
        var center = (bounds.Min + bounds.Max) * 0.5f;
        var half = (bounds.Max - bounds.Min) * 0.5f;
        var scale = MathF.Max( half.X, MathF.Max( half.Y, half.Z ) );
        if ( scale <= 0f )
            throw new FormatException( "UniRig input: mesh has zero extent." );

        var points = new Vector3[numSamples];
        var normals = new Vector3[numSamples];

        var vertexCount = mesh.Positions.Length;
        var vertexPicks = Math.Min( vertexSamples, vertexCount );
        for ( var i = 0; i < vertexPicks; i++ )
        {
            var v = (int)((long)i * vertexCount / vertexPicks);   // strided (ref: random perm)
            points[i] = (mesh.Positions[v] - center) / scale;
            normals[i] = v < mesh.Normals.Length ? mesh.Normals[v] : Vector3.UnitZ;
        }

        FillSurfaceSamples( mesh, center, scale, points, normals,
            vertexPicks, numSamples - vertexPicks );
        return (points, normals, center, scale);
    }

    /// <summary>Area-weighted surface samples with FACE normals (deterministic
    /// low-discrepancy barycentrics — the reference uses unseeded randomness).</summary>
    static void FillSurfaceSamples(
        RigMesh mesh, Vector3 center, float scale,
        Vector3[] points, Vector3[] normals, int writeFrom, int count )
    {
        var triangleCount = mesh.TriangleCount;
        var cumulative = new float[triangleCount];
        float total = 0;
        for ( var t = 0; t < triangleCount; t++ )
        {
            var a = mesh.Positions[mesh.Triangles[t * 3]];
            var b = mesh.Positions[mesh.Triangles[t * 3 + 1]];
            var c = mesh.Positions[mesh.Triangles[t * 3 + 2]];
            total += Vector3.Cross( b - a, c - a ).Length() * 0.5f;
            cumulative[t] = total;
        }
        if ( total <= 0f )
            throw new FormatException( "UniRig input: mesh has zero surface area." );

        for ( var i = 0; i < count; i++ )
        {
            var target = (i + 0.5f) / count * total;
            var t = Array.BinarySearch( cumulative, target );
            if ( t < 0 )
                t = ~t;
            t = Math.Min( t, triangleCount - 1 );

            var a = mesh.Positions[mesh.Triangles[t * 3]];
            var b = mesh.Positions[mesh.Triangles[t * 3 + 1]];
            var c = mesh.Positions[mesh.Triangles[t * 3 + 2]];
            var r1 = (i * 0.7548776662466927f) % 1f;
            var r2 = (i * 0.5698402909980532f) % 1f;
            if ( r1 + r2 > 1f )
            {
                r1 = 1f - r1;
                r2 = 1f - r2;
            }
            var p = a + (b - a) * r1 + (c - a) * r2;
            points[writeFrom + i] = (p - center) / scale;

            var faceNormal = Vector3.Cross( b - a, c - a );
            normals[writeFrom + i] = faceNormal.LengthSquared() > 1e-12f
                ? Vector3.Normalize( faceNormal )
                : Vector3.UnitZ;
        }
    }
}