namespace AutoRig.Dl.RigNet;

/// <summary>
/// SKINNET (models/SKINNING.py, nearest_bone=5, use_Dg+use_Lf): per-vertex logits
/// over the 5 nearest bones from position + 8 features per bone
/// (bone start xyz, bone end xyz, 1/geodesic-distance, is-leaf).
/// </summary>
public sealed class SkinNet
{
    public const int NearestBones = 5;
    public const int FeaturesPerBone = 8;

    public required Mlp Transform1;      // [3+40 → 128 → 64]
    public required Gcu Gcu1;            // 64 → 512
    public required Mlp Transform2;      // [512 → 512 → 1024], max-pooled global
    public required Gcu Gcu2;            // 512 → 256
    public required Gcu Gcu3;            // 256 → 256
    public required LinearReluBn Cls0;   // 1280 → 1024
    public required LinearReluBn Cls1;   // 1024 → 512
    public required Linear ClsOut;       // 512 → 5

    /// <param name="pos">[N, 3] vertex positions (normalized space).</param>
    /// <param name="skinInput">[N, 40] nearest-bone features.</param>
    /// <returns>[N, 5] logits (softmax gives weights over the 5 nearest bones).</returns>
    public Tensor Forward(
        Tensor pos, Tensor skinInput,
        (int From, int To)[] tplEdges, (int From, int To)[] geoEdges )
    {
        var x0 = Transform1.Forward( pos.Concat( skinInput, 1 ) );
        var x1 = Gcu1.Forward( x0, tplEdges, geoEdges );
        var global = PointNet.GlobalMaxPool( Transform2.Forward( x1 ) )
            .RepeatRows( pos.Shape[0] );
        var x2 = Gcu2.Forward( x1, tplEdges, geoEdges );
        var x3 = Gcu3.Forward( x2, tplEdges, geoEdges );
        return ClsOut.Forward( Cls1.Forward( Cls0.Forward( x3.Concat( global, 1 ) ) ) );
    }
}