using Fluctio.FluctioSim.Common.Configuration;
using UnityEngine;

namespace Fluctio.FluctioSim.Core.Components.MachineLearning.NumberProviders
{
    [AddComponentMenu(Config.PrefixedName+"/Machine Learning/Number Providers/Number Change Calculator", Config.ComponentMenuOrder + 100)]
    public class NumberChangeProvider : AgentNumberProvider
    {
        [field: SerializeField] public NumberProvider NumberProvider { get; set; }
        private float? _previousNumber = null;
        private float? _currentNumber = null;

        private void OnEnable()
        {
            Agent.UnityAgent.EpisodeInitializing += OnEpisodeInitializing;
        }

        private void OnDisable()
        {
            Agent.UnityAgent.EpisodeInitializing -= OnEpisodeInitializing;
        }

        private void OnEpisodeInitializing()
        {
            _previousNumber = null;
            _currentNumber = null;
        }

        private void FixedUpdate()
        {
            _previousNumber = _currentNumber;
        }

        public override float GetNumber()
        {
            // Workaround for this bug: https://github.com/Unity-Technologies/ml-agents/issues/6190
            if (Agent.UnityAgent.StepCount == 0)
            {
                return 0;
            }
            
            // _currentNumber should be calculated here instead of FixedUpdate() to preserve execution order
            // For example, if this class is used between Randomization and RewardAdder,
            // then _currentNumber is also calculated at that time, not in FixedUpdate()
            _currentNumber = NumberProvider.GetNumber();
            
            var difference = _currentNumber - _previousNumber;
            return difference ?? 0;
        }
    }
}