#include "rbdlsim.h"

#include <ccd/ccd.h>
#include <ccd/quat.h>
#include <rbdl/Dynamics.h>
#include <rbdl/Kinematics.h>
#include <rbdl/Model.h>
#include <rbdl/rbdl_math.h>

#include <iostream>
#include <vector>

using namespace std;

namespace RBDLSim {

void SimBody::step(double dt) {
  ForwardDynamics(mModel, q, qdot, tau, qddot);

  // semi-implicit eulers
  qdot += dt * qddot;

  for (int i = 1; i < mModel.mJoints.size(); ++i) {
    const Joint& joint = mModel.mJoints[i];
    if (joint.mJointType != JointTypeSpherical) {
      q.block(joint.q_index, 0, joint.mDoFCount, 1) +=
          dt * qdot.block(joint.q_index, 0, joint.mDoFCount, 1);
      continue;
    }

    // Integrate the Quaternion
    Quaternion q0 = mModel.GetQuaternion(i, q);
    Vector3d omega(qdot.block(joint.q_index, 0, 3, 1));
    Quaternion qd = q0.omegaToQDot(omega);
    Quaternion q1 = (q0 + qd * dt).normalize();
    mModel.SetQuaternion(i, q1, q);
  }
}

/** Support function for box */
void SimShapeSupport(
    const void* _shape,
    const ccd_vec3_t* _dir,
    ccd_vec3_t* v) {
  SimShape* shape = (SimShape*)_shape;

  CCD_VEC3(pos, shape->pos[0], shape->pos[1], shape->pos[2]);
  CCD_QUAT(
      quat,
      shape->orientation[0],
      shape->orientation[1],
      shape->orientation[2],
      shape->orientation[3]);
  ccd_vec3_t dir;
  ccd_quat_t qinv;

  // apply rotation on direction vector
  ccdVec3Copy(&dir, _dir);
  ccdQuatInvert2(&qinv, &quat);
  ccdQuatRotVec(&dir, &qinv);

  // compute support point in specified direction
  if (shape->mType == SimShape::Box) {
    ccdVec3Set(
        v,
        ccdSign(ccdVec3X(&dir)) * shape->scale[0] * CCD_REAL(0.5),
        ccdSign(ccdVec3Y(&dir)) * shape->scale[1] * CCD_REAL(0.5),
        ccdSign(ccdVec3Z(&dir)) * shape->scale[2] * CCD_REAL(0.5));

  } else if (shape->mType == SimShape::Sphere) {
    ccd_real_t len;
    assert((shape->scale[0] - shape->scale[1]) < 1.0e-12);
    assert((shape->scale[2] - shape->scale[1]) < 1.0e-12);

    len = ccdVec3Len2(&dir);
    if (len - CCD_EPS > CCD_ZERO) {
      ccdVec3Copy(v, &dir);
      ccdVec3Scale(v, shape->scale[0] / CCD_SQRT(len));
    } else {
      ccdVec3Set(v, CCD_ZERO, CCD_ZERO, CCD_ZERO);
    }
  }

  // transform support point according to position and rotation of object
  ccdQuatRotVec(v, &quat);
  ccdVec3Add(v, &pos);
}

bool CheckPenetration(
    const SimShape& shape_a,
    const SimShape& shape_b,
    CollisionInfo& cinfo) {
  ccd_t ccd;
  CCD_INIT(&ccd);
  ccd.support1 = SimShapeSupport;
  ccd.support2 = SimShapeSupport;
  ccd.max_iterations = 100;
  ccd.epa_tolerance = 0.0001;

  ccd_real_t depth;
  ccd_vec3_t dir, pos;
  int intersect =
      ccdGJKPenetration(&shape_a, &shape_b, &ccd, &depth, &dir, &pos);

  if (intersect == 0) {
    cinfo.pos.set(pos.v[0], pos.v[1], pos.v[2]);
    cinfo.dir.set(dir.v[0], dir.v[1], dir.v[2]);
  }

  return !intersect;
}

void SimBody::updateCollisionShapes() {
  UpdateKinematicsCustom(mModel, &q, nullptr, nullptr);

  for (SimBody::BodyCollisionInfo& body_col_info : mCollisionShapes) {
    const unsigned int& body_id = body_col_info.first;
    SimShape& shape = body_col_info.second;

    shape.pos =
        CalcBodyToBaseCoordinates(mModel, q, body_id, Vector3d::Zero(), false);
    shape.orientation.fromMatrix(
        CalcBodyWorldOrientation(mModel, q, body_id, false));
  }
}

void World::updateCollisionShapes() {
  for (SimBody& body : mBodies) {
    body.updateCollisionShapes();
  }
}

void World::detectCollisions() {
  mContactPoints.clear();

  for (SimBody& body : mBodies) {
    for (SimBody::BodyCollisionInfo& body_col_info : body.mCollisionShapes) {
      SimShape& body_shape = body_col_info.second;

      // check against all static shapes
      for (SimShape& static_shape : mStaticShapes) {
        bool has_penetration = false;
        CollisionInfo cinfo;

        has_penetration = CheckPenetration(static_shape, body_shape, cinfo);

        if (has_penetration) {
          cinfo.mBodyA = nullptr;
          cinfo.mBodyB = &body;
          mContactPoints.push_back(cinfo);
        }
      }
    }
  }
}

bool World::step(double dt) {
  for (SimBody& body : mBodies) {
    body.step(dt);
  }

  mSimTime += dt;
  return true;
}

SimBody CreateSphereBody(
    double mass,
    double radius,
    const Vector3d& pos,
    const Vector3d& vel) {
  SimBody result;

  Joint joint_trans_xyz(JointTypeTranslationXYZ);
  Joint joint_rot_quat(JointTypeSpherical);
  Body nullbody(0., Vector3d(0., 0., 0.), Matrix3d::Zero());
  Body body(
      mass,
      Vector3d(0., 0., 0.),
      Matrix3d::Identity() * 2. / 5. * mass * radius * radius);

  result.mModel.AppendBody(
      Xtrans(Vector3d(0., 0., 0.)),
      joint_trans_xyz,
      nullbody);
  unsigned int sphere_body = result.mModel.AppendBody(
      Xtrans(Vector3d(0., 0., 0.)),
      joint_rot_quat,
      body);

  result.q = VectorNd::Zero(result.mModel.q_size);
  result.q.block(0, 0, 3, 1) = pos;
  result.mModel.SetQuaternion(
      sphere_body,
      Quaternion(0., 0., 0., 1.),
      result.q);
  result.qdot = VectorNd::Zero(result.mModel.qdot_size);
  result.qddot = VectorNd::Zero(result.mModel.qdot_size);
  result.tau = VectorNd::Zero(result.mModel.qdot_size);

  SimShape shape;
  shape.mType = SimShape::Sphere;
  shape.pos = pos;
  shape.orientation.set(0., 0., 0., 1.);
  shape.scale.set(radius, radius, radius);

  result.mCollisionShapes.push_back(
      SimBody::BodyCollisionInfo(sphere_body, shape));

  return result;
}

SimBody CreateBoxBody(
    double mass,
    const Vector3d& size,
    const Vector3d& pos,
    const Vector3d& vel) {
  SimBody result;

  Joint joint_trans_xyz(JointTypeTranslationXYZ);
  Joint joint_rot_quat(JointTypeSpherical);
  Body nullbody(0., Vector3d(0., 0., 0.), Matrix3d::Zero());
  Matrix3d inertia(Matrix3d::Zero());
  inertia(0, 0) = (1. / 12.) * mass * (size[1] * size[1] + size[2] * size[2]);
  inertia(1, 1) = (1. / 12.) * mass * (size[0] * size[0] + size[2] * size[2]);
  inertia(2, 2) = (1. / 12.) * mass * (size[1] * size[1] + size[0] * size[0]);
  Body body(mass, Vector3d(0., 0., 0.), inertia);

  result.mModel.AppendBody(
      Xtrans(Vector3d(0., 0., 0.)),
      joint_trans_xyz,
      nullbody);
  unsigned int sphere_body = result.mModel.AppendBody(
      Xtrans(Vector3d(0., 0., 0.)),
      joint_rot_quat,
      body);

  result.q = VectorNd::Zero(result.mModel.q_size);
  result.q.block(0, 0, 3, 1) = pos;
  result.mModel.SetQuaternion(
      sphere_body,
      Quaternion(0., 0., 0., 1.),
      result.q);
  result.qdot = VectorNd::Zero(result.mModel.qdot_size);
  result.qddot = VectorNd::Zero(result.mModel.qdot_size);
  result.tau = VectorNd::Zero(result.mModel.qdot_size);

  SimShape shape;
  shape.mType = SimShape::Sphere;
  shape.pos = pos;
  shape.orientation.set(0., 0., 0., 1.);
  shape.scale.set(size[0], size[1], size[2]);

  result.mCollisionShapes.push_back(
      SimBody::BodyCollisionInfo(sphere_body, shape));

  return result;
}

}  // namespace RBDLSim