using System;
using System.ComponentModel;
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>
	{
		[field: SerializeField] public AngleReadingType ReadingType { get; private set; } = AngleReadingType.SinCos;
		public override AvailableSensors SensorType => AvailableSensors.JointPos;
		
		private float _barRotation;
		private float _prevBarRotation;
		private float _oldBarSpeed;
		
		protected override void SetupInternalJoint()
		{
			InternalJoint!.Configuration = InitialConfiguration;
		}

		public abstract float InitialConfiguration { get; }
		
		public override int ObservationSize => ReadingType == AngleReadingType.SinCos ? 2 : 1;

		public override void CollectObservations(VectorSensor sensor)
		{
			var reading = (float)mujocoSensor.SensorReading;
			switch (ReadingType)
			{
				case AngleReadingType.SinCos:
					// 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((float)Math.Sin(InternalJoint!.RawConfiguration));
					sensor.AddObservation((float)Math.Cos(InternalJoint!.RawConfiguration));
					break;
				case AngleReadingType.Normalized:
					sensor.AddObservation(InternalJoint!.Configuration / 360);
					break;
				case AngleReadingType.Radians:
					sensor.AddObservation(InternalJoint!.Configuration * Mathf.Deg2Rad);
					break;
				case AngleReadingType.AngleChange:
					var delta = Time.fixedDeltaTime;
					_barRotation = InternalJoint!.Configuration / 360.0f;
					var barSpeed = (_barRotation - _prevBarRotation) / delta;

					if (Mathf.Abs(barSpeed) > 20.0f)
					{
						barSpeed = _oldBarSpeed;
					}

					sensor.AddObservation(barSpeed);
					//print("delta: " + delta + "     speed: " + barSpeed);
					_prevBarRotation = _barRotation;
					_oldBarSpeed = barSpeed;
					break;
				default:
					throw new ArgumentOutOfRangeException();
			}
		}
	}

	public enum AngleReadingType
	{
		[Description("Two values: sin(angle) and cos(angle)")]
		SinCos,
		[Description("Normalized: [0; 1]")]
		Normalized,
		[Description("Radians: [0; 2pi]")]
		Radians,
		[Description("Speed calculated from angle change (normalized)")]
		AngleChange,
	}
}