using System.Linq;
using Fluctio.FluctioSim.Core.Components.Prefabs.Primitives.Base;
using JetBrains.Annotations;
using Mujoco;
using Unity.MLAgents.Sensors;
using UnityEngine;
using static Mujoco.MjJointScalarSensor;

namespace Fluctio.FluctioSim.Core.Components.MachineLearning.Sensors
{
	public abstract class JointSensor<TMjJoint> : MlMjSensor<MjJointScalarSensor>
		where TMjJoint : MjBaseJoint
	{
		[SerializeField] public PrimitiveJoint joint;
		[CanBeNull] protected TMjJoint InternalJoint => joint?.JointComponent as TMjJoint;

		private static string JointTypeName => typeof(TMjJoint).Name.Replace("Mj", "").Replace("Joint", "");
		
		protected override void InitializeOnce()
		{
			base.InitializeOnce();
			var joints = GetComponents<PrimitiveJoint>().Where(j => j.JointComponent is TMjJoint).Take(2).ToArray();
			if (joints.Length == 1)
			{
				joint = joints[0];
			} else {
				Debug.LogWarning($"Please, set the encoder's joint to a valid {JointTypeName} joint manually");	
			}
			mujocoSensor.SensorType = SensorType;
		}

		public abstract AvailableSensors SensorType { get; }

		public override void OnAnyChanged()
		{
			base.OnAnyChanged();
			if (joint != null && joint.JointComponent is not TMjJoint)
			{
				Debug.LogWarning($"Sensor's joint was not set - only {JointTypeName} joints are supported");
				joint = null;
			}
			mujocoSensor.Joint = joint?.JointComponent;
			if (InternalJoint != null)
			{
				SetupInternalJoint();
			}
		}

		protected virtual void SetupInternalJoint() {}
		
		public override int ObservationSize => 1;
		
		public override void CollectObservations(VectorSensor sensor)
		{
			sensor.AddObservation((float)mujocoSensor.SensorReading);
		}
	}
}