using AutoRig.Mesh;
using AutoRig.Voxel;

namespace AutoRig.Dl.RigNet;

using Vector3 = System.Numerics.Vector3;

/// <summary>
/// Vertex→bone volumetric geodesic distances (quick_start.calc_geodesic_matrix):
/// straight-line distance where the bone can "see" the vertex, otherwise a detour
/// over the surface to the nearest visible vertex. Visibility uses the solid
/// voxelization (segment fully inside ⇒ unoccluded) instead of the reference's
/// triangle ray casts; sub-sampling is deterministic instead of random.
/// </summary>
public static class VolumetricGeodesic
{
    /// <summary>
    /// Computes the [V, B] distance matrix. Vertices are strided down to
    /// <paramref name="maxSubsamples"/> for the heavy part and mapped back by
    /// nearest subsample, like the reference's subsampling=True path.
    /// </summary>
    public static float[,] Compute(
        RigMesh mesh, SkinBone[] bones, SurfaceGeodesic surfaceGeodesic, VoxelGrid voxels,
        int maxSubsamples = 1500 )
    {
        ArgumentNullException.ThrowIfNull( mesh );
        var vertices = mesh.Positions;

        // Deterministic strided subsample.
        var subCount = Math.Min( vertices.Length, maxSubsamples );
        var subIndices = new int[subCount];
        for ( var i = 0; i < subCount; i++ )
            subIndices[i] = (int)((long)i * vertices.Length / subCount);

        // pts2line: distance and nearest point on each bone segment.
        var distance = new float[subCount, bones.Length];
        var visible = new bool[subCount, bones.Length];
        for ( var s = 0; s < subCount; s++ )
        {
            var p = vertices[subIndices[s]];
            for ( var b = 0; b < bones.Length; b++ )
            {
                var origin = ClosestOnSegment( bones[b].Start, bones[b].End, p );
                distance[s, b] = Vector3.Distance( origin, p );
                visible[s, b] = IsUnoccluded( origin, p, voxels );
            }
        }

        // Prune "visible but far": beyond 1.3 × the 15th percentile per bone.
        var visibleDistances = new List<float>();
        for ( var b = 0; b < bones.Length; b++ )
        {
            visibleDistances.Clear();
            for ( var s = 0; s < subCount; s++ )
                if ( visible[s, b] )
                    visibleDistances.Add( distance[s, b] );
            if ( visibleDistances.Count == 0 )
                continue;
            var threshold = Percentile( visibleDistances, 15f );
            for ( var s = 0; s < subCount; s++ )
                if ( distance[s, b] > 1.3f * threshold )
                    visible[s, b] = false;
        }

        // Visible → direct distance; invisible → surface detour to the nearest
        // visible subsample; no visible at all → direct distance.
        var matrix = new float[subCount, bones.Length];
        var visibleIds = new List<int>();
        for ( var b = 0; b < bones.Length; b++ )
        {
            visibleIds.Clear();
            for ( var s = 0; s < subCount; s++ )
            {
                if ( visible[s, b] )
                {
                    visibleIds.Add( s );
                    matrix[s, b] = distance[s, b];
                }
            }
            if ( visibleIds.Count == 0 )
            {
                for ( var s = 0; s < subCount; s++ )
                    matrix[s, b] = distance[s, b];
                continue;
            }
            for ( var s = 0; s < subCount; s++ )
            {
                if ( visible[s, b] )
                    continue;
                var bestGeo = float.MaxValue;
                var bestVisible = visibleIds[0];
                foreach ( var v in visibleIds )
                {
                    var g = surfaceGeodesic.Distance( subIndices[s], subIndices[v] );
                    if ( g < bestGeo )
                    {
                        bestGeo = g;
                        bestVisible = v;
                    }
                }
                matrix[s, b] = bestGeo + matrix[bestVisible, b];
            }
        }

        if ( subCount == vertices.Length )
            return matrix;

        // Map every vertex to its nearest subsample.
        var full = new float[vertices.Length, bones.Length];
        for ( var v = 0; v < vertices.Length; v++ )
        {
            var best = 0;
            var bestDistance = float.MaxValue;
            for ( var s = 0; s < subCount; s++ )
            {
                var d = Vector3.DistanceSquared( vertices[v], vertices[subIndices[s]] );
                if ( d < bestDistance )
                {
                    bestDistance = d;
                    best = s;
                }
            }
            for ( var b = 0; b < bones.Length; b++ )
                full[v, b] = matrix[best, b];
        }
        return full;
    }

    /// <summary>Segment fully inside the solid voxelization ⇒ nothing occludes it.</summary>
    public static bool IsUnoccluded( Vector3 from, Vector3 to, VoxelGrid voxels )
    {
        var length = Vector3.Distance( from, to );
        var steps = Math.Max( 1, (int)(length / (voxels.CellSize * 0.5f)) );
        for ( var i = 0; i <= steps; i++ )
        {
            var p = Vector3.Lerp( from, to, i / (float)steps );
            var (x, y, z) = voxels.VoxelOf( p );
            if ( !voxels.IsSolid( x, y, z ) )
                return false;
        }
        return true;
    }

    static Vector3 ClosestOnSegment( Vector3 a, Vector3 b, Vector3 p )
    {
        var ab = b - a;
        var lengthSq = ab.LengthSquared();
        if ( lengthSq < 1e-8f )
            return a;
        var t = Math.Clamp( Vector3.Dot( p - a, ab ) / lengthSq, 0f, 1f );
        return a + ab * t;
    }

    /// <summary>numpy-style linear-interpolated percentile.</summary>
    public static float Percentile( List<float> values, float percent )
    {
        var sorted = values.OrderBy( v => v ).ToArray();
        if ( sorted.Length == 1 )
            return sorted[0];
        var rank = percent / 100f * (sorted.Length - 1);
        var low = (int)MathF.Floor( rank );
        var high = Math.Min( low + 1, sorted.Length - 1 );
        return sorted[low] + (sorted[high] - sorted[low]) * (rank - low);
    }
}