Editor/AutoRig/JointAdjustPanel.cs

Editor UI panel for adjusting a model rig. Displays the joint tree, supports selecting/renaming joints, nudging joints with optional L/R mirroring, adding/deleting children, live incremental re-skin for organic rigs, full re-skin, undo stack, and emits events for selection, status and rig changes.

File AccessNetworking
using System;
using System.Collections.Generic;
using System.Linq;
using AutoRig.Rig;
using AutoRig.Solve.Organic;
using AutoRig.Voxel;
using Editor;
using Sandbox;
using VecN = System.Numerics.Vector3;

namespace AutoRig.Editor;

/// <summary>
/// Rig adjustment (spec §6): joint tree with selection, rename, XYZ nudges (auto-
/// mirrored across the detected symmetry plane for _L/_R twins), delete/add child,
/// re-skin (organic), reset-to-auto, and an undo stack. Every mutation snapshots
/// first; nothing here can corrupt the rig (mutations re-validate and roll back on
/// failure).
/// </summary>
public sealed class JointAdjustPanel : Widget
{
    ModelEntry _entry;
    int _selected = -1;
    bool _mirror = true;
    OrganicSkinCache _skinCache;

    readonly Stack<(RigSkeleton Skeleton, SkinWeights Weights)> _undo = new();

    Layout _treeLayout;
    LineEdit _renameEdit;
    Label _header;

    /// <summary>Raised after any rig mutation (the window refreshes the preview).</summary>
    public Action RigChanged { get; set; }

    /// <summary>Raised when the user selects a joint (name), for preview highlighting.</summary>
    public Action<string> SelectionChanged { get; set; }

    /// <summary>Raised with a status message + tone for the window's status bar.</summary>
    public Action<string, Color> Status { get; set; }

    public JointAdjustPanel( Widget parent ) : base( parent )
    {
        Layout = Layout.Column();
        Layout.Spacing = UiSpacing.Related;

        _header = new Label( "Adjust", this );
        _header.SetStyles(
            $"font-weight: 600; color: {Theme.Blue.Hex}; "
            + $"margin-top: {UiSpacing.HeaderTop}px; margin-bottom: {UiSpacing.HeaderBottom}px;" );
        Layout.Add( _header );

        var scroll = new ScrollArea( this );
        scroll.Canvas = new Widget( scroll );
        scroll.Canvas.Layout = Layout.Column();
        scroll.Canvas.Layout.Spacing = 2;
        scroll.MinimumHeight = 120;
        _treeLayout = scroll.Canvas.Layout;
        Layout.Add( scroll, 1 );

        // ---- rename ----
        var renameRow = Layout.AddRow();
        renameRow.Spacing = UiSpacing.Related;
        _renameEdit = new LineEdit( this ) { PlaceholderText = "joint name" };
        renameRow.Add( _renameEdit, 1 );
        var renameButton = new Button( "Rename", this );
        renameButton.Clicked = Rename;
        renameRow.Add( renameButton );

        // ---- nudge ----
        var nudgeRow = Layout.AddRow();
        nudgeRow.Spacing = UiSpacing.Related;
        foreach ( var (label, delta) in new (string, VecN)[]
        {
            ("-X", new(-1, 0, 0)), ("+X", new(1, 0, 0)),
            ("-Y", new(0, -1, 0)), ("+Y", new(0, 1, 0)),
            ("-Z", new(0, 0, -1)), ("+Z", new(0, 0, 1)),
        } )
        {
            var captured = delta;
            var button = new Button( label, this );
            button.Clicked = () => Nudge( captured );
            nudgeRow.Add( button );
        }
        nudgeRow.AddStretchCell();

        // ---- structure + weights ----
        var structureRow = Layout.AddRow();
        structureRow.Spacing = UiSpacing.Related;
        var addButton = new Button( "Add Child", "add", this );
        addButton.Clicked = AddChild;
        structureRow.Add( addButton );
        var deleteButton = new Button( "Delete", "delete", this );
        deleteButton.Clicked = DeleteSelected;
        structureRow.Add( deleteButton );
        var mirrorCheck = new Checkbox( "Mirror L/R", this ) { Value = true };
        mirrorCheck.Toggled = () => _mirror = mirrorCheck.Value;
        structureRow.Add( mirrorCheck );
        structureRow.AddStretchCell();

        var actionRow = Layout.AddRow();
        actionRow.Spacing = UiSpacing.Related;
        var reskinButton = new Button( "Re-skin", "healing", this );
        reskinButton.Clicked = Reskin;
        actionRow.Add( reskinButton );
        var undoButton = new Button( "Undo", "undo", this );
        undoButton.Clicked = Undo;
        actionRow.Add( undoButton );
        var resetButton = new Button( "Reset to auto", "restart_alt", this );
        resetButton.Clicked = () => ResetRequested?.Invoke();
        actionRow.Add( resetButton );
        actionRow.AddStretchCell();
    }

    /// <summary>Raised by "Reset to auto" (the window re-runs the solver).</summary>
    public Action ResetRequested { get; set; }

    /// <summary>Rebuilds the joint tree for the CURRENT entry (rig arrived/changed).</summary>
    public void Refresh() => RebuildTree();

    public void SetEntry( ModelEntry entry )
    {
        _entry = entry;
        _selected = -1;
        _undo.Clear();
        _skinCache = null;
        RebuildTree();
    }

    void RebuildTree()
    {
        _treeLayout.Clear( true );
        if ( _entry?.Rig is null )
        {
            var empty = new Label( "Rig a model to adjust its joints.", this );
            empty.SetStyles( $"color: {Theme.TextLight.Hex}; margin: {UiSpacing.Related}px;" );
            _treeLayout.Add( empty );
            return;
        }

        var joints = _entry.Rig.Skeleton.Joints;
        var depth = new int[joints.Count];
        for ( var i = 0; i < joints.Count; i++ )
            depth[i] = joints[i].Parent < 0 ? 0 : depth[joints[i].Parent] + 1;

        for ( var i = 0; i < joints.Count; i++ )
        {
            var index = i;
            var selected = i == _selected;
            var label = new Label( joints[i].Name, this );
            label.SetStyles(
                $"margin-left: {depth[i] * UiSpacing.Group + UiSpacing.Related}px; "
                + $"color: {(selected ? Theme.Yellow.Hex : Theme.TextLight.Hex)}; "
                + $"font-weight: {(selected ? 600 : 400)};" );
            label.MouseClick = () =>
            {
                _selected = index;
                _renameEdit.Text = joints[index].Name;
                SelectionChanged?.Invoke( joints[index].Name );
                RebuildTree();
            };
            _treeLayout.Add( label );
        }
    }

    // ================================================================== mutations

    void Snapshot()
    {
        var skeleton = new RigSkeleton();
        foreach ( var j in _entry.Rig.Skeleton.Joints )
            skeleton.Joints.Add( new RigJoint
            {
                Name = j.Name,
                Parent = j.Parent,
                Position = j.Position,
                HingeAxis = j.HingeAxis,
            } );
        var weights = new SkinWeights
        {
            BoneIndices = (int[])_entry.Rig.Weights.BoneIndices.Clone(),
            Weights = (float[])_entry.Rig.Weights.Weights.Clone(),
        };
        _undo.Push( (skeleton, weights) );
    }

    void Apply( RigSkeleton skeleton, SkinWeights weights, string action )
    {
        try
        {
            skeleton.Validate();
            weights.Validate( _entry.Mesh, skeleton );
        }
        catch ( Exception e )
        {
            Status?.Invoke( $"{action} rejected: {e.Message}", Theme.Red );
            return;
        }
        _entry.Rig = new RigResult
        {
            Skeleton = skeleton,
            Weights = weights,
            SolverName = _entry.Rig.SolverName,
            Degraded = _entry.Rig.Degraded,
            Explanation = $"{_entry.Rig.Explanation} (edited)".Replace( " (edited) (edited)", " (edited)" ),
        };
        RebuildTree();
        RigChanged?.Invoke();
        Status?.Invoke( action, Theme.Green );
    }

    void Undo()
    {
        if ( _entry?.Rig is null || _undo.Count == 0 )
        {
            Status?.Invoke( "Nothing to undo.", Theme.Yellow );
            return;
        }
        var (skeleton, weights) = _undo.Pop();
        _skinCache = null;   // cached fields reflect the pre-undo joint positions
        _entry.Rig = new RigResult
        {
            Skeleton = skeleton,
            Weights = weights,
            SolverName = _entry.Rig.SolverName,
            Degraded = _entry.Rig.Degraded,
            Explanation = _entry.Rig.Explanation,
        };
        _selected = Math.Min( _selected, skeleton.Joints.Count - 1 );
        RebuildTree();
        RigChanged?.Invoke();
        Status?.Invoke( "Undone.", Theme.Green );
    }

    void Rename()
    {
        if ( _entry?.Rig is null || _selected < 0 )
            return;
        var name = _renameEdit.Text?.Trim();
        if ( string.IsNullOrEmpty( name ) )
            return;
        Snapshot();
        _entry.Rig.Skeleton.Joints[_selected].Name = name;
        try
        {
            _entry.Rig.Skeleton.Validate();
        }
        catch ( Exception e )
        {
            Undo();
            _undo.TryPop( out _ ); // drop the snapshot the failed attempt pushed
            Status?.Invoke( $"Rename rejected: {e.Message}", Theme.Red );
            return;
        }
        RebuildTree();
        RigChanged?.Invoke();
    }

    void Nudge( VecN direction )
    {
        if ( _entry?.Rig is null || _selected < 0 )
            return;
        var step = _entry.Mesh.ComputeBounds().Size.Length() / 200f;
        Snapshot();

        var joints = _entry.Rig.Skeleton.Joints;
        joints[_selected].Position += direction * step;
        var changed = new List<int> { _selected };

        // Mirror the nudge onto the _L/_R twin across the symmetry plane.
        if ( _mirror && _entry.Analysis?.Symmetry is { } plane )
        {
            var twin = TwinOf( _selected );
            if ( twin >= 0 )
            {
                joints[twin].Position = CategoryDetector.Mirror( joints[_selected].Position, plane );
                changed.Add( twin );
            }
        }
        LiveReskin( changed );
        RebuildTree();
        RigChanged?.Invoke();
    }

    /// <summary>
    /// Incremental re-skin after joints moved (organic rigs): only the moved bones'
    /// geodesic fields recompute, so dragging stays interactive. The first move
    /// pays the one-time cache build. Failures keep the previous weights (the
    /// Re-skin button remains the exact full pass).
    /// </summary>
    void LiveReskin( List<int> movedJoints )
    {
        if ( _entry?.Rig is null || _entry.Rig.SolverName != "organic" )
            return;
        try
        {
            var skeleton = _entry.Rig.Skeleton;
            if ( _skinCache is null || !_skinCache.Matches( _entry.Mesh, skeleton ) )
            {
                var grid = VoxelGrid.Build( _entry.Mesh, 64 );
                _skinCache = OrganicSkinner.BuildCache( _entry.Mesh, grid, skeleton );
            }

            var affected = new HashSet<int>();
            foreach ( var joint in movedJoints )
                foreach ( var bone in OrganicSkinner.BonesAffectedByJoint( skeleton, joint ) )
                    affected.Add( bone );

            var weights = OrganicSkinner.ReskinBones( _skinCache, skeleton, affected );
            weights.Validate( _entry.Mesh, skeleton );
            _entry.Rig = new RigResult
            {
                Skeleton = skeleton,
                Weights = weights,
                SolverName = _entry.Rig.SolverName,
                Degraded = _entry.Rig.Degraded,
                Explanation = _entry.Rig.Explanation,
            };
            Status?.Invoke( "Moved and re-skinned live.", Theme.Green );
        }
        catch ( Exception e )
        {
            _skinCache = null;
            Status?.Invoke( $"Live re-skin skipped: {e.Message}", Theme.Yellow );
        }
    }

    int TwinOf( int joint )
    {
        var name = _entry.Rig.Skeleton.Joints[joint].Name;
        string twinName = name.EndsWith( "_L", StringComparison.Ordinal )
            ? name[..^2] + "_R"
            : name.EndsWith( "_R", StringComparison.Ordinal ) ? name[..^2] + "_L" : null;
        return twinName is null ? -1 : _entry.Rig.Skeleton.IndexOf( twinName );
    }

    void AddChild()
    {
        if ( _entry?.Rig is null || _selected < 0 )
            return;
        Snapshot();
        var joints = _entry.Rig.Skeleton.Joints;
        var parent = joints[_selected];
        var baseName = $"{parent.Name}_child";
        var name = baseName;
        var n = 2;
        while ( _entry.Rig.Skeleton.IndexOf( name ) >= 0 )
            name = $"{baseName}_{n++}";

        var offset = _entry.Mesh.ComputeBounds().Size.Length() / 40f;
        joints.Add( new RigJoint
        {
            Name = name,
            Parent = _selected,
            Position = parent.Position + new VecN( 0, offset, 0 ),
        } );
        Apply( _entry.Rig.Skeleton, _entry.Rig.Weights, $"Added '{name}'." );
    }

    void DeleteSelected()
    {
        if ( _entry?.Rig is null || _selected < 0 )
            return;
        var joints = _entry.Rig.Skeleton.Joints;
        if ( joints.Count <= 1 || joints[_selected].Parent < 0 )
        {
            Status?.Invoke( "The root joint cannot be deleted.", Theme.Yellow );
            return;
        }
        Snapshot();

        var removed = _selected;
        var parentOfRemoved = joints[removed].Parent;

        // Rebuild with the joint removed: children reparent, indices shift down.
        var skeleton = new RigSkeleton();
        var remap = new int[joints.Count];
        for ( var i = 0; i < joints.Count; i++ )
        {
            if ( i == removed )
            {
                remap[i] = remap[parentOfRemoved];
                continue;
            }
            var parent = joints[i].Parent;
            if ( parent == removed )
                parent = parentOfRemoved;
            remap[i] = skeleton.Joints.Count;
            skeleton.Joints.Add( new RigJoint
            {
                Name = joints[i].Name,
                Parent = parent < 0 ? -1 : remap[parent],
                Position = joints[i].Position,
                HingeAxis = joints[i].HingeAxis,
            } );
        }

        // Weights: influences of the removed bone go to its parent; others remap.
        var boneIndices = (int[])_entry.Rig.Weights.BoneIndices.Clone();
        for ( var k = 0; k < boneIndices.Length; k++ )
            boneIndices[k] = remap[boneIndices[k]];
        var weights = new SkinWeights
        {
            BoneIndices = boneIndices,
            Weights = (float[])_entry.Rig.Weights.Weights.Clone(),
        };

        _selected = -1;
        Apply( skeleton, weights, "Joint deleted (children reparented)." );
    }

    void Reskin()
    {
        if ( _entry?.Rig is null )
            return;
        if ( _entry.Rig.SolverName != "organic" )
        {
            Status?.Invoke( "Rigid binding follows its parts automatically - no re-skin needed.", Theme.Yellow );
            return;
        }
        Snapshot();
        try
        {
            var grid = VoxelGrid.Build( _entry.Mesh, 64 );
            // Rebuild the live cache alongside the exact pass so later drags stay in sync.
            _skinCache = OrganicSkinner.BuildCache( _entry.Mesh, grid, _entry.Rig.Skeleton );
            var weights = OrganicSkinner.SkinFromCache( _skinCache, _entry.Rig.Skeleton );
            Apply( _entry.Rig.Skeleton, weights, "Re-skinned." );
        }
        catch ( Exception e )
        {
            Undo();
            _undo.TryPop( out _ );
            Status?.Invoke( $"Re-skin failed: {e.Message}", Theme.Red );
        }
    }
}