using AGXUnity;
using AGXUnity.IO.OpenPLX;
using AGXUnity.Utils;
using System;
using System.Collections.Generic;
using System.Linq;
using UnityEngine;

/// <summary>
/// This script is a ported version of the controller provided in the t06-urdf-robot OpenPLX tutorial
/// </summary>
public class PandaController : MonoBehaviour
{
  /// <summary>
  /// A wrapper around a simple double array which allows for operator overloading to mimic numpy arrays
  /// </summary>
  struct VectorWrapper
  {
    double[] values;

    public double this[ int index ]
    {
      get => values[ index ];
      set => values[ index ] = value;
    }

    public static implicit operator double[]( VectorWrapper v ) => v.values;

    public static VectorWrapper operator +( VectorWrapper a, VectorWrapper b )
    {
      var ret = new double[ a.values.Length ];
      for ( int i = 0; i < a.values.Length; i++ )
        ret[ i ] = a.values[ i ] + b.values[ i ];
      return new VectorWrapper( ret );
    }

    public static VectorWrapper operator -( VectorWrapper a, VectorWrapper b )
    {
      var ret = new double[ a.values.Length ];
      for ( int i = 0; i < a.values.Length; i++ )
        ret[ i ] = a.values[ i ] - b.values[ i ];
      return new VectorWrapper( ret );
    }

    public static VectorWrapper operator *( VectorWrapper a, double s )
    {
      var ret = new double[ a.values.Length ];
      for ( int i = 0; i < a.values.Length; i++ )
        ret[ i ] = a.values[ i ] * s;
      return new VectorWrapper( ret );
    }

    public static VectorWrapper operator /( VectorWrapper a, double s )
    {
      var ret = new double[ a.values.Length ];
      for ( int i = 0; i < a.values.Length; i++ )
        ret[ i ] = a.values[ i ] / s;
      return new VectorWrapper( ret );
    }

    public VectorWrapper( IEnumerable<double> val )
    {
      values = val.ToArray();
    }
    public VectorWrapper( int size )
    {
      values = new double[ size ];
    }
  }

  VectorWrapper m_InitialPose = new VectorWrapper(7);
  OpenPLXRoot panda;

  OutputSource[] m_AngleOutputs = new OutputSource[7];
  OutputSource[] m_AngularVelocityOutputs = new OutputSource[7];
  InputTarget[] m_TorqueInputs = new InputTarget[7];

  double kp,kd,ki;
  VectorWrapper e_int;

  Trajectory m_trajectory;
  agxModel.SerialKinematicChain chain;

  /// <summary>
  /// SerialKinematicChain currently requires that ranges are specified on the mate which is not the case for OpenPLX imported models.
  /// This method moves the range specification from the mapped range onto the mapped mate.
  /// </summary>
  private void MigrateRanges()
  {
    for ( int i = 1; i < 8; i++ ) {
      var src = panda.FindMappedObject( $"panda.panda_joint{i}_range" ).GetComponent<Constraint>().GetController<RangeController>();
      var target = panda.FindMappedObject( $"panda.panda_joint{i}" ).GetComponent<Constraint>().GetController<RangeController>();
      target.Enable = true;
      src.Enable = false;
      target.Range = src.Range;
    }
  }

  void Start()
  {
    panda = GetComponent<OpenPLXRoot>().GetInitialized<OpenPLXRoot>();
    panda.gameObject.InitializeAll();

    var signals = panda.GetComponent<OpenPLXSignals>().GetInitialized<OpenPLXSignals>();

    var rootName = panda.Native.getName();

    // Find initial angles and setup input/output signals
    for ( int i = 0; i < 7; i++ ) {
      m_InitialPose[ i ]            = panda.Native.getObject( "panda" ).getObject( $"panda_joint{i+1}" ).getNumber( "initial_angle" );
      m_AngleOutputs[ i ]           = signals.FindOutputSource( $"{rootName}.panda.panda_joint{i+1}.angle_output" );
      m_AngularVelocityOutputs[ i ] = signals.FindOutputSource( $"{rootName}.panda.panda_joint{i+1}.angular_velocity_output" );
      m_TorqueInputs[ i ]           = signals.FindInputTarget( $"{rootName}.panda.panda_joint{i+1}_actuator.torque_input" );
    }

    // Specify the start and end bodies in the chain
    var baseBody = panda.FindMappedObject("panda.panda_link0").GetComponent<RigidBody>().GetInitialized();
    var eeBody = panda.FindMappedObject("panda.panda_link7").GetComponent<RigidBody>().GetInitialized();

    // Calculate the tool offset in local coordinates.
    var agx_local_offset = panda.FindMappedObject( "panda.panda_link7.panda_joint8_mc" ).transform.localPosition.ToVec3();
    var local_transform = agx.AffineMatrix4x4.translate( agx_local_offset );

    // Migrate Range specification onto mate secondary constraints and setup the kinematic chain
    MigrateRanges();
    chain = new agxModel.SerialKinematicChain( Simulation.Instance.Native, baseBody.Native, eeBody.Native, local_transform );

    // Setup PID parameters for the position control
    var numJoints = chain.getNumJoints();
    var step = AGXUnity.Simulation.Instance.TimeStep;
    kp = numJoints * 3.0 / step;
    kd = 1.0 / ( 10.0 * step );
    ki = 0;

    // Setup the target tool position in the base's coordinate system
    var x_ready = eeBody.Native.getFrame().getLocalMatrix() * baseBody.Native.getFrame().getMatrix().inverse();
    var x_ready_offset = x_ready.getTranslate() + agx_local_offset;
    var x_d = x_ready_offset + new agx.Vec3( 0.5, 0.4, -0.4 );
    var rot_d = new agx.Quat(Mathf.PI / 2, agx.Vec3.Y_AXIS() );
    var x_t = new agx.AffineMatrix4x4( rot_d, x_d )  ;

    // Compute the pose for the target tool position and calculate a trajectory
    agx.RealVector output = new agx.RealVector(7);
    agx.RealVector input = new agx.RealVector((double[])m_InitialPose);
    var res = chain.computeInverseKinematics( x_t, input, output );
    var desiredPose = new VectorWrapper(output.ToArray());
    m_trajectory = new Trajectory( m_InitialPose, desiredPose, 7 );

    e_int = new VectorWrapper( 7 );

    Simulation.Instance.StepCallbacks.PreStepForward += PreStep;
  }

  void PreStep()
  {
    // Calculate the desired position given the trajectory and the time
    var h = Simulation.Instance.TimeStep;
    var (q_d, qd_d, qdd_d) = m_trajectory.calc( (float)Simulation.Instance.Native.getTimeStamp() );

    // Get the current position and velocities.
    var q  = new VectorWrapper(m_AngleOutputs.Select(o => o.HasSentSignal ? o.GetValue<double>() : 0.0f));
    var qd = new VectorWrapper(m_AngularVelocityOutputs.Select(o => o.HasSentSignal ? o.GetValue<double>() : 0.0f));

    // Calculate PID errors
    var e = q_d - q;
    var e_dot = qd_d - qd;
    e_int = e_int + e * h;

    // Calculate feedback terms
    var qdd_ff_fb = qdd_d + e * kp;
    qdd_ff_fb     = qdd_ff_fb + e_dot * kd;
    qdd_ff_fb     = qdd_ff_fb + e_int * ki;

    // Calculate the torques required to correct the errors
    var torques = new agx.RealVector();
    var targets = new agx.RealVector((double[])qdd_ff_fb);
    var res = chain.computeInverseDynamics( targets, torques );

    // Send desired torques to motors via OpenPLX signals
    for ( int i = 0; i < torques.Count; i++ )
      m_TorqueInputs[ i ].SendSignal( torques[ i ] );
  }

  /// <summary>
  /// Utility class compute a trajectory given a start and end pose along with a duration.
  /// </summary>
  class Trajectory
  {
    float m_duration;
    VectorWrapper m_start;
    VectorWrapper m_final;

    public Trajectory( VectorWrapper start, VectorWrapper final, float duration = 5.0f )
    {
      m_duration = duration;

      m_start = start;
      m_final = final;
    }

    VectorWrapper q_d( float time )
    {
      var t = time/m_duration;
      var scaled = ( 1.0 - Mathf.Cos( t * Mathf.PI) ) * 0.5;
      return m_start + ( m_final - m_start ) * scaled;
    }

    VectorWrapper qd_d( float t )
    {
      var h = AGXUnity.Simulation.Instance.TimeStep;
      return ( q_d( t + h ) - q_d( t - h ) ) / 2 * h;
    }

    VectorWrapper qdd_d( float t )
    {
      var h = AGXUnity.Simulation.Instance.TimeStep;
      return ( ( q_d( t + h ) - q_d( t ) * 2 ) + q_d( t - h ) ) / ( h * h );
    }

    public Tuple<VectorWrapper, VectorWrapper, VectorWrapper> calc( float time )
    {
      return Tuple.Create( q_d( time ), qd_d( time ), qdd_d( time ) );
    }
  }
}
