using AutoRig.Analyze;
using AutoRig.Dl;
using AutoRig.Dl.RigAnything;
using AutoRig.Dl.UniRig;
using AutoRig.Rig;

namespace AutoRig.Solve;

using Vector3 = System.Numerics.Vector3;

/// <summary>
/// RigAnything inference (deterministic modules golden-verified vs PyTorch):
/// point-cloud conditioned autoregressive DIFFUSION rigging. Joint positions
/// are sampled from the AdaLN diffusion head (seeded RNG — the reference is
/// stochastic by design), parents come from the pairwise decoder with the
/// self-parent stop rule, and skin weights are PREDICTED by the network's
/// skinning MLP (top influences, 0.068 threshold, renormalized) — the first
/// model in the catalog with learned skinning. NONCOMMERCIAL research license.
/// </summary>
public static class RigAnythingSolver
{
    public const int NumSamples = 1024;
    public const float SkinThreshold = 0.068f;

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

    public static RigResult Rig( AnalysisResult analysis, RigAnythingModel model )
    {
        ArgumentNullException.ThrowIfNull( analysis );
        ArgumentNullException.ThrowIfNull( model );
        var mesh = analysis.Mesh;

        // ---- input: 1024 surface samples in [-1,1] (max-|coord| convention) ----
        Progress?.Invoke( "input: sampling 1024-point cloud" );
        var (points, normals, center, scale) =
            UniRigInput.SampleCloud( mesh, NumSamples, vertexSamples: 0 );

        // ---- stage 1: pc + start through the transformer ----
        Progress?.Invoke( "transformer: encoding point cloud" );
        var cache = new OptKvCache( model.Blocks.Length );
        var out1 = model.Forward(
            model.PcTokenize( points, normals ).Concat( model.StartToken, 0 ),
            cache, pcCount: NumSamples );
        var pcOut = out1.Slice( 0, 0, NumSamples );
        var z = out1.Slice( 0, NumSamples, 1 );

        // ---- autoregressive joint generation ----
        var rng = new Random( 0 );   // reference sampling is stochastic; we pin the seed
        var joints = new List<Vector3>();
        var parents = new List<int>();
        var fused = new List<Tensor>();

        while ( joints.Count < RigAnythingModel.MaxJoints )
        {
            var xyz = model.Diffusion.Sample( z, rng );
            joints.Add( xyz );
            var index = joints.Count - 1;

            var positionToken = model.JointTokenize( new[] { xyz } );
            fused.Add( model.FuseJoint( z, index, positionToken ) );

            var scores = model.ParentScores( fused[index], fused );
            var parent = 0;
            for ( var candidate = 1; candidate < scores.Length; candidate++ )
                if ( scores[candidate] > scores[parent] )
                    parent = candidate;

            // Stop: the joint claimed itself as parent (and it is not the root).
            if ( parent == index && joints.Count > 1 )
            {
                joints.RemoveAt( index );
                fused.RemoveAt( index );
                break;
            }
            parents.Add( parent );
            Progress?.Invoke( $"decode: joint {joints.Count} (parent {parent})" );

            var parentToken = model.JointTokenize( new[] { joints[parent] } );
            var nextInput = model.ComposeInputToken( index, positionToken, parent, parentToken );
            z = model.Forward( nextInput, cache, pcCount: 0 );
        }

        if ( joints.Count < 2 )
            throw new FormatException(
                $"RigAnything produced {joints.Count} joint(s) - too few for a rig." );

        // ---- skeleton (parents are explicit and topological) ----
        var skeleton = new RigSkeleton();
        for ( var j = 0; j < joints.Count; j++ )
            skeleton.Joints.Add( new RigJoint
            {
                Name = j == 0 ? "root" : $"bone_{j}",
                Parent = j == 0 ? -1 : parents[j],
                Position = joints[j] * scale + center,
            } );
        skeleton.Validate();

        // ---- NEURAL skinning: per-sample scores → nearest-sample transfer ----
        Progress?.Invoke( $"skinning: neural weights for {joints.Count} joints" );
        var jointScores = new float[joints.Count][];
        for ( var j = 0; j < joints.Count; j++ )
            jointScores[j] = model.SkinScores( pcOut, fused[j] );

        var weights = TransferWeights( mesh, points, center, scale, jointScores );
        weights.Validate( mesh, skeleton );

        return new RigResult
        {
            Skeleton = skeleton,
            Weights = weights,
            SolverName = "deep-learning/riganything",
            Degraded = false,
            Explanation = $"RigAnything predicted {joints.Count} joints with "
                + "NEURAL skin weights (diffusion sampling, seeded).",
        };
    }

    /// <summary>Reference postprocess: per point keep the top joint scores,
    /// softmax, drop &lt; 0.068, renormalize; mesh vertices take their nearest
    /// sampled point's weights (we keep 4 influences — the engine's limit).</summary>
    static SkinWeights TransferWeights(
        AutoRig.Mesh.RigMesh mesh, Vector3[] samples,
        Vector3 center, float scale, float[][] jointScores )
    {
        var vertexCount = mesh.Positions.Length;
        var jointCount = jointScores.Length;
        var boneIndices = new int[vertexCount * 4];
        var boneWeights = new float[vertexCount * 4];

        Concurrency.For( 0, vertexCount, v =>
        {
            var p = (mesh.Positions[v] - center) / scale;
            var nearest = 0;
            var bestDistance = float.MaxValue;
            for ( var s = 0; s < samples.Length; s++ )
            {
                var d = Vector3.DistanceSquared( p, samples[s] );
                if ( d < bestDistance )
                {
                    bestDistance = d;
                    nearest = s;
                }
            }

            // Top-4 joints by raw score.
            Span<int> top = stackalloc int[4];
            Span<float> topScore = stackalloc float[4];
            var used = 0;
            for ( var j = 0; j < jointCount; j++ )
            {
                var score = jointScores[j][nearest];
                if ( used < 4 )
                {
                    top[used] = j;
                    topScore[used++] = score;
                }
                else
                {
                    var weakest = 0;
                    for ( var k = 1; k < 4; k++ )
                        if ( topScore[k] < topScore[weakest] )
                            weakest = k;
                    if ( score > topScore[weakest] )
                    {
                        top[weakest] = j;
                        topScore[weakest] = score;
                    }
                }
            }

            // Softmax over the kept scores, threshold, renormalize.
            var max = float.MinValue;
            for ( var k = 0; k < used; k++ )
                max = MathF.Max( max, topScore[k] );
            float total = 0;
            for ( var k = 0; k < used; k++ )
            {
                topScore[k] = MathF.Exp( topScore[k] - max );
                total += topScore[k];
            }
            float kept = 0;
            for ( var k = 0; k < used; k++ )
            {
                topScore[k] /= total;
                if ( topScore[k] < SkinThreshold )
                    topScore[k] = 0;
                kept += topScore[k];
            }
            if ( kept <= 0f )
            {
                // Everything thresholded away — keep the single strongest.
                var strongest = 0;
                for ( var k = 1; k < used; k++ )
                    if ( topScore[k] > topScore[strongest] )
                        strongest = k;
                topScore[strongest] = 1f;
                kept = 1f;
            }
            for ( var k = 0; k < used; k++ )
            {
                boneIndices[v * 4 + k] = top[k];
                boneWeights[v * 4 + k] = topScore[k] / kept;
            }
        } );

        return new SkinWeights { BoneIndices = boneIndices, Weights = boneWeights };
    }
}