using System.Collections.Generic;
using Fluctio.FluctioSim.Common.Configuration;
using Fluctio.FluctioSim.Core.Components.Prefabs.Primitives.Base;
using Mujoco;
using Unity.MLAgents.Actuators;
using UnityEngine;

namespace Fluctio.FluctioSim.Core.Components.MachineLearning.Actuators
{
	[AddComponentMenu(Config.PrefixedName+"/Machine Learning/Actuator", Config.ComponentMenuOrder + 30)]
	public class MlMjActuator : Actuator
	{
		#region Inspector settings
		[SerializeField] public PrimitiveJoint joint;
		[SerializeField] public bool isHeuristicsEnabled;
		[SerializeField] public float power = 100;
		#endregion
		
		#region Components setup
		[SerializeField, HideInInspector] protected GameObject actuatorObject;
		[SerializeField, HideInInspector] protected ProxyMlActuatorComponent mlagentsActuator;
		[SerializeField, HideInInspector] protected MjActuator mujocoActuator;

		protected override void InitializeOnce()
		{
			base.InitializeOnce();
			
			actuatorObject = new GameObject("Actuator");
			actuatorObject.transform.SetParent(transform);
			mlagentsActuator = actuatorObject.AddComponent<ProxyMlActuatorComponent>();
			mujocoActuator = actuatorObject.AddComponent<MjActuator>();

			mlagentsActuator.originalActuator = this;
			mujocoActuator.CommonParams.CtrlLimited = true;
			mujocoActuator.CommonParams.CtrlRange = new Vector2(-1, 1);
			
			var joints = GetComponents<PrimitiveJoint>();
			if (joints.Length == 1)
			{
				joint = joints[0];
			} else {
				Debug.LogWarning("Please, set the actuator's joint manually");	
			}
		}

		protected override IEnumerable<GameObject> GetInternalObjects() => new[] {actuatorObject};
		protected override IEnumerable<Object> GetCreatedInternals() => new[] {actuatorObject};

		public override void OnAnyChanged()
		{
			base.OnAnyChanged();
			mujocoActuator.Joint = joint?.JointComponent;
		}

		public override void OnSelfChanged()
		{
			base.OnSelfChanged();
			power = Mathf.Max(power, Config.MinFloat);
			mujocoActuator.CustomParams.GainPrm[0] = power;
		}

		public override void OnOtherChanged()
		{
			base.OnOtherChanged();
			power = mujocoActuator.CustomParams.GainPrm[0];
		}
		#endregion

		#region Abstract methods implementation
		public override bool IsHeuristicsEnabled => isHeuristicsEnabled;

		public override ActionSpec ActionSpec => ActionSpec.MakeContinuous(1);
		
		public override void Heuristic(ActionSegment<float> continiousActions, ActionSegment<int> discreteActions)
		{
			//TODO: use InputSystem and make action configurable
			continiousActions[0] = Input.GetAxis("Horizontal");
		}
		
		public override void PerformAction(ActionSegment<float> continiousActions, ActionSegment<int> discreteActions)
		{
			mujocoActuator.Control = continiousActions[0];
		}
		#endregion
	}
}