streamingle-unity-utilities/Editor/Gaze/GazeCalibrationPreviewSession.cs

using System;
using System.Collections.Generic;
using UnityEditor;
using UnityEditor.SceneManagement;
using UnityEngine;
using UnityEngine.SceneManagement;

namespace Streamingle.Gaze.Editor
{
    /// <summary>
    /// Owns the temporary blendshape writes used by the calibration UI.
    /// The session never modifies the profile or scene and always attempts to
    /// restore the weights that were present when it started.
    /// </summary>
    internal sealed class GazeCalibrationPreviewSession : IDisposable
    {
        private int[] blendShapeIndices = Array.Empty<int>();
        private string[] resolvedBlendShapeNames = Array.Empty<string>();
        private float[] originalWeights = Array.Empty<float>();
        private float[] previewWeights = Array.Empty<float>();
        private Mesh sourceMesh;
        private bool hasPreview;
        private bool eventsSubscribed;
        private bool driverOverrideAcquired;
        private BlendshapeGazeProfile driverProfile;

        internal event Action Stopped;

        public bool IsActive { get; private set; }

        public BlendshapeGazeDriver Driver { get; private set; }

        public SkinnedMeshRenderer Renderer { get; private set; }

        public BlendshapeGazeProfile Profile { get; private set; }

        public Vector2 PreviewCoordinates { get; private set; }

        public bool IsFreePreview { get; private set; }

        public bool Start(BlendshapeGazeDriver driver, out string error)
        {
            return Start(driver, null, out error);
        }

        internal bool Start(
            BlendshapeGazeDriver driver,
            BlendshapeGazeProfile previewProfile,
            out string error)
        {
            Stop();
            error = null;

            if (Application.isPlaying || EditorApplication.isPlayingOrWillChangePlaymode)
            {
                error = "Calibration preview is only available in Edit Mode.";
                return false;
            }

            if (driver == null)
            {
                error = "A Blendshape Gaze Driver is required.";
                return false;
            }

            var renderer = driver.TargetRenderer;
            if (renderer == null)
            {
                error = "The driver does not have a Target Renderer.";
                return false;
            }

            var mesh = renderer.sharedMesh;
            if (mesh == null)
            {
                error = "The Target Renderer does not have a shared mesh.";
                return false;
            }

            var assignedProfile = driver.Profile;
            if (assignedProfile == null)
            {
                error = "The driver does not have a Gaze Profile.";
                return false;
            }

            var profile = previewProfile != null ? previewProfile : assignedProfile;

            if (!profile.IsCompatibleWith(mesh))
            {
                error = "The Gaze Profile was calibrated for a different source mesh.";
                return false;
            }

            var channels = profile.Channels;
            if (channels == null || channels.Count == 0)
            {
                error = "The Gaze Profile does not contain any blendshape channels.";
                return false;
            }

            var indices = new int[channels.Count];
            var names = new string[channels.Count];
            var originals = new float[channels.Count];
            for (var channelIndex = 0; channelIndex < channels.Count; channelIndex++)
            {
                var channel = channels[channelIndex];
                if (channel == null || string.IsNullOrWhiteSpace(channel.BlendShapeName))
                {
                    error = $"Gaze channel {channelIndex + 1} does not have a blendshape name.";
                    return false;
                }

                var blendShapeIndex = profile.FindBlendShapeIndex(mesh, channelIndex);
                if (blendShapeIndex < 0 || blendShapeIndex >= mesh.blendShapeCount)
                {
                    error = $"Blendshape '{channel.BlendShapeName}' was not found on '{mesh.name}'.";
                    return false;
                }

                indices[channelIndex] = blendShapeIndex;
                names[channelIndex] = mesh.GetBlendShapeName(blendShapeIndex);
            }

            // The driver relinquishes its current gaze result only after the entire
            // profile-to-mesh mapping has passed validation. Capture the baseline
            // after that release so Stop restores the actual external facial pose.
            driver.BeginCalibrationPreviewOverride();
            try
            {
                for (var channelIndex = 0; channelIndex < channels.Count; channelIndex++)
                {
                    var blendShapeIndex = indices[channelIndex];
                    originals[channelIndex] = renderer.GetBlendShapeWeight(blendShapeIndex);
                }

                Driver = driver;
                Renderer = renderer;
                Profile = profile;
                driverProfile = assignedProfile;
                sourceMesh = mesh;
                blendShapeIndices = indices;
                resolvedBlendShapeNames = names;
                originalWeights = originals;
                previewWeights = new float[channels.Count];
                PreviewCoordinates = Vector2.zero;
                IsFreePreview = false;
                hasPreview = false;
                driverOverrideAcquired = true;
                IsActive = true;
                SubscribeEvents();
                return true;
            }
            catch (Exception exception)
            {
                driver.EndCalibrationPreviewOverride();
                error = $"Could not start calibration preview: {exception.Message}";
                return false;
            }
        }

        public void PreviewSample(GazeCalibrationPoint point)
        {
            if (!IsActive || !IsSessionValid())
                return;

            var sampleIndex = (int)point;
            var samples = Profile.Samples;
            if (sampleIndex < 0 || sampleIndex >= samples.Count)
                return;

            var sample = samples[sampleIndex];
            if (sample == null || !sample.Captured || sample.Weights.Count != previewWeights.Length)
                return;

            PreviewWeights(sample.Weights, point);
        }

        public void PreviewWeights(IReadOnlyList<float> weights, GazeCalibrationPoint point)
        {
            if (!IsActive || !IsSessionValid() || !IsValidPoint(point)
                || !TrySetPreviewWeights(weights, point == GazeCalibrationPoint.Center))
                return;

            PreviewCoordinates = CoordinatesFor(point);
            IsFreePreview = false;
            hasPreview = true;
            ApplyPreviewWeights();
        }

        public void PreviewNormalized(Vector2 normalized)
        {
            if (!IsActive || !IsSessionValid())
                return;

            normalized = new Vector2(
                Mathf.Clamp(normalized.x, -1f, 1f),
                Mathf.Clamp(normalized.y, -1f, 1f));

            var evaluated = new float[previewWeights.Length];
            var direction = new GazeDirectionResult(
                Vector3.forward,
                Vector3.forward,
                0f,
                0f,
                0f,
                0f,
                normalized.x,
                normalized.y);

            // TryEvaluate also guarantees that all nine calibration samples exist.
            // Correctives are intentionally included so this is an exact preview of
            // the completed profile rather than only its eye-look subset.
            if (!Profile.TryEvaluate(direction, true, evaluated))
                return;

            if (!TrySetPreviewWeights(evaluated, false))
                return;

            PreviewCoordinates = normalized;
            IsFreePreview = true;
            hasPreview = true;
            ApplyPreviewWeights();
        }

        public float[] GetOriginalWeightsCopy()
        {
            return (float[])originalWeights.Clone();
        }

        public void Stop()
        {
            var wasActive = IsActive;
            UnsubscribeEvents();
            var overriddenDriver = Driver;

            try
            {
                if (IsActive)
                    RestoreOriginalWeights();
            }
            finally
            {
                // Restoring while the override is still held guarantees that a
                // Timeline/editor evaluation cannot race the final face reset.
                if (driverOverrideAcquired && overriddenDriver != null)
                    overriddenDriver.EndCalibrationPreviewOverride();

                IsActive = false;
                Driver = null;
                Renderer = null;
                Profile = null;
                driverProfile = null;
                sourceMesh = null;
                blendShapeIndices = Array.Empty<int>();
                resolvedBlendShapeNames = Array.Empty<string>();
                originalWeights = Array.Empty<float>();
                previewWeights = Array.Empty<float>();
                PreviewCoordinates = Vector2.zero;
                IsFreePreview = false;
                hasPreview = false;
                driverOverrideAcquired = false;

                if (wasActive)
                    Stopped?.Invoke();
            }
        }

        public void Dispose()
        {
            Stop();
        }

        private void OnEditorUpdate()
        {
            if (!IsActive)
                return;

            if (EditorApplication.isPlayingOrWillChangePlaymode || !IsSessionValid())
            {
                Stop();
                return;
            }

            if (hasPreview)
                ApplyPreviewWeights();
        }

        private void OnBeforeAssemblyReload()
        {
            Stop();
        }

        private void OnEditorQuitting()
        {
            Stop();
        }

        private void OnPlayModeStateChanged(PlayModeStateChange state)
        {
            Stop();
        }

        private void OnSceneSaving(Scene scene, string path)
        {
            Stop();
        }

        private void OnPrefabSaving(GameObject root)
        {
            Stop();
        }

        private void SubscribeEvents()
        {
            if (eventsSubscribed)
                return;

            EditorApplication.update += OnEditorUpdate;
            EditorApplication.quitting += OnEditorQuitting;
            EditorApplication.playModeStateChanged += OnPlayModeStateChanged;
            AssemblyReloadEvents.beforeAssemblyReload += OnBeforeAssemblyReload;
            EditorSceneManager.sceneSaving += OnSceneSaving;
            PrefabStage.prefabSaving += OnPrefabSaving;
            eventsSubscribed = true;
        }

        private void UnsubscribeEvents()
        {
            if (!eventsSubscribed)
                return;

            EditorApplication.update -= OnEditorUpdate;
            EditorApplication.quitting -= OnEditorQuitting;
            EditorApplication.playModeStateChanged -= OnPlayModeStateChanged;
            AssemblyReloadEvents.beforeAssemblyReload -= OnBeforeAssemblyReload;
            EditorSceneManager.sceneSaving -= OnSceneSaving;
            PrefabStage.prefabSaving -= OnPrefabSaving;
            eventsSubscribed = false;
        }

        private bool IsSessionValid()
        {
            if (!IsActive || Driver == null || Renderer == null || Profile == null
                || driverProfile == null || sourceMesh == null)
                return false;

            // Cached indices belong exclusively to this exact mesh. Never use them
            // after a renderer mesh swap, even if the replacement has matching names.
            if (Renderer.sharedMesh != sourceMesh
                || Driver.TargetRenderer != Renderer
                || Driver.Profile != driverProfile
                || !Profile.IsCompatibleWith(sourceMesh))
                return false;

            var channels = Profile.Channels;
            if (channels == null || channels.Count != blendShapeIndices.Length
                                 || channels.Count != resolvedBlendShapeNames.Length)
                return false;

            for (var channelIndex = 0; channelIndex < channels.Count; channelIndex++)
            {
                var channel = channels[channelIndex];
                var blendShapeIndex = blendShapeIndices[channelIndex];
                if (channel == null || blendShapeIndex < 0 || blendShapeIndex >= sourceMesh.blendShapeCount
                    || !string.Equals(
                        sourceMesh.GetBlendShapeName(blendShapeIndex),
                        resolvedBlendShapeNames[channelIndex],
                        StringComparison.Ordinal)
                    || Profile.FindBlendShapeIndex(sourceMesh, channelIndex) != blendShapeIndices[channelIndex])
                    return false;
            }

            return previewWeights.Length == blendShapeIndices.Length
                   && originalWeights.Length == blendShapeIndices.Length;
        }

        private void ApplyPreviewWeights()
        {
            if (!hasPreview || !IsSessionValid())
                return;

            for (var channelIndex = 0; channelIndex < blendShapeIndices.Length; channelIndex++)
                Renderer.SetBlendShapeWeight(blendShapeIndices[channelIndex], previewWeights[channelIndex]);
        }

        private void RestoreOriginalWeights()
        {
            if (Renderer == null || sourceMesh == null || Renderer.sharedMesh != sourceMesh)
                return;

            var count = Mathf.Min(blendShapeIndices.Length, originalWeights.Length);
            if (count != resolvedBlendShapeNames.Length)
                return;

            // Validate the complete cached index/name mapping before writing any
            // value. A mesh rebuilt in place may reuse an index for another shape.
            for (var channelIndex = 0; channelIndex < count; channelIndex++)
            {
                var blendShapeIndex = blendShapeIndices[channelIndex];
                if (blendShapeIndex < 0 || blendShapeIndex >= sourceMesh.blendShapeCount
                    || !string.Equals(
                        sourceMesh.GetBlendShapeName(blendShapeIndex),
                        resolvedBlendShapeNames[channelIndex],
                        StringComparison.Ordinal))
                    return;
            }

            for (var channelIndex = 0; channelIndex < count; channelIndex++)
                Renderer.SetBlendShapeWeight(blendShapeIndices[channelIndex], originalWeights[channelIndex]);
        }

        private bool TrySetPreviewWeights(IReadOnlyList<float> weights, bool inputRepresentsCenter)
        {
            var channels = Profile != null ? Profile.Channels : null;
            if (weights == null || channels == null || weights.Count != previewWeights.Length
                || channels.Count != previewWeights.Length || originalWeights.Length != previewWeights.Length)
                return false;

            var translated = new float[weights.Count];
            for (var channelIndex = 0; channelIndex < weights.Count; channelIndex++)
            {
                var sampleWeight = weights[channelIndex];
                if (!IsFinite(sampleWeight))
                    return false;

                var channel = channels[channelIndex];
                if (channel == null)
                    return false;

                if (channel.Usage == GazeChannelUsage.Corrective
                    && channel.CorrectiveBlendMode == GazeCorrectiveBlendMode.AdditiveFromCenter)
                {
                    var centerWeight = inputRepresentsCenter
                        ? sampleWeight
                        : Profile.IsSampleCaptured(GazeCalibrationPoint.Center)
                            ? Profile.GetSampleWeight(GazeCalibrationPoint.Center, channelIndex)
                            : originalWeights[channelIndex];
                    if (!IsFinite(centerWeight))
                        return false;
                    translated[channelIndex] = originalWeights[channelIndex] + sampleWeight - centerWeight;
                }
                else
                {
                    translated[channelIndex] = sampleWeight;
                }

                if (!IsFinite(translated[channelIndex]))
                    return false;
            }

            Array.Copy(translated, previewWeights, translated.Length);
            return true;
        }

        private static Vector2 CoordinatesFor(GazeCalibrationPoint point)
        {
            var index = (int)point;
            return new Vector2(index % 3 - 1, index / 3 - 1);
        }

        private static bool IsValidPoint(GazeCalibrationPoint point)
        {
            var index = (int)point;
            return index >= 0 && index < 9;
        }

        private static bool IsFinite(float value)
        {
            return !float.IsNaN(value) && !float.IsInfinity(value);
        }
    }
}