#nullable enable annotations

using System;
using System.Collections.Generic;
using System.Linq;
using System.Numerics;
using System.Threading;
using HumanoidMocap.Core.Mapping;
using HumanoidMocap.Core.Maths;
using HumanoidMocap.Core.Motion;

namespace HumanoidMocap.EditorTools.MotionBricks;

using Vector3 = System.Numerics.Vector3;

/// <summary>
/// Rebuilds a capture's body motion with MotionBricks, the way an animator keys a shot: key poses are
/// taken from the capture where the body briefly settles (0.6–1.9 s apart, at lengths the model plays at
/// its own timing), and the motion between them is generated by the model, each in-between continuing
/// the previous one and landing exactly on the next key pose. Where an in-between strays more than <see cref="Tolerance"/> from the
/// capture (a kick between two settled poses), the capture's pose at the worst frame becomes another
/// key and both halves are generated again, so the rebuild keeps the performance. The model has no fingers, neck or head, so those keep the capture. The
/// source document is not changed; the rebuilt one keeps its timestamps, contacts and evidence.
/// </summary>
public static class MotionBricksRebuild
{
    /// <summary>The model's frame rate.</summary>
    const float Fps = G1Motion.Fps;
    /// <summary>Key pose spacing in model frames: long enough for the model to shape the motion, short enough to keep the performance.</summary>
    public const int MinimumGap = 17, MaximumGap = 57;
    /// <summary>The in-between lengths the model generates at its own timing: 6–16 tokens of 4 frames, less the 7 boundary frames.</summary>
    public static readonly int[] NativeLengths = Enumerable.Range( MotionBricksModel.MinTokens, MotionBricksModel.MaxTokens - MotionBricksModel.MinTokens + 1 )
        .Select( t => t * MotionBricksModel.FramesPerToken - 7 ).ToArray();
    /// <summary>How many frames an in-between of <paramref name="length"/> must be stretched or squeezed by (0 at a native length).</summary>
    public static int Warp( int length ) => NativeLengths.Min( n => Math.Abs( n - length ) );
    /// <summary>Shortest gap the model fills (from AI Animator); shorter tails keep the capture.</summary>
    const int ModelMinimumLength = 10;
    /// <summary>How far (metres, relative to the hips) a body joint of an in-between may stray from the capture before a key pose
    /// is added there: enough for the model's own timing and style, not for a missing kick or swing.</summary>
    public const float Tolerance = .15f;

    /// <summary>Tuning, for the offline comparison harness; the defaults are what the editor uses.</summary>
    public sealed record Options
    {
        /// <summary>Take the model's most likely motion instead of a random sample (measured smoother, and repeatable).</summary>
        public bool Argmax { get; init; } = true;
        /// <summary>Gaussian smoothing (frames) of the capture the key poses and targets come from, so the model does not aim at
        /// the capture's frame-to-frame jitter; 0 = raw.</summary>
        public float KeySmoothing { get; init; } = 1.5f;
        /// <summary>Instead of <see cref="KeySmoothing"/>: a quadratic Savitzky–Golay filter of this half-width (frames), which removes
        /// jitter but keeps the height of quick peaks such as kicks; 0 = off.</summary>
        public int KeyPeakKeeping { get; init; } = 4;
        /// <summary>Shortest in-between a split may leave.</summary>
        public int MinimumSplit { get; init; } = ModelMinimumLength;
        public bool SmoothJoins { get; init; } = true;
        public float Tolerance { get; init; } = MotionBricksRebuild.Tolerance;
        /// <summary>Blend back toward the (smoothed) capture where the model still misses a move it could not be split around.</summary>
        public bool RecoverMisses { get; init; } = true;
        /// <summary>Raise the AI's leg to the capture's height at each kick's peak (the thigh only, eased in and out), keeping its timing and knee.</summary>
        public bool KeepKickHeight { get; init; } = true;
        /// <summary>Pin each foot where the capture's foot is planted (leg IK), so the AI's feet do not slide.</summary>
        public bool PlantFeet { get; init; } = true;
        /// <summary>Degrees the AI's pelvis may lean or turn away from the capture's before the pelvis and spine are eased back
        /// toward it (fully at <see cref="TorsoFull"/>); 0 turns it off.</summary>
        public float TorsoFrom { get; init; } = 4f;
        public float TorsoFull { get; init; } = 10f;
        /// <summary>A move counts as missed beyond this distance (metres, relative to the hips), and is fully recovered at <see cref="RecoverFull"/>.</summary>
        public float RecoverFrom { get; init; } = .06f;
        public float RecoverFull { get; init; } = .14f;
        /// <summary>Frames either side a miss is held, and the smoothing (frames) of the blend weight.</summary>
        public int RecoverHold { get; init; }
        public float RecoverSigma { get; init; } = 2f;
    }
    public static Options Settings { get; set; } = new();

    public sealed record Result( MotionDocument Motion, int KeyPoses, int InBetweens, TimeSpan Elapsed, IReadOnlyList<double> KeyTimes = null, float AiShare = 1 );

    public static Result Rebuild( MotionDocument source, MotionBricksModel model, RetargeterLibrary library, IProgress<float>? progress = null, CancellationToken cancel = default )
    {
        var watch = System.Diagnostics.Stopwatch.StartNew();
        var bones = source.Bones;
        if ( !new[] { BoneRole.Hips, BoneRole.UpperLegL, BoneRole.UpperLegR, BoneRole.FootL, BoneRole.FootR, BoneRole.HandL, BoneRole.HandR }.All( r => bones.Any( b => b.Role == r ) ) )
            throw new InvalidOperationException( "MotionBricks rebuilds whole-body captures; this capture has no legs." );
        var root = bones.FindIndex( b => b.Parent < 0 );
        var bridge = new CaptureMotionBridge( library, bones, model.Skeleton );

        // Model frames: the capture resampled to 30 fps, levelled with gravity and standing on y = 0.
        var start = source.Frames[0].Time;
        var count = Math.Max( 2, (int)Math.Floor( (source.Frames[^1].Time - start) * Fps + 1e-6 ) + 1 );
        var level = Level( source );
        var original = new XForm[count][];
        for ( var k = 0; k < count; k++ )
        {
            var pose = Sample( source, start + k / (double)Fps );
            pose[root] = new XForm( Vector3.Transform( pose[root].Pos, level ), Quaternion.Normalize( level * pose[root].Rot ) );
            original[k] = pose;
        }
        var floor = Floor( bones, original );
        for ( var k = 0; k < count; k++ )
            original[k][root] = new XForm( original[k][root].Pos - new Vector3( 0, floor, 0 ), original[k][root].Rot );

        var capture = original;
        if ( Settings.KeyPeakKeeping > 0 )
            original = SmoothPoses( bones, original, SavitzkyGolay( Settings.KeyPeakKeeping ) );
        else if ( Settings.KeySmoothing > 0 )
            original = SmoothPoses( bones, original, Gaussian( Settings.KeySmoothing ) );
        var anchors = Anchors( bones, original );
        var working = original.Select( p => (XForm[])p.Clone() ).ToArray();
        var detail = DetailBones( bones );
        var body = Enumerable.Range( 0, bones.Count ).Where( j => !detail.Contains( j ) || bones[j].Role is BoneRole.Head or BoneRole.HandL or BoneRole.HandR ).ToArray();
        var originalWorld = capture.Select( p => World( bones, p ) ).ToArray();
        var generated = 0;
        var keys = new SortedSet<int>( anchors );
        // In order, so every in-between continues the one before it; a split re-queues both halves at the front.
        var pending = new Stack<(int A, int B)>();
        for ( var s = anchors.Count - 2; s >= 0; s-- )
            pending.Push( (anchors[s], anchors[s + 1]) );
        while ( pending.Count > 0 )
        {
            cancel.ThrowIfCancellationRequested();
            var (a, b) = pending.Pop();
            progress?.Report( a / (float)(count - 1) );
            var length = b - a;
            if ( length < ModelMinimumLength )
                continue;
            var context = Enumerable.Range( a - 3, 4 ).Select( f => working[Math.Clamp( f, 0, count - 1 )] ).ToList();
            var target = Enumerable.Range( b, 4 ).Select( f => original[Math.Clamp( f, 0, count - 1 )] ).ToList();
            var frames = Synthesize( model, bridge, context, target, length, 1UL + (ulong)a * 7919UL, cancel, Settings.Argmax );
            HitKeys( frames, working[a], original[b] );
            for ( var t = 0; t <= length; t++ )
                foreach ( var bone in detail )
                    frames[t][bone] = capture[a + t][bone];
            // The worst frame, comparing the pose shape (relative to the hips).
            var worst = -1;
            var worstDistance = Settings.Tolerance;
            for ( var t = 1; t < length; t++ )
            {
                var w = World( bones, frames[t] );
                var o = originalWorld[a + t];
                var d = body.Max( j => Vector3.Distance( w[j] - w[root], o[j] - o[root] ) );
                if ( d > worstDistance )
                {
                    worstDistance = d;
                    worst = t;
                }
            }
            var split = Math.Max( ModelMinimumLength, Settings.MinimumSplit );
            if ( worst >= 0 && length >= 2 * split )
            {
                // Split near the worst frame, preferring halves the model plays at its own timing (no stretching);
                // halves shorter than its shortest in-between are squeezed only when the gap allows nothing else.
                var m = a + Math.Clamp( worst, split, length - split );
                var best = float.MaxValue;
                for ( var c = a + split; c <= b - split; c++ )
                {
                    var cost = Math.Abs( c - (a + worst) ) + 2f * (Warp( c - a ) + Warp( b - c ));
                    if ( cost < best ) { best = cost; m = c; }
                }
                {
                    keys.Add( m );
                    pending.Push( (m, b) );
                    pending.Push( (a, m) );
                    continue;
                }
            }
            for ( var t = 0; t <= length; t++ )
                working[a + t] = frames[t];
            generated++;
        }
        anchors = keys.ToList();
        if ( Settings.SmoothJoins )
            SmoothJoins( bones, working, anchors, detail, root );
        var captureShare = Settings.RecoverMisses ? RecoverMisses( bones, working, original, originalWorld, detail, root ) : 0;
        if ( Settings.TorsoFrom > 0 )
            RecoverTorso( bones, working, original, detail, root );
        var kicks = Settings.KeepKickHeight ? KeepKickHeight( bones, working, originalWorld ) : 0;
        // Last, so nothing above moves a planted foot again.
        if ( Settings.PlantFeet )
            PlantFeet( bones, working, capture, originalWorld );

        // Back to the capture's own frames, floor and tilt.
        var unlevel = Quaternion.Conjugate( level );
        var rebuilt = source.Copy();
        for ( var i = 0; i < rebuilt.Frames.Count; i++ )
        {
            var frame = rebuilt.Frames[i];
            var g = Math.Clamp( (frame.Time - start) * Fps, 0, count - 1 );
            var lo = (int)Math.Floor( g );
            var hi = Math.Min( lo + 1, count - 1 );
            var blend = (float)(g - lo);
            for ( var j = 0; j < bones.Count; j++ )
            {
                // Bones the model does not drive keep their exact captured values (no resampling).
                if ( j != root && detail.Contains( j ) )
                    continue;
                var p = Vector3.Lerp( working[lo][j].Pos, working[hi][j].Pos, blend );
                var q = Quaternion.Normalize( Quaternion.Slerp( working[lo][j].Rot, working[hi][j].Rot, blend ) );
                if ( j == root )
                {
                    p = Vector3.Transform( p + new Vector3( 0, floor, 0 ), unlevel );
                    q = Quaternion.Normalize( unlevel * q );
                }
                frame.Positions[j] = MotionDocument.A( p );
                frame.Rotations[j] = MotionDocument.A( q );
            }
        }
        rebuilt.Diagnostics.Add( $"Rebuilt with MotionBricks: {anchors.Count} key poses from the capture, {generated} in-betweens generated{(captureShare >= .005f ? FormattableString.Invariant( $", {captureShare * 100:F0}% of the limb motion taken back from the capture where the model missed a move" ) : "")}{(kicks > 0 ? $", {kicks} kicks raised to the captured height" : "")}; fingers, wrists, neck and head kept from the capture." );
        rebuilt.Corrections.Add( new MotionCorrection { Type = "motionbricks-rebuild", Start = source.Frames[0].Time, End = source.Frames[^1].Time,
            Settings = new Dictionary<string, float> { ["keyPoses"] = anchors.Count, ["inBetweens"] = generated } } );
        progress?.Report( 1 );
        return new Result( rebuilt, anchors.Count, generated, watch.Elapsed, anchors.Select( k => start + k / (double)Fps ).ToList(), 1 - captureShare );
    }

    /// <summary>One in-between: port of AI Animator's <c>MotionBricksSynthesizer</c>. The boundary frames go to the G1 in the model's
    /// canonical frame (first context frame at the origin, heading zero), are generated, restored and time-warped onto the gap.</summary>
    static XForm[][] Synthesize( MotionBricksModel model, CaptureMotionBridge bridge, IReadOnlyList<XForm[]> context, IReadOnlyList<XForm[]> target, int length, ulong seed, CancellationToken cancel, bool argmax = false )
    {
        bridge.ResetProjection();
        var boundary = bridge.ToG1( context.Concat( target ).ToList() );
        var origin = boundary.Root[0];
        var heading = MotionRepresentation.Heading( boundary.Local[0][0] );
        var toCanonical = Quaternion.CreateFromAxisAngle( Vector3.UnitY, -heading );
        var fromCanonical = Quaternion.Conjugate( toCanonical );
        var roots = boundary.Root.Select( p => Vector3.Transform( new Vector3( p.X - origin.X, p.Y, p.Z - origin.Z ), toCanonical ) ).ToArray();
        var locals = boundary.Local.Select( l =>
        {
            var copy = (Quaternion[])l.Clone();
            copy[0] = Quaternion.Normalize( toCanonical * copy[0] );
            return copy;
        } ).ToArray();
        var representation = model.Representation;
        var encodedContext = representation.Encode( roots.AsSpan( 0, 4 ), locals.Take( 4 ).ToList(), repeatLastVelocity: false );
        var encodedTarget = representation.Encode( roots.AsSpan( 4, 4 ), locals.Skip( 4 ).ToList(), repeatLastVelocity: true );
        var tokens = Math.Clamp( (int)MathF.Round( (length + 7) / 4f ), MotionBricksModel.MinTokens, MotionBricksModel.MaxTokens );
        var motion = model.RunTransition( new TransitionConstraints { Context = encodedContext, Target = encodedTarget, Tokens = tokens, Seed = seed, Argmax = argmax }, cancel );
        for ( var f = 0; f < motion.FrameCount; f++ )
        {
            var p = Vector3.Transform( motion.Root[f], fromCanonical );
            motion.Root[f] = new Vector3( p.X + origin.X, p.Y, p.Z + origin.Z );
            motion.Local[f][0] = Quaternion.Normalize( fromCanonical * motion.Local[f][0] );
        }
        // Generated frame 3 is the gap's first frame and frame F-4 its last.
        var span = motion.FrameCount - 7;
        var warped = new G1Motion( length + 1 );
        for ( var t = 0; t <= length; t++ )
        {
            var g = 3f + t * span / (float)length;
            var lo = Math.Clamp( (int)MathF.Floor( g ), 0, motion.FrameCount - 1 );
            var hi = Math.Min( lo + 1, motion.FrameCount - 1 );
            var blend = g - lo;
            warped.Root[t] = Vector3.Lerp( motion.Root[lo], motion.Root[hi], blend );
            for ( var j = 0; j < G1Skeleton.JointCount; j++ )
                warped.Local[t][j] = Quaternion.Normalize( Quaternion.Slerp( motion.Local[lo][j], motion.Local[hi][j], blend ) );
        }
        cancel.ThrowIfCancellationRequested();
        return bridge.ToCapture( warped );
    }

    static float[] Gaussian( float sigma )
    {
        var radius = (int)MathF.Ceiling( 3 * sigma );
        return Enumerable.Range( -radius, 2 * radius + 1 ).Select( k => MathF.Exp( -k * k / (2 * sigma * sigma) ) ).ToArray();
    }
    /// <summary>Quadratic Savitzky–Golay smoothing weights over 2·<paramref name="half"/>+1 frames (they sum to 1; the outer ones are negative).</summary>
    static float[] SavitzkyGolay( int half )
    {
        float m = half;
        var norm = (2 * m + 1) * (4 * m * m + 4 * m - 3) / 3f;
        return Enumerable.Range( -half, 2 * half + 1 ).Select( k => (3 * m * m + 3 * m - 1 - 5f * k * k) / norm ).ToArray();
    }

    /// <summary>The capture filtered with <paramref name="kernel"/> (centred) on the bones the model drives and the hips' path.</summary>
    static XForm[][] SmoothPoses( IReadOnlyList<MotionBone> bones, XForm[][] frames, float[] kernel )
    {
        var count = frames.Length;
        var radius = kernel.Length / 2;
        var result = frames.Select( f => (XForm[])f.Clone() ).ToArray();
        for ( var t = 0; t < count; t++ )
            for ( var j = 0; j < bones.Count; j++ )
            {
                if ( bones[j].Group == "fingers" )
                    continue;
                var position = Vector3.Zero;
                var rotation = new Quaternion( 0, 0, 0, 0 );
                var total = 0f;
                for ( var k = -radius; k <= radius; k++ )
                {
                    var f = Math.Clamp( t + k, 0, count - 1 );
                    var w = kernel[k + radius];
                    var q = frames[f][j].Rot;
                    if ( Quaternion.Dot( q, frames[t][j].Rot ) < 0 )
                        q = Quaternion.Negate( q );
                    position += frames[f][j].Pos * w;
                    rotation += q * w;
                    total += w;
                }
                result[t][j] = new XForm( position / total, Quaternion.Normalize( rotation ) );
            }
        return result;
    }

    /// <summary>Where the rebuild still misses a move (one quicker than the model's shortest in-between, or a kick the model keeps
    /// low), that limb is blended toward the smoothed capture. Each leg, each arm and the spine is judged on its own, by where its
    /// end (foot, hand, neck) sits relative to where the limb starts: a miss beyond <see cref="Options.RecoverFrom"/> starts the
    /// blend, fully at <see cref="Options.RecoverFull"/>, the weight smoothed over time so the hand-over never pops. The rest of the
    /// body keeps the model's motion. Returns the share of the motion (limbs × frames) taken from the capture.</summary>
    static float RecoverMisses( IReadOnlyList<MotionBone> bones, XForm[][] frames, XForm[][] target, Vector3[][] captureWorld, HashSet<int> detail, int root )
    {
        var count = frames.Length;
        int Role( BoneRole r ) => bones.ToList().FindIndex( b => b.Role == r );
        var groups = new List<(int Start, int End, int[] Bones)>();
        void Group( int start, int end, Func<int, bool> include )
        {
            if ( start < 0 || end < 0 )
                return;
            var members = Enumerable.Range( 0, bones.Count ).Where( j => j != root && !detail.Contains( j ) && include( j ) ).ToArray();
            if ( members.Length > 0 )
                groups.Add( (start, end, members) );
        }
        var legL = Role( BoneRole.UpperLegL );
        var legR = Role( BoneRole.UpperLegR );
        var armL = Role( BoneRole.ClavicleL ) is var cl and >= 0 ? cl : Role( BoneRole.UpperArmL );
        var armR = Role( BoneRole.ClavicleR ) is var cr and >= 0 ? cr : Role( BoneRole.UpperArmR );
        var limbs = new[] { legL, legR, armL, armR };
        Group( legL, Role( BoneRole.ToeL ) is var tl and >= 0 ? tl : Role( BoneRole.FootL ), j => Under( bones, j, legL ) );
        Group( legR, Role( BoneRole.ToeR ) is var tr and >= 0 ? tr : Role( BoneRole.FootR ), j => Under( bones, j, legR ) );
        Group( armL, Role( BoneRole.HandL ), j => Under( bones, j, armL ) );
        Group( armR, Role( BoneRole.HandR ), j => Under( bones, j, armR ) );
        Group( root, Role( BoneRole.Neck ), j => !limbs.Any( l => l >= 0 && Under( bones, j, l ) ) );

        var sigma = Settings.RecoverSigma;
        var hold = Settings.RecoverHold;
        var radius = (int)MathF.Ceiling( 3 * sigma );
        var rebuiltWorld = frames.Select( f => World( bones, f ) ).ToArray();
        var capture = 0f;
        foreach ( var (start, end, members) in groups )
        {
            var raw = new float[count];
            for ( var t = 0; t < count; t++ )
            {
                var d = Vector3.Distance( rebuiltWorld[t][end] - rebuiltWorld[t][start], captureWorld[t][end] - captureWorld[t][start] );
                raw[t] = Math.Clamp( (d - Settings.RecoverFrom) / Math.Max( .01f, Settings.RecoverFull - Settings.RecoverFrom ), 0, 1 );
            }
            var held = Enumerable.Range( 0, count ).Select( t => Enumerable.Range( Math.Max( 0, t - hold ), Math.Min( count, t + hold + 1 ) - Math.Max( 0, t - hold ) ).Max( k => raw[k] ) ).ToArray();
            for ( var t = 0; t < count; t++ )
            {
                float sum = 0, total = 0;
                for ( var k = -radius; k <= radius; k++ )
                {
                    var w = MathF.Exp( -k * k / (2 * sigma * sigma) );
                    sum += held[Math.Clamp( t + k, 0, count - 1 )] * w;
                    total += w;
                }
                var weight = Math.Min( 1f, sum / total );
                capture += weight;
                if ( weight < .01f )
                    continue;
                foreach ( var j in members )
                    frames[t][j] = new XForm( Vector3.Lerp( frames[t][j].Pos, target[t][j].Pos, weight ),
                        Quaternion.Normalize( Quaternion.Slerp( frames[t][j].Rot, target[t][j].Rot, weight ) ) );
            }
        }
        return groups.Count == 0 ? 0 : capture / (groups.Count * count);
    }

    /// <summary>The model keeps the pelvis more upright than a performer who leans into a move (measured on a dance: 9° typical,
    /// 23° at worst), which squares the whole body up against the video. Where the pelvis leans or turns more than
    /// <see cref="Options.TorsoFrom"/> away from the capture's, its rotation and the spine's are eased toward the capture's, fully at
    /// <see cref="Options.TorsoFull"/>, the weight smoothed over time. The pelvis position (the path) stays the model's.</summary>
    static void RecoverTorso( IReadOnlyList<MotionBone> bones, XForm[][] frames, XForm[][] target, HashSet<int> detail, int root )
    {
        var count = frames.Length;
        var limbs = new[] { BoneRole.UpperLegL, BoneRole.UpperLegR, BoneRole.ClavicleL, BoneRole.ClavicleR, BoneRole.UpperArmL, BoneRole.UpperArmR }
            .Select( r => bones.ToList().FindIndex( b => b.Role == r ) ).Where( i => i >= 0 ).ToArray();
        var spine = Enumerable.Range( 0, bones.Count ).Where( j => j != root && !detail.Contains( j ) && !limbs.Any( l => Under( bones, j, l ) ) ).ToArray();
        float Degrees( Quaternion a, Quaternion b ) => 2 * MathF.Acos( Math.Clamp( MathF.Abs( Quaternion.Dot( Quaternion.Normalize( a ), Quaternion.Normalize( b ) ) ), 0f, 1f ) ) * 180 / MathF.PI;
        var raw = Enumerable.Range( 0, count ).Select( t => Math.Clamp( (Degrees( frames[t][root].Rot, target[t][root].Rot ) - Settings.TorsoFrom)
            / Math.Max( .1f, Settings.TorsoFull - Settings.TorsoFrom ), 0, 1 ) ).ToArray();
        var sigma = Settings.RecoverSigma;
        var radius = (int)MathF.Ceiling( 3 * sigma );
        for ( var t = 0; t < count; t++ )
        {
            float sum = 0, total = 0;
            for ( var k = -radius; k <= radius; k++ )
            {
                var w = MathF.Exp( -k * k / (2 * sigma * sigma) );
                sum += raw[Math.Clamp( t + k, 0, count - 1 )] * w;
                total += w;
            }
            var weight = Math.Min( 1f, sum / total );
            if ( weight < .01f )
                continue;
            frames[t][root] = new XForm( frames[t][root].Pos, Quaternion.Normalize( Quaternion.Slerp( frames[t][root].Rot, target[t][root].Rot, weight ) ) );
            foreach ( var j in spine )
                frames[t][j] = new XForm( Vector3.Lerp( frames[t][j].Pos, target[t][j].Pos, weight ),
                    Quaternion.Normalize( Quaternion.Slerp( frames[t][j].Rot, target[t][j].Rot, weight ) ) );
        }
    }

    /// <summary>A foot counts as planted in the capture within this height (metres) of its lowest and below this horizontal speed (m/s).</summary>
    const float PlantHeight = .03f, PlantSpeed = .2f;
    /// <summary>Frames over which a plant lock eases in and out (measured: over 3 frames the foot snapped onto its spot at 40 cm/s),
    /// and the longest lift inside a plant that is bridged.</summary>
    const int PlantBlend = 6, PlantBridge = 3;

    /// <summary>The model moves its own legs between key poses, and its hips wander a few centimetres from the capture's, so a foot the
    /// performer kept planted slid (measured on a floss dance: 3 cm/s in the capture, 50 cm/s rebuilt). Wherever the capture's foot is
    /// planted, the rebuilt leg is re-solved (two-bone IK, the knee kept on its side) so the ankle sits where the capture's does and the
    /// foot and toe take the capture's angle, easing in and out over <see cref="PlantBlend"/> frames. Everything above the feet stays
    /// the model's.</summary>
    static void PlantFeet( IReadOnlyList<MotionBone> bones, XForm[][] frames, XForm[][] capture, Vector3[][] captureWorld )
    {
        var count = frames.Length;
        var chains = new[]
        {
            (BoneRole.UpperLegL, BoneRole.LowerLegL, BoneRole.FootL, BoneRole.ToeL),
            (BoneRole.UpperLegR, BoneRole.LowerLegR, BoneRole.FootR, BoneRole.ToeR),
        };
        int Role( BoneRole r ) => bones.ToList().FindIndex( b => b.Role == r );
        var weights = chains.Select( c => PlantWeights( captureWorld, Role( c.Item3 ) ) ).ToArray();
        // A planted leg can only reach the capture's spot if the hips are near where the capture had them (the model's hips
        // wander up to 10-20 cm): while either foot is planted the hips ease to the capture's position.
        var root = bones.ToList().FindIndex( b => b.Parent < 0 );
        if ( root >= 0 )
            for ( var t = 0; t < count; t++ )
            {
                var w = Math.Max( weights[0]?[t] ?? 0, weights[1]?[t] ?? 0 );
                if ( w > 0 )
                    frames[t][root] = new XForm( Vector3.Lerp( frames[t][root].Pos, capture[t][root].Pos, w ), frames[t][root].Rot );
            }
        for ( var c = 0; c < chains.Length; c++ )
        {
            var (thighRole, shinRole, footRole, toeRole) = chains[c];
            int thigh = Role( thighRole ), shin = Role( shinRole ), foot = Role( footRole ), toe = Role( toeRole );
            if ( weights[c] is not { } plant || thigh < 0 || shin < 0 || bones[shin].Parent != thigh || bones[foot].Parent != shin )
                continue;
            for ( var t = 0; t < count; t++ )
            {
                var weight = plant[t];
                if ( weight <= 0 )
                    continue;
                var world = WorldTransforms( bones, frames[t] );
                var target = Vector3.Lerp( world[foot].Pos, captureWorld[t][foot], weight );
                Vector3 s = world[thigh].Pos, e = world[shin].Pos, w = world[foot].Pos;
                float upper = Vector3.Distance( s, e ), lower = Vector3.Distance( e, w );
                if ( upper < 1e-4f || lower < 1e-4f )
                    continue;
                var toGoal = target - s;
                var d = Math.Clamp( toGoal.Length(), MathF.Abs( upper - lower ) + 1e-4f, upper + lower - 1e-4f );
                var direction = Vector3.Normalize( toGoal );
                var pole = e - s;
                pole -= direction * Vector3.Dot( pole, direction );
                if ( pole.LengthSquared() < 1e-10f )
                    continue;
                pole = Vector3.Normalize( pole );
                var along = (upper * upper - lower * lower + d * d) / (2 * d);
                var knee = s + direction * along + pole * MathF.Sqrt( Math.Max( 0, upper * upper - along * along ) );
                var ankle = s + direction * d;
                var bend = Between( e - s, knee - s );
                var thighWorld = Quaternion.Normalize( bend * world[thigh].Rot );
                var shinTurned = Quaternion.Normalize( bend * world[shin].Rot );
                var shinWorld = Quaternion.Normalize( Between( Vector3.Transform( w - e, bend ), ankle - knee ) * shinTurned );
                // The foot takes the capture's angle as the lock comes in.
                var captureFoot = WorldTransforms( bones, capture[t] )[foot].Rot;
                var footWorld = Quaternion.Normalize( Quaternion.Slerp( world[foot].Rot, captureFoot, weight ) );
                var parentWorld = bones[thigh].Parent < 0 ? Quaternion.Identity : world[bones[thigh].Parent].Rot;
                frames[t][thigh] = new XForm( frames[t][thigh].Pos, Quaternion.Normalize( Quaternion.Conjugate( parentWorld ) * thighWorld ) );
                frames[t][shin] = new XForm( frames[t][shin].Pos, Quaternion.Normalize( Quaternion.Conjugate( thighWorld ) * shinWorld ) );
                frames[t][foot] = new XForm( frames[t][foot].Pos, Quaternion.Normalize( Quaternion.Conjugate( shinWorld ) * footWorld ) );
                if ( toe >= 0 && bones[toe].Parent == foot )
                    frames[t][toe] = new XForm( frames[t][toe].Pos, Quaternion.Normalize( Quaternion.Slerp( frames[t][toe].Rot, capture[t][toe].Rot, weight ) ) );
            }
        }
    }

    /// <summary>Per frame, how firmly the capture plants <paramref name="foot"/>: 1 while it is low and still, easing to 0 over
    /// <see cref="PlantBlend"/> frames either side; lifts up to <see cref="PlantBridge"/> frames inside a plant are bridged. Null when there is no such bone.</summary>
    static float[] PlantWeights( Vector3[][] captureWorld, int foot )
    {
        if ( foot < 0 )
            return null;
        var count = captureWorld.Length;
        var heights = captureWorld.Select( w => w[foot].Y ).OrderBy( y => y ).ToArray();
        var lowest = heights[Math.Min( heights.Length - 1, heights.Length / 20 )];
        var planted = new bool[count];
        for ( var t = 0; t < count; t++ )
        {
            int a = Math.Max( 0, t - 1 ), b = Math.Min( count - 1, t + 1 );
            var step = captureWorld[b][foot] - captureWorld[a][foot];
            var speed = MathF.Sqrt( step.X * step.X + step.Z * step.Z ) * Fps / Math.Max( 1, b - a );
            planted[t] = captureWorld[t][foot].Y < lowest + PlantHeight && speed < PlantSpeed;
        }
        var bridged = (bool[])planted.Clone();
        for ( var t = 0; t < count; t++ )
            if ( !planted[t] )
            {
                var before = Enumerable.Range( 1, PlantBridge ).Any( k => t - k >= 0 && planted[t - k] );
                var after = Enumerable.Range( 1, PlantBridge ).Any( k => t + k < count && planted[t + k] );
                if ( before && after )
                    bridged[t] = true;
            }
        var result = new float[count];
        for ( var t = 0; t < count; t++ )
        {
            if ( bridged[t] ) { result[t] = 1; continue; }
            var distance = Enumerable.Range( 1, PlantBlend ).FirstOrDefault( k => (t - k >= 0 && bridged[t - k]) || (t + k < count && bridged[t + k]) );
            if ( distance == 0 )
                continue;
            var w = 1f - distance / (PlantBlend + 1f);
            result[t] = w * w * (3 - 2 * w);
        }
        return result;
    }

    static XForm[] WorldTransforms( IReadOnlyList<MotionBone> bones, XForm[] pose )
    {
        var world = new XForm[bones.Count];
        for ( var j = 0; j < bones.Count; j++ )
            world[j] = bones[j].Parent < 0 ? pose[j] : XForm.Compose( world[bones[j].Parent], pose[j] );
        return world;
    }

    static Quaternion Between( Vector3 from, Vector3 to )
    {
        from = Vector3.Normalize( from );
        to = Vector3.Normalize( to );
        var dot = Vector3.Dot( from, to );
        if ( dot > .999999f )
            return Quaternion.Identity;
        if ( dot < -.999999f )
        {
            var axis = Vector3.Cross( Vector3.UnitX, from );
            if ( axis.LengthSquared() < 1e-6f )
                axis = Vector3.Cross( Vector3.UnitY, from );
            return Quaternion.CreateFromAxisAngle( Vector3.Normalize( axis ), MathF.PI );
        }
        var c = Vector3.Cross( from, to );
        return Quaternion.Normalize( new Quaternion( c.X, c.Y, c.Z, 1 + dot ) );
    }

    /// <summary>Degrees a leg must swing from hanging straight down for its peak to count as a kick, and the least shortfall corrected.</summary>
    const float KickDegrees = 30, KickShortfallDegrees = 3;
    /// <summary>Frames either side of a kick's peak over which the correction eases in and out.</summary>
    const int KickReach = 6;
    /// <summary>Frames either side of the capture's kick peak searched for the AI's own peak.</summary>
    const int KickSearch = 4;

    /// <summary>The model tends to kick lower than the performer (its training motion and the G1's hips). For each peak of a leg's swing
    /// in the capture (the leg's angle from straight down, at least <see cref="KickDegrees"/>), the AI's own peak nearby is found and its
    /// thigh raised within its own swing plane by the missing angle, fully at the AI's peak and eased (Hann) over <see cref="KickReach"/>
    /// frames either side. The timing, direction and knee bend stay the model's; only the height changes. Returns the kicks corrected.</summary>
    static int KeepKickHeight( IReadOnlyList<MotionBone> bones, XForm[][] frames, Vector3[][] captureWorld )
    {
        var count = frames.Length;
        var corrected = 0;
        foreach ( var (thighRole, footRole) in new[] { (BoneRole.UpperLegL, BoneRole.FootL), (BoneRole.UpperLegR, BoneRole.FootR) } )
        {
            var thigh = bones.ToList().FindIndex( b => b.Role == thighRole );
            var foot = bones.ToList().FindIndex( b => b.Role == footRole );
            if ( thigh < 0 || foot < 0 )
                continue;
            float Swing( Vector3[] w ) => MathF.Acos( Math.Clamp( Vector3.Dot( Vector3.Normalize( w[foot] - w[thigh] ), -Vector3.UnitY ), -1f, 1f ) ) * 180 / MathF.PI;
            var captured = captureWorld.Select( Swing ).ToArray();
            // Each correction: the AI's peak frame and the world rotation raising its leg by the shortfall.
            var rebuilt = frames.Select( f => World( bones, f ) ).ToArray();
            var swing = rebuilt.Select( Swing ).ToArray();
            var peaks = new List<(int Frame, Quaternion Turn)>();
            for ( var t = 2; t < count - 2; t++ )
            {
                if ( captured[t] < KickDegrees || captured[t] < captured[t - 1] || captured[t] < captured[t + 1] || captured[t] < captured[t - 2] || captured[t] < captured[t + 2] )
                    continue;
                var own = t;
                for ( var k = Math.Max( 0, t - KickSearch ); k <= Math.Min( count - 1, t + KickSearch ); k++ )
                    if ( swing[k] > swing[own] )
                        own = k;
                var shortfall = captured[t] - swing[own];
                if ( shortfall < KickShortfallDegrees || peaks.Any( p => Math.Abs( p.Frame - own ) <= KickReach ) )
                    continue;
                var leg = Vector3.Normalize( rebuilt[own][foot] - rebuilt[own][thigh] );
                var axis = Vector3.Cross( -Vector3.UnitY, leg );
                if ( axis.LengthSquared() < 1e-6f )
                    continue;
                peaks.Add( (own, Quaternion.CreateFromAxisAngle( Vector3.Normalize( axis ), shortfall * MathF.PI / 180 )) );
            }
            foreach ( var (peak, turn) in peaks )
            {
                corrected++;
                for ( var t = Math.Max( 0, peak - KickReach ); t <= Math.Min( count - 1, peak + KickReach ); t++ )
                {
                    var weight = .5f + .5f * MathF.Cos( MathF.PI * (t - peak) / (KickReach + 1f) );
                    weight *= weight * (3 - 2 * weight);
                    var partial = Quaternion.Slerp( Quaternion.Identity, turn, weight );
                    // World turn of the thigh, written back as a local rotation under its (unchanged) parent.
                    var parentWorld = Quaternion.Identity;
                    for ( var b = bones[thigh].Parent; b >= 0; b = bones[b].Parent )
                        parentWorld = Quaternion.Normalize( frames[t][b].Rot * parentWorld );
                    var worldThigh = Quaternion.Normalize( parentWorld * frames[t][thigh].Rot );
                    var local = Quaternion.Normalize( Quaternion.Conjugate( parentWorld ) * Quaternion.Normalize( partial * worldThigh ) );
                    frames[t][thigh] = new XForm( frames[t][thigh].Pos, local );
                }
            }
        }
        return corrected;
    }

    /// <summary>Frames either side of a join that are eased.</summary>
    const float JoinSigma = 2f, JoinReach = 4f;

    /// <summary>Each in-between is time-warped onto its own gap, so neighbouring ones play at different speeds and the body
    /// hitches where they meet (measured: acceleration at the joins 2–3× that inside the in-betweens). Around every join the body
    /// bones and the hips' path are smoothed (Gaussian, <see cref="JoinSigma"/> frames), fading out over <see cref="JoinReach"/>
    /// frames, so the generated motion elsewhere is left as the model made it.</summary>
    static void SmoothJoins( IReadOnlyList<MotionBone> bones, XForm[][] frames, IReadOnlyList<int> anchors, HashSet<int> detail, int root )
    {
        var count = frames.Length;
        var joins = anchors.Where( k => k > 0 && k < count - 1 ).ToArray();
        if ( joins.Length == 0 )
            return;
        var source = frames.Select( f => (XForm[])f.Clone() ).ToArray();
        var radius = (int)MathF.Ceiling( 3 * JoinSigma );
        for ( var t = 1; t < count - 1; t++ )
        {
            var distance = joins.Min( k => Math.Abs( k - t ) );
            if ( distance > JoinReach )
                continue;
            var weight = MathF.Exp( -distance * distance / (JoinReach * JoinReach / 2) );
            for ( var j = 0; j < bones.Count; j++ )
            {
                if ( j != root && detail.Contains( j ) )
                    continue;
                // Weighted average: positions directly, rotations as a sign-aligned normalized sum (close rotations).
                var position = Vector3.Zero;
                var rotation = new Quaternion( 0, 0, 0, 0 );
                var reference = source[t][j].Rot;
                var total = 0f;
                for ( var k = -radius; k <= radius; k++ )
                {
                    var f = Math.Clamp( t + k, 0, count - 1 );
                    var w = MathF.Exp( -k * k / (2 * JoinSigma * JoinSigma) );
                    var q = source[f][j].Rot;
                    if ( Quaternion.Dot( q, reference ) < 0 )
                        q = Quaternion.Negate( q );
                    position += source[f][j].Pos * w;
                    rotation += q * w;
                    total += w;
                }
                var smoothed = new XForm( position / total, Quaternion.Normalize( rotation ) );
                frames[t][j] = new XForm( Vector3.Lerp( source[t][j].Pos, smoothed.Pos, weight ),
                    Quaternion.Normalize( Quaternion.Slerp( source[t][j].Rot, smoothed.Rot, weight ) ) );
            }
        }
    }

    /// <summary>Spreads the difference between the generated and the captured pose at both ends over the in-between (per bone,
    /// parent-local), so the key poses are reproduced exactly and nothing pops.</summary>
    static void HitKeys( XForm[][] frames, XForm[] first, XForm[] last )
    {
        var n = frames.Length - 1;
        for ( var bone = 0; bone < first.Length; bone++ )
        {
            var rotA = first[bone].Rot * Quaternion.Conjugate( frames[0][bone].Rot );
            var rotB = last[bone].Rot * Quaternion.Conjugate( frames[n][bone].Rot );
            var posA = first[bone].Pos - frames[0][bone].Pos;
            var posB = last[bone].Pos - frames[n][bone].Pos;
            for ( var t = 0; t <= n; t++ )
            {
                var s = n == 0 ? 0f : t / (float)n;
                s = s * s * (3f - 2f * s);
                var correction = Quaternion.Normalize( Quaternion.Slerp( rotA, rotB, s ) );
                frames[t][bone] = new XForm( frames[t][bone].Pos + Vector3.Lerp( posA, posB, s ), Quaternion.Normalize( correction * frames[t][bone].Rot ) );
            }
        }
    }

    /// <summary>Key poses: the first and last frame, and in between the frame where the body moves least within
    /// <see cref="MinimumGap"/>–<see cref="MaximumGap"/> frames of the previous key (a settle, a turn-around, an extreme).</summary>
    public static List<int> Anchors( IReadOnlyList<MotionBone> bones, IReadOnlyList<XForm[]> poses )
    {
        var count = poses.Count;
        var body = Enumerable.Range( 0, bones.Count ).Where( j => bones[j].Group != "fingers" ).ToArray();
        var world = poses.Select( p => World( bones, p ) ).ToArray();
        var speed = new float[count];
        for ( var k = 0; k < count; k++ )
        {
            int a = Math.Max( 0, k - 1 ), b = Math.Min( count - 1, k + 1 );
            speed[k] = body.Sum( j => Vector3.Distance( world[a][j], world[b][j] ) );
        }
        var anchors = new List<int> { 0 };
        var last = 0;
        while ( count - 1 - last > MaximumGap )
        {
            // Leave room for a whole in-between after the next key.
            var from = last + MinimumGap;
            var to = Math.Min( last + MaximumGap, count - 1 - MinimumGap );
            if ( to < from )
                break;
            // The stillest frame at a length the model plays at its own timing; any frame if none fits.
            var next = -1;
            for ( var k = from; k <= to; k++ )
                if ( Warp( k - last ) == 0 && (next < 0 || speed[k] < speed[next]) )
                    next = k;
            if ( next < 0 )
            {
                next = from;
                for ( var k = from + 1; k <= to; k++ )
                    if ( speed[k] < speed[next] )
                        next = k;
            }
            anchors.Add( next );
            last = next;
        }
        anchors.Add( count - 1 );
        return anchors;
    }

    /// <summary>Bones kept from the capture: fingers, neck and head (the model has none), and the wrists. The model's G1 wrist is three
    /// robot hinges it turns freely (measured: the rebuilt wrist turned nearly twice as fast as the capture's and wobbled 3×), while the
    /// capture's wrist was fitted to the hand seen in the video. The hand still follows the model's forearm; only its bend at the wrist
    /// is the capture's.</summary>
    static HashSet<int> DetailBones( IReadOnlyList<MotionBone> bones )
    {
        var neck = bones.ToList().FindIndex( b => b.Role == BoneRole.Neck );
        var head = bones.ToList().FindIndex( b => b.Role == BoneRole.Head );
        var hands = bones.Select( ( b, i ) => (b, i) ).Where( x => x.b.Role is BoneRole.HandL or BoneRole.HandR ).Select( x => x.i ).ToArray();
        var result = new HashSet<int>();
        for ( var j = 0; j < bones.Count; j++ )
            if ( Under( bones, j, neck ) || Under( bones, j, head ) || hands.Any( h => Under( bones, j, h ) ) || bones[j].Group == "fingers" )
                result.Add( j );
        return result;
    }

    static bool Under( IReadOnlyList<MotionBone> bones, int bone, int ancestor )
    {
        if ( ancestor < 0 )
            return false;
        for ( var b = bone; b >= 0; b = bones[b].Parent )
            if ( b == ancestor )
                return true;
        return false;
    }

    /// <summary>The capture's pose at <paramref name="time"/>, per bone parent-local, interpolated between its frames.</summary>
    static XForm[] Sample( MotionDocument doc, double time )
    {
        var frames = doc.Frames;
        var lo = 0;
        var hi = frames.Count - 1;
        while ( lo < hi )
        {
            var mid = (lo + hi + 1) / 2;
            if ( frames[mid].Time <= time ) lo = mid; else hi = mid - 1;
        }
        var a = frames[lo];
        var b = frames[Math.Min( lo + 1, frames.Count - 1 )];
        var t = b.Time > a.Time ? (float)Math.Clamp( (time - a.Time) / (b.Time - a.Time), 0, 1 ) : 0f;
        var pose = new XForm[doc.Bones.Count];
        for ( var j = 0; j < pose.Length; j++ )
            pose[j] = new XForm( Vector3.Lerp( MotionDocument.V( a.Positions[j] ), MotionDocument.V( b.Positions[j] ), t ),
                Quaternion.Normalize( Quaternion.Slerp( MotionDocument.Q( a.Rotations[j] ), MotionDocument.Q( b.Rotations[j] ), t ) ) );
        return pose;
    }

    /// <summary>Turns a capture that stays relative to a tilted camera upright (the body network's gravity estimate); identity otherwise.</summary>
    static Quaternion Level( MotionDocument doc )
    {
        if ( doc.Space != MotionSpace.CameraRelative || doc.Cameras.FirstOrDefault( c => c.Id == "video" )?.Up is not { Length: 3 } values )
            return Quaternion.Identity;
        var up = MotionDocument.V( values );
        if ( !(up.LengthSquared() > 1e-6f) )
            return Quaternion.Identity;
        up = Vector3.Normalize( up );
        var axis = Vector3.Cross( up, Vector3.UnitY );
        if ( axis.LengthSquared() < 1e-8f )
            return Quaternion.Identity;
        return Quaternion.CreateFromAxisAngle( Vector3.Normalize( axis ), MathF.Acos( Math.Clamp( up.Y, -1f, 1f ) ) );
    }

    /// <summary>The ground: the height the feet and toes reach in the lowest 2% of frames.</summary>
    static float Floor( IReadOnlyList<MotionBone> bones, IReadOnlyList<XForm[]> poses )
    {
        var feet = Enumerable.Range( 0, bones.Count ).Where( j => bones[j].Role is BoneRole.FootL or BoneRole.FootR or BoneRole.ToeL or BoneRole.ToeR ).ToArray();
        var lows = poses.Select( p => { var w = World( bones, p ); return feet.Min( j => w[j].Y ); } ).OrderBy( y => y ).ToArray();
        return lows[Math.Min( lows.Length - 1, lows.Length / 50 )];
    }

    static Vector3[] World( IReadOnlyList<MotionBone> bones, XForm[] pose )
    {
        var world = new XForm[bones.Count];
        var result = new Vector3[bones.Count];
        for ( var j = 0; j < bones.Count; j++ )
        {
            world[j] = bones[j].Parent < 0 ? pose[j] : XForm.Compose( world[bones[j].Parent], pose[j] );
            result[j] = world[j].Pos;
        }
        return result;
    }
}