#include "rbdlsim.h"
#include "utils.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) {
  // Prerequisite: qdot has already been updated by resolveCollisions();
  calcNextPositions(dt, qdot, q);
}

void SimBody::calcNextPositions(
    double dt,
    const VectorNd& in_qdot,
    VectorNd& io_q) {
  for (int i = 1; i < mModel.mJoints.size(); ++i) {
    const Joint& joint = mModel.mJoints[i];
    if (joint.mJointType != JointTypeSpherical) {
      io_q.block(joint.q_index, 0, joint.mDoFCount, 1) +=
          dt * in_qdot.block(joint.q_index, 0, joint.mDoFCount, 1);
      continue;
    }

    // Integrate the Quaternion
    Quaternion q0 = mModel.GetQuaternion(i, io_q);
    Vector3d omega(in_qdot.block(joint.q_index, 0, 3, 1));
    Quaternion qd = q0.omegaToQDot(omega);
    Quaternion q1 = (q0 + qd * dt).normalize();
    mModel.SetQuaternion(i, q1, io_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);
    }
  } else {
    cerr << "Unknown shape type: " << shape->mType << endl;
    abort();
  }

  // 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) {
  cinfo.mShapeA = &shape_a;
  cinfo.mShapeB = &shape_b;
  if (shape_a.mType == SimShape::Sphere && shape_b.mType == SimShape::Plane) {
    return CheckPenetrationSphereVsPlane(shape_a, shape_b, cinfo);
  }
  if (shape_b.mType == SimShape::Sphere && shape_a.mType == SimShape::Plane) {
    bool result = CheckPenetrationSphereVsPlane(shape_b, shape_a, cinfo);
    cinfo.dir *= -1.;
    return result;
  }
  if (shape_a.mType == SimShape::Sphere && shape_b.mType == SimShape::Sphere) {
    return CheckPenetrationSphereVsSphere(shape_b, shape_a, 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.posA.set(pos.v[0], pos.v[1], pos.v[2]);
    cinfo.posB.set(pos.v[0], pos.v[1], pos.v[2]);
    cinfo.dir.set(dir.v[0], dir.v[1], dir.v[2]);
    cinfo.depth = depth;
  }

  return !intersect;
}

bool CheckPenetrationSphereVsSphere(
    const SimShape& shape_a,
    const SimShape& shape_b,
    CollisionInfo& cinfo) {
  assert(shape_a.mType == SimShape::Sphere);
  assert(
      shape_a.scale[0] == shape_a.scale[1]
      && shape_a.scale[1] == shape_a.scale[2]);
  assert(shape_b.mType == SimShape::Sphere);
  assert(
      shape_b.scale[0] == shape_b.scale[1]
      && shape_b.scale[1] == shape_b.scale[2]);

  Vector3d diff_pos = shape_b.pos - shape_a.pos;
  double diff_pos_norm = diff_pos.norm();
  double distance = diff_pos_norm - (shape_a.scale[0] + shape_b.scale[0]) * 0.5;

  if (distance < cCollisionEps) {
    cinfo.dir = diff_pos / diff_pos_norm;
    cinfo.posA = shape_a.pos + cinfo.dir * shape_a.scale[0] * 0.5;
    cinfo.posB = shape_b.pos - cinfo.dir * shape_b.scale[0] * 0.5;
    cinfo.depth = fabs(distance);

    return true;
  }

  return false;
}

bool CheckPenetrationSphereVsPlane(
    const SimShape& shape_a,
    const SimShape& shape_b,
    CollisionInfo& cinfo) {
  assert(shape_a.mType == SimShape::Sphere);
  assert(shape_b.mType == SimShape::Plane);

  // For now only support aligned spheres
  assert(
      (shape_a.orientation - Quaternion(0., 0., 0., 1.)).squaredNorm()
      < cCollisionEps);

  Vector3d plane_normal = shape_b.orientation.toMatrix().block(0, 1, 3, 1);
  Vector3d plane_point = shape_b.pos;
  Vector3d sphere_point_to_plane =
      shape_a.pos - plane_normal * shape_a.scale[0] * 0.5;

  double sphere_center_height =
      plane_normal.transpose() * (sphere_point_to_plane - plane_point);
  if (sphere_center_height < cCollisionEps) {
    cinfo.dir = -plane_normal;
    cinfo.depth = sphere_center_height;

    cinfo.posA = sphere_point_to_plane;
    cinfo.posB = sphere_point_to_plane - sphere_center_height * plane_normal;

    return 1;
  }

  return 0;
}

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));
  }
}

bool SolveGaussSeidelProj(
    MatrixNd& A,
    VectorNd& b,
    VectorNd& x,
    VectorNd& xlo,
    VectorNd& xhi,
    double tol = 1.0e-6,
    int maxiter = 100) {
  int n = A.cols();
  for (int iter = 0; iter < maxiter; iter++) {
    double residual = 0.;
    for (int i = 0; i < n; i++) {
      double sigma = 0.;
      for (int j = 0; j < i; j++) {
        sigma += A(i, j) * x[j];
      }
      for (int j = i + 1; j < n; j++) {
        sigma += A(i, j) * x[j];
      }
      double x_old = x[i];
      x[i] = (b[i] - sigma) / A(i, i);

      // Projection onto admissible values
      if (x[i] < xlo[i]) {
        x[i] = xlo[i];
      }
      if (x[i] > xhi[i]) {
        x[i] = xhi[i];
      }

      double diff = x[i] - x_old;
      residual += diff * diff;
    }

    if (residual < tol) {
      cout << "Numiter: " << iter << endl;
      return true;
    }

    if (iter > maxiter) {
      break;
    }
  }

  return false;
}

void CalcCollisions(
    SimBody& body_a,
    SimBody& body_b,
    std::vector<CollisionInfo>& collisions) {
  collisions.clear();

  for (int i = 0, ni = body_a.mCollisionShapes.size(); i < ni; ++i) {
    const SimShape& shape_a = body_a.mCollisionShapes[i].second;

    for (int j = 0, nj = body_b.mCollisionShapes.size(); j < nj; ++j) {
      const SimShape& shape_b = body_b.mCollisionShapes[j].second;

      CollisionInfo cinfo;
      bool has_penetration = false;
      has_penetration = CheckPenetration(shape_a, shape_b, cinfo);
      cinfo.effectiveRestitution = shape_a.restitution * shape_b.restitution;

      if (has_penetration) {
        cinfo.mBodyA = &body_a;
        cinfo.mBodyAIndex = body_a.mCollisionShapes[i].first;
        cinfo.mBodyB = &body_b;
        cinfo.mBodyBIndex = body_b.mCollisionShapes[j].first;
        collisions.push_back(cinfo);
      }
    }
  }
}

void CalcImpulseVariables(
    SimBody* body,
    unsigned int body_index,
    const Vector3d& pos,
    const Vector3d& dir,
    MatrixNd* MInv,
    VectorNd* jac,
    double* G_MInv_GT) {
  if (body == nullptr || body->mIsStatic) {
    jac->setZero();
    *G_MInv_GT = 0.;
    return;
  }

  Model* model = &body->mModel;
  int ndof = model->qdot_size;
  const VectorNd& q = body->q;
  const VectorNd& qdot = body->qdot;

  // Calculate local coordinates of the contact point
  UpdateKinematicsCustom(*model, &q, nullptr, nullptr);
  Vector3d point_local_b =
      CalcBodyToBaseCoordinates(*model, q, body_index, pos, false);

  // Compute vectors and matrices of the contact system
  MatrixNd M(MatrixNd::Zero(ndof, ndof));
  CompositeRigidBodyAlgorithm(*model, q, M, false);
  *MInv = M.inverse();

  MatrixNd G_constr(MatrixNd::Zero(3, ndof));
  CalcPointJacobian(*model, q, body_index, point_local_b, G_constr, false);
  (*jac) = dir.transpose() * G_constr;

  *G_MInv_GT = (*jac) * (*MInv) * (*jac).transpose();
}

void PrepareConstraintImpulse(
    SimBody* body_a,
    SimBody* body_b,
    CollisionInfo& cinfo) {
  CalcImpulseVariables(
      body_a,
      cinfo.mBodyAIndex,
      cinfo.posA,
      -cinfo.dir,
      &cinfo.MInvA,
      &cinfo.jacA,
      &cinfo.GMInvGTA);
  CalcImpulseVariables(
      body_b,
      cinfo.mBodyBIndex,
      cinfo.posB,
      cinfo.dir,
      &cinfo.MInvB,
      &cinfo.jacB,
      &cinfo.GMInvGTB);
  cout << "jacA = " << cinfo.jacA << endl;
  cout << "jacB = " << cinfo.jacB << endl;

}

void CalcConstraintImpulse(
    SimBody* body_a,
    SimBody* body_b,
    CollisionInfo& cinfo,
    const double dt) {
  // Todo: add nonlinear effects * dt

  double rhs = 0.;
  if (body_a && !body_a->mIsStatic) {
    rhs += (1.0 + cinfo.effectiveRestitution) * cinfo.jacA * body_a->qdot;
  }
  if (body_b && !body_b->mIsStatic) {
    rhs -= (1.0 + cinfo.effectiveRestitution) * cinfo.jacB * (body_b->qdot);
  }

  double denom = cinfo.GMInvGTA + cinfo.GMInvGTB;

  if (body_a && !body_a->mIsStatic) {
    double old_impulse = cinfo.accumImpulseA;
    cinfo.deltaImpulseA = rhs / denom;
    cinfo.accumImpulseA += cinfo.deltaImpulseA;
    cinfo.accumImpulseA = std::max(0., cinfo.accumImpulseA);
    cinfo.deltaImpulseA = cinfo.accumImpulseA - old_impulse;
  }

  if (body_b && !body_b->mIsStatic) {
    double old_impulse = cinfo.accumImpulseB;
    cinfo.deltaImpulseB = rhs / denom;
    cinfo.accumImpulseB += cinfo.deltaImpulseB;
    cinfo.accumImpulseB = std::max(0., cinfo.accumImpulseB);
    cinfo.deltaImpulseB = cinfo.accumImpulseB - old_impulse;
  }
}

void ApplyConstraintImpulse(
    SimBody* body_a,
    SimBody* body_b,
    CollisionInfo& cinfo) {
  if (body_a && !body_a->mIsStatic) {
    body_a->qdot += cinfo.MInvA * cinfo.jacA.transpose()
                    * (cinfo.deltaImpulseA);
  }

  if (body_b && !body_b->mIsStatic) {
    body_b->qdot -= cinfo.MInvB * cinfo.jacB.transpose()
                    * (cinfo.deltaImpulseB);
  }
}

void World::calcUnconstrainedVelUpdate(double dt) {
  for (SimBody& body : mBodies) {
    ForwardDynamics(body.mModel, body.q, body.qdot, body.tau, body.qddot);

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

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

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

  for (int i = 0, n = mBodies.size(); i < n; ++i) {
    SimBody& ref_body = mBodies[i];

    // Check collisions against moving bodies
    for (int j = i + 1; j < n; ++j) {
      SimBody& other_body = mBodies[j];
      CalcCollisions(ref_body, other_body, mContactPoints);
    }

    // Check collisions against static bodies
    for (SimBody::BodyCollisionInfo& body_col_info :
         ref_body.mCollisionShapes) {
      SimShape& ref_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, ref_body_shape, cinfo);
        cinfo.effectiveRestitution = ref_body_shape.restitution;

        if (has_penetration) {
          cinfo.mBodyA = nullptr;
          cinfo.mBodyAIndex = -1;
          cinfo.mBodyB = &ref_body;
          cinfo.mBodyBIndex = body_col_info.first;
          mContactPoints.push_back(cinfo);
        }
      }
    }
  }
}

void World::resolveCollisions(double dt) {
  for (CollisionInfo& cinfo : mContactPoints) {
    PrepareConstraintImpulse(cinfo.mBodyA, cinfo.mBodyB, cinfo);
  }

  int num_iter = 20;
  for (int i = 0; i < num_iter; i++) {
    for (CollisionInfo& cinfo : mContactPoints) {
      CalcConstraintImpulse(cinfo.mBodyA, cinfo.mBodyB, cinfo, dt);
      ApplyConstraintImpulse(cinfo.mBodyA, cinfo.mBodyB, cinfo);
    }
  }
}

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

  mSimTime += dt;
  return true;
}

SimBody CreateSphereBody(
    double mass,
    double radius,
    double restitution,
    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);
  shape.restitution = restitution;

  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