namespace AutoRig.Dl.RigNet;

using Vector3 = System.Numerics.Vector3;

/// <summary>A skinning bone: a joint-to-joint segment, or a virtual leaf bone.</summary>
public readonly record struct SkinBone( Vector3 Start, Vector3 End, int StartJoint, int EndJoint, bool IsLeaf );

/// <summary>
/// The skinning stage around SKINNET, transcribed from quick_start.predict_skinning
/// and run_skinning.post_filter: bone extraction with virtual leaf bones, the
/// 5-nearest-bone feature rows, and neighbor smoothing + thresholding of the
/// predicted weights.
/// </summary>
public static class SkinPipeline
{
    /// <summary>
    /// get_bones: level-order traversal emitting (parent → child) bones, plus a
    /// zero-length virtual leaf bone at every childless joint.
    /// </summary>
    public static SkinBone[] GetBones( Vector3[] joints, int[] parents )
    {
        var root = Array.IndexOf( parents, -1 );
        var children = new List<int>[joints.Length];
        for ( var i = 0; i < joints.Length; i++ )
            children[i] = new List<int>();
        for ( var i = 0; i < joints.Length; i++ )
            if ( parents[i] >= 0 )
                children[parents[i]].Add( i );

        var bones = new List<SkinBone>();
        var level = new List<int> { root };
        while ( level.Count > 0 )
        {
            var next = new List<int>();
            foreach ( var p in level )
            {
                next.AddRange( children[p] );
                foreach ( var c in children[p] )
                {
                    bones.Add( new SkinBone( joints[p], joints[c], p, c, IsLeaf: false ) );
                    if ( children[c].Count == 0 )
                        bones.Add( new SkinBone( joints[c], joints[c], c, c, IsLeaf: true ) );
                }
            }
            level = next;
        }
        return bones.ToArray();
    }

    /// <summary>
    /// Builds SKINNET's [N, 40] input plus the per-vertex nearest-bone ids and the
    /// loss mask, from a [N, B] vertex→bone volumetric geodesic distance matrix.
    /// Bones sort ascending by distance (ties → lower bone id).
    /// </summary>
    public static (Tensor Input, int[,] NearestBones, float[,] Mask) BuildSkinInput(
        SkinBone[] bones, float[,] geoDist )
    {
        var vertexCount = geoDist.GetLength( 0 );
        var input = new float[vertexCount * SkinNet.NearestBones * SkinNet.FeaturesPerBone];
        var nearest = new int[vertexCount, SkinNet.NearestBones];
        var mask = new float[vertexCount, SkinNet.NearestBones];

        var order = new int[bones.Length];
        var distances = new float[bones.Length];
        for ( var v = 0; v < vertexCount; v++ )
        {
            for ( var b = 0; b < bones.Length; b++ )
            {
                order[b] = b;
                distances[b] = geoDist[v, b];
            }
            Array.Sort( order, ( a, b ) =>
                distances[a] != distances[b] ? distances[a].CompareTo( distances[b] ) : a.CompareTo( b ) );

            var baseIndex = v * SkinNet.NearestBones * SkinNet.FeaturesPerBone;
            for ( var i = 0; i < SkinNet.NearestBones; i++ )
            {
                // Fewer bones than slots → repeat the nearest with mask 0.
                var bone = i < bones.Length ? order[i] : order[0];
                var at = baseIndex + i * SkinNet.FeaturesPerBone;
                input[at + 0] = bones[bone].Start.X;
                input[at + 1] = bones[bone].Start.Y;
                input[at + 2] = bones[bone].Start.Z;
                input[at + 3] = bones[bone].End.X;
                input[at + 4] = bones[bone].End.Y;
                input[at + 5] = bones[bone].End.Z;
                input[at + 6] = 1f / (distances[bone] + 1e-10f);
                input[at + 7] = bones[bone].IsLeaf ? 1f : 0f;
                nearest[v, i] = i < bones.Length ? bone : 0;
                mask[v, i] = i < bones.Length ? 1f : 0f;
            }
        }
        return (Tensor.From( input, vertexCount, SkinNet.NearestBones * SkinNet.FeaturesPerBone ),
            nearest, mask);
    }

    /// <summary>
    /// predict_skinning post-processing: softmax logits → masked weights scattered
    /// to the full bone set → one-ring neighbor smoothing (post_filter) →
    /// zero out &lt; 0.35 × row max → renormalize. Returns [N, B] bone weights.
    /// </summary>
    public static float[,] FinalizeWeights(
        Tensor logits, int[,] nearestBones, float[,] mask, int boneCount,
        (int From, int To)[] tplEdges )
    {
        var vertexCount = logits.Shape[0];
        var probabilities = logits.Softmax( 1 );

        var full = new float[vertexCount, boneCount];
        for ( var v = 0; v < vertexCount; v++ )
            for ( var i = 0; i < SkinNet.NearestBones; i++ )
                full[v, nearestBones[v, i]] = probabilities.Data[v * SkinNet.NearestBones + i] * mask[v, i];

        full = SmoothOneRing( full, vertexCount, boneCount, tplEdges );

        for ( var v = 0; v < vertexCount; v++ )
        {
            var rowMax = 0f;
            for ( var b = 0; b < boneCount; b++ )
                rowMax = MathF.Max( rowMax, full[v, b] );
            var sum = 0f;
            for ( var b = 0; b < boneCount; b++ )
            {
                if ( full[v, b] < rowMax * 0.35f )
                    full[v, b] = 0f;
                sum += full[v, b];
            }
            for ( var b = 0; b < boneCount; b++ )
                full[v, b] /= sum + 1e-10f;
        }
        return full;
    }

    /// <summary>post_filter(num_ring=1): replace each row with the mean of its
    /// one-ring neighbors (self excluded; kept as-is when isolated).</summary>
    static float[,] SmoothOneRing(
        float[,] weights, int vertexCount, int boneCount, (int From, int To)[] tplEdges )
    {
        var neighbors = new HashSet<int>[vertexCount];
        for ( var v = 0; v < vertexCount; v++ )
            neighbors[v] = new HashSet<int>();
        foreach ( var (from, to) in tplEdges )
        {
            if ( from == to )
                continue;
            neighbors[from].Add( to );
            neighbors[to].Add( from );
        }

        var smoothed = new float[vertexCount, boneCount];
        for ( var v = 0; v < vertexCount; v++ )
        {
            if ( neighbors[v].Count == 0 )
            {
                for ( var b = 0; b < boneCount; b++ )
                    smoothed[v, b] = weights[v, b];
                continue;
            }
            foreach ( var n in neighbors[v] )
                for ( var b = 0; b < boneCount; b++ )
                    smoothed[v, b] += weights[n, b];
            for ( var b = 0; b < boneCount; b++ )
                smoothed[v, b] /= neighbors[v].Count;
        }
        return smoothed;
    }
}