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

namespace Fluctio.FluctioSim.Core.Components.MachineLearning.NumberProviders
{
	[AddComponentMenu(Config.PrefixedName+"/Machine Learning/Number Providers/Agent's Reward Number", Config.ComponentMenuOrder + 90)]
	[DefaultExecutionOrder(40)] // between AddReward/SetReward and EndEpisode to capture all reward changes
	public class RewardNumberProvider : AgentNumberProvider
	{
		private float _previousEpisodeReward = 0;
		private float _currentEpisodeReward = 0;
		
		private void OnEnable()
		{
			CheckAgentExists();
			Agent.UnityAgent.EpisodeInitializing += OnEpisodeStart;
		}

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

		private void OnEpisodeStart()
		{
			_previousEpisodeReward = _currentEpisodeReward;
			_currentEpisodeReward = 0;
		}

		private void FixedUpdate()
		{
			_currentEpisodeReward = Agent.UnityAgent.GetCumulativeReward();
		}

		public override float GetNumber()
		{
			CheckAgentExists();
			return _previousEpisodeReward;
		}
	}
}