using Fluctio.FluctioSim.Core.Components.MachineLearning.NumberProviders;
using Fluctio.FluctioSim.Core.Components.Prefabs.Primitives.Base;
using Mujoco;
using UnityEngine;

namespace Fluctio.FluctioSim.Core.Components.Prefabs.Complex.CartPole
{
    public class CartPoleReward : NumberProvider
    {
        [SerializeField] protected PrimitiveJoint sliderJoint;
        [SerializeField] protected PrimitiveJoint spinnerJoint;

        [SerializeField] protected float upwardsLookingFactor = 0.1f;
        [SerializeField] protected float smallRotationSpeedFactor = 0.001f;
        [SerializeField] protected float nearCenterFactor = 0.001f;

        private MjSlideJoint _internalSliderJoint;
        private MjHingeJoint _internalSpinnerJoint;

        public void Start()
        {
            _internalSliderJoint = sliderJoint.GetJointComponent<MjSlideJoint>();
            _internalSpinnerJoint = spinnerJoint.GetJointComponent<MjHingeJoint>();
        }

        public override float GetNumber()
        {
            var rotation = Mathf.Abs(_internalSpinnerJoint.Configuration / 360) - 0.5f;
            var rotationSpeed = Mathf.Abs(_internalSpinnerJoint.Velocity);
            var centerDistance = Mathf.Abs(_internalSliderJoint.Configuration);
            return upwardsLookingFactor * rotation * rotation - smallRotationSpeedFactor * rotationSpeed - nearCenterFactor * centerDistance;
        }
    }
}