#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) {
  // 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(
    const SimBody& body_a,
    const 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);

      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->mIsStatic) {
    *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);
}

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

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

  double denom = cinfo.GMInvGTA + cinfo.GMInvGTB;

  if (!body_a->mIsStatic) {
    cinfo.accumImpulseA = -rhs / denom * cinfo.dir;
  }

  if (!body_b->mIsStatic) {
    cinfo.accumImpulseB = rhs / denom * cinfo.dir;
  }
}

void ApplyConstraintImpulse(
    SimBody* body,
    const MatrixNd& MInv,
    const VectorNd& jac,
    double impulse) {}

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

  if (!body_b->mIsStatic) {
    body_b->qdot += cinfo.MInvB * cinfo.jacB.transpose()
                    * (-cinfo.dir.transpose() * cinfo.accumImpulseB);
  }
}

void SimBody::resolveCollisions(
    double dt,
    std::vector<CollisionInfo>& collisions) {
  if (collisions.size() == 0) {
    // No contacts, calculate new qdot using simple forward
    // dynamics
    ForwardDynamics(mModel, q, qdot, tau, qddot);

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

    return;
  }

  int ndof = mModel.qdot_size;
  int nconstraints = collisions.size();

  // Allocate space for the constraint system
  MatrixNd M(MatrixNd::Zero(ndof, ndof));
  MatrixNd Minv(MatrixNd::Zero(ndof, ndof));

  VectorNd N(VectorNd::Zero(ndof));
  MatrixNd G(MatrixNd::Zero(nconstraints, ndof));
  VectorNd gamma(VectorNd::Zero(nconstraints));

  // Calculate local coordinates of the contact point
  std::vector<Vector3d> pos_local;
  VectorNd constr_value(VectorNd::Zero(nconstraints));
  UpdateKinematicsCustom(mModel, &q, nullptr, nullptr);
  for (int i = 0; i < nconstraints; i++) {
    constr_value[i] = -collisions[i].depth;
    pos_local.push_back(CalcBaseToBodyCoordinates(
        mModel,
        q,
        collisions[i].mBodyAIndex,
        collisions[i].posA,
        false));
  }

  // Calculate predicted position state
  VectorNd q_next_pred(q);
  calcNextPositions(dt, qdot, q_next_pred);

  // Compute vectors and matrices of the contact system
  NonlinearEffects(mModel, q_next_pred, qdot, N);
  CompositeRigidBodyAlgorithm(mModel, q_next_pred, M, false);
  Minv = M.inverse();

  // Calculate contact Jacobians
  for (int i = 0; i < nconstraints; i++) {
    MatrixNd G_constr(MatrixNd::Zero(3, ndof));
    CalcPointJacobian(
        mModel,
        q_next_pred,
        collisions[i].mBodyAIndex,
        pos_local[i],
        G_constr,
        false);
    G.block(i, 0, 1, ndof) = collisions[i].dir.transpose() * G_constr;
  }

  MatrixNd A(MatrixNd::Zero(nconstraints, nconstraints));
  VectorNd b(VectorNd::Zero(nconstraints));

  // Solve for the impules hlambda
  A = G * Minv * G.transpose();
  b = (constr_value)*1. / dt + G * (qdot + Minv * dt * (tau - N));

  VectorNd hlambda(VectorNd::Zero(nconstraints));
  VectorNd hlambda_lo(VectorNd::Constant(nconstraints, 0.));
  VectorNd hlambda_hi(VectorNd::Constant(nconstraints, 1.0e9));

  bool solve_result = false;
  solve_result = SolveGaussSeidelProj(A, b, hlambda, hlambda_lo, hlambda_hi);

  //  VectorNd hlambda = A.colPivHouseholderQr().solve(b) * -1.;

  if (!solve_result) {
    cout << "Impulse Solve Failed!! " << endl;
  }
  cout << "Contact impulse: " << hlambda.transpose() << endl;

  // solve for the new velocity
  qdot = qdot + Minv * (dt * tau + dt * N + G.transpose() * hlambda);
}

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.mBodyAIndex = -1;
          cinfo.mBodyB = &body;
          cinfo.mBodyBIndex = body_col_info.first;
          mContactPoints.push_back(cinfo);
        }
      }
    }
  }
}

void World::resolveCollisions(double dt) {
  // so far only solve Body vs World collisions
  for (SimBody& body : mBodies) {
    // collect all collisions for current body
    std::vector<CollisionInfo> body_collisions;

    for (const CollisionInfo cinfo : mContactPoints) {
      if (cinfo.mBodyA == &body) {
        body_collisions.push_back(cinfo);
      } else if (cinfo.mBodyB == &body) {
        // Make sure the collision info is expressed in terms
        // of mBodyA.
        CollisionInfo rev_cinfo;

        rev_cinfo.mBodyA = &body;
        rev_cinfo.mShapeA = cinfo.mShapeB;
        rev_cinfo.mBodyAIndex = cinfo.mBodyBIndex;

        rev_cinfo.mBodyB = cinfo.mBodyA;
        rev_cinfo.mShapeB = cinfo.mShapeA;
        rev_cinfo.mBodyBIndex = cinfo.mBodyAIndex;

        Vector3d temp = cinfo.posA;
        rev_cinfo.posA = cinfo.posB;
        rev_cinfo.posA = temp;
        rev_cinfo.dir = -cinfo.dir;
        rev_cinfo.depth = cinfo.depth;

        body_collisions.push_back(rev_cinfo);
      }
    }

    if (body_collisions.size() > 0) {
      cout << "Collision at t = " << mSimTime
           << ", pos = " << body_collisions[0].posA.transpose()
           << ", depth = " << body_collisions[0].depth << endl
           << " qdotpre  = " << mBodies[0].qdot.transpose() << endl;
    }
    body.resolveCollisions(dt, body_collisions);
    if (body_collisions.size() > 0) {
      cout << " qdotpost = " << mBodies[0].qdot.transpose() << endl;
    }
  }
}

bool World::integrateWorld(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