using System;
using Unity.MLAgents.Sensors;
using UnityEngine;

namespace Fluctio.FluctioSim.Core.Components.MachineLearning.Sensors
{
	[AddComponentMenu("/")] // hide component from menu
	public class ProxyMlSensorComponent : SensorComponent
	{
		[SerializeField] public ObservationType observationType = ObservationType.Default;
		[SerializeField] public Sensor originalSensor;
		
		public override ISensor[] CreateSensors()
		{
			if (originalSensor == null)
			{
				Debug.LogError($"{nameof(originalSensor)} is not set");
				return Array.Empty<ISensor>();
			}
			var proxySensor = new ProxyMlSensor(originalSensor, observationType);
			return new ISensor[]
			{
				originalSensor.StackedSize > 1
					? new StackingSensor(proxySensor, originalSensor.StackedSize)
					: proxySensor
			};
		}
	}


	public class ProxyMlSensor : VectorSensor, ISensor
	{
		private readonly Sensor _sensor;
		
		public ProxyMlSensor(Sensor sensor, ObservationType observationType = ObservationType.Default) : base(sensor.ObservationSize, sensor.SensorName, observationType)
		{
			_sensor = sensor;
		}

		public new void Update()
		{
			base.Update();
			_sensor.CollectObservations(this);
		}
	}
}