using Mujoco;
using Unity.MLAgents.Sensors;
using UnityEngine;
using static Mujoco.MjJointScalarSensor;

namespace Fluctio.FluctioSim.Core.Components.MachineLearning.Sensors.Hinge
{
	public abstract class Encoder : JointSensor<MjHingeJoint>
	{
		public override AvailableSensors SensorType => AvailableSensors.JointPos;
		
		protected override void SetupInternalJoint()
		{
			InternalJoint!.Configuration = InitialConfiguration;
		}

		public abstract float InitialConfiguration { get; }
		
		public override int ObservationSize => 2;

		public override void CollectObservations(VectorSensor sensor)
		{
			var reading = (float)mujocoSensor.SensorReading;
			// splitting to sin and cos to avoid jumping from 2pi to 0
			// idea taken from: https://github.com/Balint-H/modular-agents/blob/927a4ead43c8b8e4db902657f37550db3b7d45e9/UnityMjExamples/Assets/Environments/0.CartPole/RL%20Scripts/HingeObservations.cs#L44
			sensor.AddObservation(Mathf.Sin(reading));
			sensor.AddObservation(Mathf.Cos(reading));
		}
	}
}