using AutoRig.Dl.Nn;
using AutoRig.Dl.UniRig;

namespace AutoRig.Dl.Anymate;

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

/// <summary>The three Anymate heads over their PointBERT latents.</summary>
public sealed class AnymateModel
{
    public required PointBert JointEncoder, ConnEncoder, SkinEncoder;

    // ---- joint: 96 queries → cross → ln_post → xyz → cluster ----
    public required Tensor JointQuery;                 // [96, 768]
    public required PerceiverCrossBlock JointCross;
    public required Tensor JointLnPostG, JointLnPostB, JointOutW, JointOutB;

    // ---- conn: dual stream, 2×(cross+self), dot-product parent logits ----
    public required Tensor ConnCoProjW, ConnCoProjB;   // 51 → 768
    public required PerceiverCrossBlock[] ConnCross;   // 2
    public required PerceiverSelfBlock[] ConnSelf;     // 4
    public required Tensor ConnKProjW, ConnKProjB, ConnQProjW, ConnQProjB;
    public required Tensor ConnLn1G, ConnLn1B, ConnLn2G, ConnLn2B;

    // ---- skin: verts+bones streams, 2×(cross latents + cross bones) ----
    public required Tensor SkinCoProjW, SkinCoProjB;       // 51 → 768
    public required Tensor SkinBoneProjW, SkinBoneProjB;   // 102 → 768
    public required Tensor SkinNormalProjW, SkinNormalProjB;   // 25 → 768
    public required PerceiverCrossBlock[] SkinCross;       // 2
    public required PerceiverCrossBlock[] SkinCrossJoint;  // 2
    public required Tensor SkinKProjW, SkinKProjB, SkinQProjW, SkinQProjB;
    public required Tensor SkinLn1G, SkinLn1B, SkinLn2G, SkinLn2B;

    public const int Heads = 12;
    public const float ClusterEps = 0.03f;

    public Action<string> Progress { get; set; }

    /// <summary>Joints from the cloud: 96 samples → eps-graph connected
    /// components (DBSCAN min_samples=1) → cluster means.</summary>
    public List<Vector3> Joints( Vector3[] points )
    {
        Progress?.Invoke( "anymate: joint encoder" );
        var latents = JointEncoder.Encode( points );
        var x = JointCross.Forward( JointQuery, latents, Heads );
        x = TransformerOps.LayerNorm( x, JointLnPostG, JointLnPostB );
        var samples = x.MatMul( JointOutW.Transposed ).Add( JointOutB );   // (96,3)

        var count = samples.Shape[0];
        var cluster = new int[count];
        Array.Fill( cluster, -1 );
        var clusters = 0;
        for ( var i = 0; i < count; i++ )
        {
            if ( cluster[i] >= 0 )
                continue;
            var queue = new Queue<int>();
            queue.Enqueue( i );
            cluster[i] = clusters;
            while ( queue.Count > 0 )
            {
                var a = queue.Dequeue();
                for ( var b = 0; b < count; b++ )
                {
                    if ( cluster[b] >= 0 )
                        continue;
                    var dx = samples.Data[a * 3] - samples.Data[b * 3];
                    var dy = samples.Data[a * 3 + 1] - samples.Data[b * 3 + 1];
                    var dz = samples.Data[a * 3 + 2] - samples.Data[b * 3 + 2];
                    if ( dx * dx + dy * dy + dz * dz <= ClusterEps * ClusterEps )
                    {
                        cluster[b] = clusters;
                        queue.Enqueue( b );
                    }
                }
            }
            clusters++;
        }
        var sums = new Vector3[clusters];
        var counts = new int[clusters];
        for ( var i = 0; i < count; i++ )
        {
            sums[cluster[i]] += new Vector3(
                samples.Data[i * 3], samples.Data[i * 3 + 1], samples.Data[i * 3 + 2] );
            counts[cluster[i]]++;
        }
        var joints = new List<Vector3>();
        for ( var c = 0; c < clusters; c++ )
            joints.Add( sums[c] / counts[c] );
        return joints;
    }

    /// <summary>Parent index per joint (argmax row; self = root).</summary>
    public int[] Parents( Vector3[] points, IReadOnlyList<Vector3> joints )
    {
        Progress?.Invoke( "anymate: connectivity encoder" );
        var latents = ConnEncoder.Encode( points );
        var jointCount = joints.Count;
        var embed = new float[jointCount * 51];
        for ( var j = 0; j < jointCount; j++ )
            Array.Copy( AnymateEmbed.FourierPi( joints[j] ), 0, embed, j * 51, 51 );
        var je = Tensor.From( embed, jointCount, 51 )
            .MatMul( ConnCoProjW.Transposed ).Add( ConnCoProjB );

        var streams = new[] { je, je };
        for ( var i = 0; i < 2; i++ )
            for ( var j = 0; j < 2; j++ )
            {
                streams[i] = ConnCross[j].Forward( streams[i], latents, Heads );
                streams[i] = ConnSelf[2 * i + j].Forward( streams[i], Heads );
            }

        var k = TransformerOps.LayerNorm( streams[0], ConnLn1G, ConnLn1B )
            .MatMul( ConnKProjW.Transposed ).Add( ConnKProjB );
        var q = TransformerOps.LayerNorm( streams[1], ConnLn2G, ConnLn2B )
            .MatMul( ConnQProjW.Transposed ).Add( ConnQProjB );

        var parents = new int[jointCount];
        for ( var i = 0; i < jointCount; i++ )
        {
            var best = 0;
            var bestScore = float.MinValue;
            for ( var j = 0; j < jointCount; j++ )
            {
                float dot = 0;
                for ( var d = 0; d < 768; d++ )
                    dot += k.Data[i * 768 + d] * q.Data[j * 768 + d];
                if ( dot > bestScore )
                {
                    bestScore = dot;
                    best = j;
                }
            }
            parents[i] = best;   // self = root
        }
        return parents;
    }

    /// <summary>Per-query-point weights over bones ([parent|joint] rows).</summary>
    public float[] Skin(
        Vector3[] points, Vector3[] queryPoints, Vector3[] queryNormals,
        float[] bones6, int boneCount )
    {
        Progress?.Invoke( "anymate: skin encoder" );
        var latents = SkinEncoder.Encode( points );

        var qn = queryPoints.Length;
        var pcEmbed = new float[qn * 51];
        var shEmbed = new float[qn * 25];
        for ( var i = 0; i < qn; i++ )
        {
            Array.Copy( AnymateEmbed.FourierPi( queryPoints[i] ), 0, pcEmbed, i * 51, 51 );
            Array.Copy( AnymateEmbed.Sh25( queryNormals[i] ), 0, shEmbed, i * 25, 25 );
        }
        var pce = Tensor.From( pcEmbed, qn, 51 )
            .MatMul( SkinCoProjW.Transposed ).Add( SkinCoProjB )
            .Add( Tensor.From( shEmbed, qn, 25 )
                .MatMul( SkinNormalProjW.Transposed ).Add( SkinNormalProjB ) );

        var boneEmbed = new float[boneCount * 102];
        for ( var b = 0; b < boneCount; b++ )
        {
            var parent = new Vector3( bones6[b * 6], bones6[b * 6 + 1], bones6[b * 6 + 2] );
            var joint = new Vector3( bones6[b * 6 + 3], bones6[b * 6 + 4], bones6[b * 6 + 5] );
            Array.Copy( AnymateEmbed.FourierPi( parent ), 0, boneEmbed, b * 102, 51 );
            Array.Copy( AnymateEmbed.FourierPi( joint ), 0, boneEmbed, b * 102 + 51, 51 );
        }
        var be = Tensor.From( boneEmbed, boneCount, 102 )
            .MatMul( SkinBoneProjW.Transposed ).Add( SkinBoneProjB );

        var x = pce.Concat( be, 0 );
        for ( var i = 0; i < 2; i++ )
        {
            x = SkinCross[i].Forward( x, latents, Heads );
            x = SkinCrossJoint[i].Forward( x, x.Slice( 0, qn, boneCount ), Heads );
        }
        var pcf = TransformerOps.LayerNorm( x.Slice( 0, 0, qn ), SkinLn1G, SkinLn1B )
            .MatMul( SkinKProjW.Transposed ).Add( SkinKProjB );
        var bef = TransformerOps.LayerNorm( x.Slice( 0, qn, boneCount ), SkinLn2G, SkinLn2B )
            .MatMul( SkinQProjW.Transposed ).Add( SkinQProjB );

        var weights = new float[qn * boneCount];
        Concurrency.For( 0, qn, i =>
        {
            Span<float> row = stackalloc float[128];
            var max = float.MinValue;
            for ( var b = 0; b < boneCount; b++ )
            {
                float dot = 0;
                for ( var d = 0; d < 768; d++ )
                    dot += pcf.Data[i * 768 + d] * bef.Data[b * 768 + d];
                row[b] = dot;
                max = MathF.Max( max, dot );
            }
            float total = 0;
            for ( var b = 0; b < boneCount; b++ )
            {
                row[b] = MathF.Exp( row[b] - max );
                total += row[b];
            }
            for ( var b = 0; b < boneCount; b++ )
                weights[i * boneCount + b] = row[b] / total;
        } );
        return weights;
    }
}