using System;
using Fluctio.FluctioSim.Common.Configuration;
using Unity.MLAgents.Actuators;
using UnityEngine;

namespace Fluctio.FluctioSim.Core.Components.MachineLearning.Actuating
{
	[AddComponentMenu(Config.PrefixedName+"/Machine Learning/Actuating/Discrete Control", Config.ComponentMenuOrder + 540)]
	public class DiscreteControl : DirectControl
	{
		
		public override ActionSpec ActionSpec => ActionSpec.MakeDiscrete(3);
		
		public override ActionBuffers GetActions()
		{
			var modelActionValue = base.GetActions().DiscreteActions[0];
			var remappedValue = modelActionValue switch
			{
				0 => -1f,
				1 => 0f,
				2 => 1f,
				_ => throw new ArgumentOutOfRangeException($"Discrete action value {modelActionValue} is invalid (expected 0, 1 or 2)")
			};
			return ActionBuffersExtensions.FromOneContinuous(remappedValue);
		}

		public override void Heuristic(in ActionBuffers actionBuffers)
		{
			var axisValue = GetManualValue();
			var discreteActions = actionBuffers.DiscreteActions;
			discreteActions[0] = axisValue switch
			{
				<0 => 0,
				0 => 1,
				>0 => 2,
				float.NaN => throw new ArgumentException("Manual axis value is NaN"),
			};
		}
		
	}
}