using AGXUnity;
using AGXUnity.Utils;
using agxUtil;
using UnityEngine;

public class MLMachineResetter : ScriptComponent
{
  private agx.RigidBody m_mainBody;
  private agx.AffineMatrix4x4 m_mainTransform;

  private agxSDK.Collection m_collection;
  private agxModel.ConstraintPositionVector m_constraintPositions;

  private agxSDK.Assembly m_jumpAssembly;

  public int NumBodies => m_collection?.getRigidBodies().Count ?? 0;
  public int NumConstraints => m_collection?.getConstraints().Count ?? 0;
  public Vector3 MainBodyPos => m_mainTransform.getTranslate().ToHandedVector3();
  public Vector3 MainBodyRot => m_mainTransform.getRotate().ToHandedQuaternion().eulerAngles;

  public void SaveCurrentPose()
  {
    foreach ( var sc in GetComponentsInChildren<ScriptComponent>() )
      if ( sc != this )
        sc.GetInitialized();

    m_collection = new agxSDK.Collection();
    m_constraintPositions = new agxModel.ConstraintPositionVector();
    m_jumpAssembly = new agxSDK.Assembly();

    var bodies = GetComponentsInChildren<RigidBody>();
    var constraints = GetComponentsInChildren<Constraint>();
    var tires = GetComponentsInChildren<AGXUnity.Model.TwoBodyTire>();

    // Add initial native components properties to native collections
    foreach ( var bod in bodies ) {
      m_collection.add( bod.GetInitialized().Native );
      m_jumpAssembly.add( bod.Native );
    }
    foreach ( var constraint in constraints ) {
      if ( constraint.isActiveAndEnabled ) {
        var init = constraint.GetInitialized();
        m_collection.add( init.Native );
        m_constraintPositions.Add( new agxModel.ConstraintPosition( constraint.Native, constraint.GetCurrentAngle() ) );
      }
    }
    foreach ( var tire in tires )
      m_collection.add( tire.GetInitialized().Native );

    // save the initial transform of the main body
    m_mainBody = bodies[ 0 ].Native;
    m_mainTransform = m_mainBody.getTransform();
  }

  public void LoadSavedPose()
  {
    // Apply transform for main body
    agxUtilSWIG.jumpRequest( m_jumpAssembly, m_mainBody, m_mainTransform );

    // Compute and apply new configuration
    var reconfRequest = new agxModel.ReconfigureRequest();
    agxModel.BodyTransformVector result = new agxModel.BodyTransformVector();
    reconfRequest.computeTransforms( m_collection, m_mainBody, m_constraintPositions, result );
    reconfRequest.applyTransforms( m_collection, m_constraintPositions, result );

    // Applying new transforms to a machine might cause enabled LockControllers to be violated.
    // Here, we set the locked position for all enabled locks after applying the transformation 
    // avoid any potential violations.
    foreach ( var constraint in GetComponentsInChildren<Constraint>() ) {
      var transLock = constraint.GetController<LockController>(Constraint.ControllerType.Translational);
      if ( transLock != null )
        if ( transLock.Enable )
          transLock.Position = constraint.GetCurrentAngle( Constraint.ControllerType.Rotational );

      var rotLock = constraint.GetController<LockController>(Constraint.ControllerType.Rotational);
      if ( rotLock != null )
        if ( rotLock.Enable )
          rotLock.Position = constraint.GetCurrentAngle( Constraint.ControllerType.Rotational );
    }
  }
}
