Vector2s RigidBody2DSim::computeTotalMomentum() const
{
  VectorXs p;
  computeMomentum( m_state.v(), p );
  assert( p.size() == 2 );
  return p;
}

// This method and the smooth version share the second half of code. Abstract that out.
void StaticPlaneSphereConstraint::computeGeneralizedFrictionDisk( const VectorXs& q, const VectorXs& v, const int start_column, const int num_samples, SparseMatrixsc& D, VectorXs& drel ) const
{
  assert( start_column >= 0 );
  assert( start_column < D.cols() );
  assert( num_samples > 0 );
  assert( start_column + num_samples - 1 < D.cols() );
  assert( q.size() % 12 == 0 );
  assert( q.size() == 2 * v.size() );

  const Vector3s n{ m_plane.n() };
  assert( fabs( n.norm() - 1.0 ) <= 1.0e-6 );

  std::vector<Vector3s> friction_disk( static_cast<std::vector<Vector3s>::size_type>( num_samples ) );
  assert( friction_disk.size() == std::vector<Vector3s>::size_type( num_samples ) );
  {
    // Compute the relative velocity
    Vector3s tangent_suggestion{ computeRelativeVelocity( q, v ) };
    if( tangent_suggestion.cross( n ).squaredNorm() < 1.0e-9 )
    {
      tangent_suggestion = FrictionUtilities::orthogonalVector( n );
    }
    tangent_suggestion *= -1.0;

    // Sample the friction disk
    FrictionUtilities::generateOrthogonalVectors( n, friction_disk, tangent_suggestion );
  }
  assert( unsigned( num_samples ) == friction_disk.size() );

  // Compute the displacement from the center of mass to the point of contact
  assert( fabs( n.norm() - 1.0 ) <= 1.0e-10 ); assert( m_r >= 0.0 );
  const Vector3s r_world{ - m_r * n };

  // Cache the velocity of the collision point on the plane
  const Vector3s plane_collision_point_vel{ computePlaneCollisionPointVelocity( q ) };

  // For each sample of the friction disk
  const unsigned nbodies{ static_cast<unsigned>( q.size() / 12 ) };
  for( unsigned friction_sample = 0; friction_sample < unsigned( num_samples ); ++friction_sample )
  {
    const unsigned cur_col{ start_column + friction_sample };
    assert( cur_col < unsigned( D.cols() ) );

    // Effect on center of mass
    D.insert( 3 * m_sphere_idx + 0, cur_col ) = friction_disk[friction_sample].x();
    D.insert( 3 * m_sphere_idx + 1, cur_col ) = friction_disk[friction_sample].y();
    D.insert( 3 * m_sphere_idx + 2, cur_col ) = friction_disk[friction_sample].z();

    // Effect on orientation
    {
      const Vector3s ntilde{ r_world.cross( friction_disk[friction_sample] ) };
      D.insert( 3 * ( nbodies + m_sphere_idx ) + 0, cur_col ) = ntilde.x();
      D.insert( 3 * ( nbodies + m_sphere_idx ) + 1, cur_col ) = ntilde.y();
      D.insert( 3 * ( nbodies + m_sphere_idx ) + 2, cur_col ) = ntilde.z();
    }

    // Relative velocity contribution from kinematic scripting
    assert( cur_col < drel.size() );
    drel( cur_col ) = - friction_disk[friction_sample].dot( plane_collision_point_vel );
  }
}

void RigidBody2DSim::computeBodyBodyActiveSetSpatialGrid( const VectorXs& q0, const VectorXs& q1, const VectorXs& v, std::vector<std::unique_ptr<Constraint>>& active_set ) const
{
  assert( q0.size() % 3 == 0 ); assert( q0.size() == q1.size() );

  const unsigned nbodies{ static_cast<unsigned>( q0.size() / 3 ) };

  // Candidate bodies that might overlap
  std::set<std::pair<unsigned,unsigned>> possible_overlaps;
  {
    // Compute an AABB for each body
    std::vector<AABB> aabbs;
    aabbs.reserve( nbodies );
    for( unsigned bdy_idx = 0; bdy_idx < nbodies; ++bdy_idx )
    {
      Array2s min;
      Array2s max;
      m_state.bodyGeometry( bdy_idx )->computeCollisionAABB( q0.segment<2>( 3 * bdy_idx ), q0( 3 * bdy_idx + 2 ), q1.segment<2>( 3 * bdy_idx ), q1( 3 * bdy_idx + 2 ), min, max );
      aabbs.emplace_back( min, max );
    }
    assert( aabbs.size() == nbodies );

    // Determine which bodies possibly overlap
    SpatialGrid::getPotentialOverlaps( aabbs, possible_overlaps );
  }

  // Create constraints for bodies that actually overlap
  for( const auto& possible_overlap_pair : possible_overlaps )
  {
    assert( possible_overlap_pair.first < nbodies );
    assert( possible_overlap_pair.second < nbodies );

    // We can run standard narrow phase
    dispatchNarrowPhaseCollision( possible_overlap_pair.first, possible_overlap_pair.second, q0, q1, v, active_set );
  }
}

void IntegrationTools::exponentialEuler( const VectorXs& q0, const VectorXs& v0, const std::vector<bool>& fixed, const scalar& dt, VectorXs& q1 )
{
  assert( q0.size() == q1.size() );
  assert( q0.size() % 12 == 0 );
  assert( v0.size() * 2 == q0.size() );
  assert( 12 * fixed.size() == static_cast<unsigned long>(q0.size()) );

  const int nbodies{ static_cast<int>( fixed.size() ) };

  // Center of mass update
  for( int bdy_num = 0; bdy_num < nbodies; bdy_num++ )
  {
    if( !fixed[bdy_num] )
    {
      q1.segment<3>( 3 * bdy_num ) = q0.segment<3>( 3 * bdy_num ) + dt * v0.segment<3>( 3 * bdy_num );
    }
    else
    {
      q1.segment<3>( 3 * bdy_num ) = q0.segment<3>( 3 * bdy_num );
    }
  }

  // Orientation update
  for( int bdy_num = 0; bdy_num < nbodies; bdy_num++ )
  {
    if( !fixed[bdy_num] )
    {
      updateOrientation( nbodies, bdy_num, q0, v0, dt, q1 );
    }
    else
    {
      q1.segment<9>( 3 * nbodies + 9 * bdy_num ) = q0.segment<9>( 3 * nbodies + 9 * bdy_num );
    }
  }
}

void SpringForce::addGradEToTotal( const VectorXs& x, const VectorXs& v, const VectorXs& m, VectorXs& gradE )
{
  assert( x.size() == v.size() );
  assert( x.size() == gradE.size() );
  assert( x.size()%2 == 0 );
  assert( m_endpoints.first >= 0 );  assert( m_endpoints.first < x.size()/2 );
  assert( m_endpoints.second >= 0 ); assert( m_endpoints.second < x.size()/2 );

  //GET PARTICLES
  int size = 2;
  int xistart = m_endpoints.first*2;
  int xjstart = m_endpoints.second*2;  
  VectorXs xi = x.segment(xistart, size);
  VectorXs xj = x.segment(xjstart, size);
	
  //SPRING FORCE
  VectorXs nhat = xi-xj;  
  double lxixj = nhat.norm();
  nhat /= lxixj;	
  VectorXs dxiU = m_k*(lxixj-m_l0)*nhat;
  VectorXs dxjU = -dxiU;  
    
  //SPRING DAMPING FORCE 
  VectorXs vi = v.segment(xistart, size);
  VectorXs vj = v.segment(xjstart, size);
  double lvivj = (vi-vj).dot(nhat);
  dxiU += m_b * nhat * lvivj; //F-
  dxjU += -m_b * nhat * lvivj; //F+
	
  //CHANGE SYSTEM
  gradE[xistart+0] += dxiU[0];
  gradE[xistart+1] += dxiU[1];	
  gradE[xjstart+0] += dxjU[0];
  gradE[xjstart+1] += dxjU[1];
}

// Responds to a collision detected between a particle and a half-plane by 
// applying an impulse to the velocity of the particle.
// Inputs:
//   scene: The scene data structure.
//   vidx:  The index of the particle.
//   pidx:  The index of the half-plane.
//   n:     The shortest vector between the particle and the half-plane.
// Outputs:
//   None.
void SimpleCollisionHandler::respondParticleHalfplane(TwoDScene &scene, int vidx, int pidx, const VectorXs &n)
{
  VectorXs nhat = n;
  nhat.normalize();
  double cfactor = (1.0+getCOR())/2.0;
  scene.getV().segment<2>(2*vidx) -= 2*cfactor*scene.getV().segment<2>(2*vidx).dot(nhat)*nhat;
}

// Adds the gradient of the penalty potential (-1 * force) for a pair of 
// particles to the total.
// Read the positions of the particles from the input variable x. Radii can
// be obtained from the member variable m_scene, the penalty force stiffness 
// from member variable m_k, and penalty force thickness from member variable
// m_thickness.
// Inputs:
//   x:    The positions of the particles in the scene. 
//   idx1: The index of the first particle, i.e. the position of this particle
//         is ( x[2*idx1], x[2*idx1+1] ).
//   idx2: The index of the second particle.
// Outputs:
//   gradE: The total gradient of penalty force. *ADD* the particle-particle
//          gradient to this total gradient.
void PenaltyForce::addParticleParticleGradEToTotal(const VectorXs &x, int idx1, int idx2, VectorXs &gradE)
{
//    VectorXs x1 = x.segment<2>(2*idx1);
//    VectorXs x2 = x.segment<2>(2*idx2);
//    
//    double r1 = m_scene.getRadius(idx1);
//    double r2 = m_scene.getRadius(idx2);
//    
//    // your implementation here
//    
    VectorXs x1 = x.segment<2>(2*idx1);
    VectorXs x2 = x.segment<2>(2*idx2);
    
    double r1 = m_scene.getRadius(idx1);
    double r2 = m_scene.getRadius(idx2);
    
    VectorXs n = x2-x1;
    VectorXs nhat = n;
    nhat.normalize();
    if(n.norm() < r1 + r2 + m_thickness)
    {
        gradE.segment<2>(2*idx1) -= m_k * (n.norm() - r1 - r2 - m_thickness) * nhat;
        gradE.segment<2>(2*idx2) += m_k * (n.norm() - r1 - r2 - m_thickness) * nhat;
    }
}

void BodyBodyConstraint::evalgradg( const VectorXs& q, const int col, SparseMatrixsc& G, const FlowableSystem& fsys ) const
{
  assert( q.size() % 12 == 0 );
  assert( col >= 0 );
  assert( col < G.cols() );

  const unsigned nbodies{ static_cast<unsigned>( q.size() / 12 ) };

  // MUST BE ADDED GOING DOWN THE COLUMN. DO NOT TOUCH ANOTHER COLUMN.
  {
    assert( 3 * nbodies + 3 * m_idx0 + 2 < unsigned( G.rows() ) );
    G.insert( 3 * m_idx0 + 0, col ) = m_n.x();
    G.insert( 3 * m_idx0 + 1, col ) = m_n.y();
    G.insert( 3 * m_idx0 + 2, col ) = m_n.z();
    const Vector3s ntilde_0{ m_r0.cross( m_n ) };
    G.insert( 3 * ( m_idx0 + nbodies ) + 0, col ) = ntilde_0.x();
    G.insert( 3 * ( m_idx0 + nbodies ) + 1, col ) = ntilde_0.y();
    G.insert( 3 * ( m_idx0 + nbodies ) + 2, col ) = ntilde_0.z();
  }

  {
    assert( 3 * nbodies + 3 * m_idx1 + 2 < unsigned( G.rows() ) );
    G.insert( 3 * m_idx1 + 0, col ) = - m_n.x();
    G.insert( 3 * m_idx1 + 1, col ) = - m_n.y();
    G.insert( 3 * m_idx1 + 2, col ) = - m_n.z();
    const Vector3s ntilde_1{ m_r1.cross( m_n ) };
    G.insert( 3 * ( m_idx1 + nbodies ) + 0, col ) = - ntilde_1.x();
    G.insert( 3 * ( m_idx1 + nbodies ) + 1, col ) = - ntilde_1.y();
    G.insert( 3 * ( m_idx1 + nbodies ) + 2, col ) = - ntilde_1.z();
  }
}

VectorXs Constraint::computeWorldSpaceContactNormal( const VectorXs& q ) const
{
  VectorXs n;
  getWorldSpaceContactNormal( q, n );
  assert( fabs( n.norm() - 1.0 ) <= 1.0e-6 );
  return n;
}

void VortexForce::addHessXToTotal( const VectorXs& x, const VectorXs& v, const VectorXs& m, MatrixXs& hessE )
{
  assert( x.size() == v.size() );
  assert( x.size() == m.size() );
  assert( x.size() == hessE.rows() );
  assert( x.size() == hessE.cols() );
  assert( x.size()%2 == 0 );

//  scalar m1 = m(2*m_particles.second);
//  scalar m2 = m(2*m_particles.first);
//  
//  Vector2s nhat = x.segment<2>(2*m_particles.second)-x.segment<2>(2*m_particles.first); 
//  scalar r = nhat.norm(); 
//  assert( r != 0.0 ); 
//  nhat /= r;
//
//  Matrix2s entry = Matrix2s::Identity()-3.0*nhat*nhat.transpose();
//  entry *= m_G*m1*m2/r*r*r;
//  
//  hessE.block<2,2>(2*m_particles.first,2*m_particles.first)   += entry;
//  hessE.block<2,2>(2*m_particles.second,2*m_particles.second) += entry;
//  hessE.block<2,2>(2*m_particles.first,2*m_particles.second)  -= entry;
//  hessE.block<2,2>(2*m_particles.second,2*m_particles.first)  -= entry;

  std::cerr << outputmod::startred << "ERROR IN VORTEXFORCE: " << outputmod::endred << "No addHessXToTotal defined for VortexForce." << std::endl;
  exit(1);
}

scalar RigidBody2DSim::computeTotalAngularMomentum() const
{
  VectorXs L;
  computeAngularMomentum( m_state.v(), L );
  assert( L.size() == 1 );
  return L( 0 );
}

scalar HertzianPenaltyForce::computePotential( const VectorXs& q, const SparseMatrixsc& M, const VectorXs& r ) const
{
  assert( q.size() % 2 == 0 ); assert( q.size() == M.rows() ); assert( q.size() == M.cols() ); assert( r.size() == q.size() / 2 );

  scalar U{ 0.0 };
  // For each ball
  for( unsigned ball0 = 0; ball0 < r.size(); ++ball0 )
  {
    // For each subsequent ball
    for( unsigned ball1 = ball0 + 1; ball1 < r.size(); ++ball1 )
    {
      // Compute the total radius
      const scalar total_radius{ r(ball0) + r(ball1) };
      // Compute a vector pointing from ball0 to ball1
      const Vector2s n{ q.segment<2>( 2 * ball1 ) - q.segment<2>( 2 * ball0 ) };
      // If the squared distance is greater or equal to the sum of the radii squared, no force
      if( n.squaredNorm() > total_radius * total_radius )
      {
        continue;
      }
      // Compute the penetration depth
      const scalar delta{ n.norm() - total_radius };
      assert( delta < 0.0 );
      // U = 0.5 * k * pen_depth^(5/2)
      U += 0.5 * m_k * std::pow( -delta, scalar( 2.5 ) );
    }
  }

  return U;
}

void DragDampingForce::addEnergyToTotal( const VectorXs& x, const VectorXs& v, const VectorXs& m, scalar& E )
{
  assert( x.size() == v.size() );
  assert( x.size() == m.size() );
  assert( x.size()%3 == 0 );

  std::cerr << outputmod::startred << "WARNING IN DRAGDAMPINGFORCE: " << outputmod::endred << "No energy defined for DragDampingForce." << std::endl;
}

void FrictionOperatorUtilities::computeMDPLambda( const VectorXs& vrel, VectorXs& lambda )
{
  assert( vrel.size() % 2 == 0 );
  lambda.conservativeResize( vrel.size() / 2 );
  for( int con_num = 0; con_num < lambda.size(); ++con_num )
  {
    lambda( con_num ) = vrel.segment<2>( 2 * con_num ).norm();
  }
}

void SpringForce::addEnergyToTotal( const VectorXs& x, const VectorXs& v, const VectorXs& m, scalar& E )
{
  assert( x.size() == v.size() );
  assert( x.size()%2 == 0 );
  assert( m_endpoints.first >= 0 );  assert( m_endpoints.first < x.size()/2 );
  assert( m_endpoints.second >= 0 ); assert( m_endpoints.second < x.size()/2 );

  // Add milestone 2 code here.
}

void TwoDScene::accumulateddUdxdv( MatrixXs& A, const VectorXs& dx, const VectorXs& dv )
{
  assert( A.rows() == m_x.size() );
  assert( A.cols() == m_x.size() );
  assert( dx.size() == dv.size() );
  assert( dx.size() == 0 || dx.size() == A.rows() );

  if( dx.size() == 0 ) for( std::vector<Force*>::size_type i = 0; i < m_forces.size(); ++i ) m_forces[i]->addHessVToTotal( m_x, m_v, m_m, A );
  else                 for( std::vector<Force*>::size_type i = 0; i < m_forces.size(); ++i ) m_forces[i]->addHessVToTotal( m_x+dx, m_v+dv, m_m, A );
}

void DragDampingForce::addHessXToTotal( const VectorXs& x, const VectorXs& v, const VectorXs& m, MatrixXs& hessE )
{
  assert( x.size() == v.size() );
  assert( x.size() == m.size() );
  assert( x.size() == hessE.rows() );
  assert( x.size() == hessE.cols() );
  assert( x.size()%3 == 0 );

  // Nothing to do.
}

void Constraint::evalKinematicRelVelGivenBases( const VectorXs& q, const VectorXs& v, const std::vector<std::unique_ptr<Constraint>>& constraints, const MatrixXXsc& bases, VectorXs& nrel, VectorXs& drel )
{
  assert( bases.cols() == nrel.size() + drel.size() );
  assert( nrel.size() % constraints.size() == 0 );
  assert( drel.size() % constraints.size() == 0 );

  // Number of constraints in the system
  const unsigned ncons{ static_cast<unsigned>( constraints.size() ) };

  // Number of tangent friction samples per constraint in the system
  const unsigned friction_vectors_per_con{ static_cast<unsigned>( drel.size() / ncons ) };

  for( unsigned con_num = 0; con_num < ncons; ++con_num )
  {
    // Grab the kinematic relative velocity
    const VectorXs kinematic_rel_vel{ constraints[con_num]->computeKinematicRelativeVelocity( q, v ) };
    assert( kinematic_rel_vel.size() == bases.rows() );

    // Compute the column of the normal in the bases matrix
    const unsigned n_idx{ ( friction_vectors_per_con + 1 ) * con_num };
    assert( n_idx < bases.cols() ); assert( fabs( bases.col( n_idx ).norm() - 1.0 ) <= 1.0e-6 );
    // Project the relative velocity onto the normal
    nrel( con_num ) = - kinematic_rel_vel.dot( bases.col( n_idx ) );
    // For each tangent friction sample
    for( unsigned friction_sample = 0; friction_sample < friction_vectors_per_con; ++friction_sample )
    {
      // Compute the column of the current friction sample in the bases matrix
      const unsigned f_idx{ ( friction_vectors_per_con + 1 ) * con_num + friction_sample + 1 };
      assert( f_idx < bases.cols() );
      assert( fabs( bases.col( f_idx ).norm() - 1.0 ) <= 1.0e-6 );
      assert( fabs( bases.col( n_idx ).dot( bases.col( f_idx ) ) ) <= 1.0e-6 );
      drel( friction_vectors_per_con * con_num + friction_sample ) = - kinematic_rel_vel.dot( bases.col( f_idx ) );
    }
  }
}

void VortexForce::addHessVToTotal( const VectorXs& x, const VectorXs& v, const VectorXs& m, MatrixXs& hessE )
{
  assert( x.size() == v.size() );
  assert( x.size() == m.size() );
  assert( x.size() == hessE.rows() );
  assert( x.size() == hessE.cols() );
  assert( x.size()%2 == 0 );

  std::cerr << outputmod::startred << "ERROR IN VORTEXFORCE: " << outputmod::endred << "No addHessXToTotal defined for VortexForce." << std::endl;
  exit(1);
}

void DragDampingForce::addHessVToTotal( const VectorXs& x, const VectorXs& v, const VectorXs& m, MatrixXs& hessE )
{
  assert( x.size() == v.size() );
  assert( x.size() == m.size() );
  assert( x.size() == hessE.rows() );
  assert( x.size() == hessE.cols() );
  assert( x.size()%2 == 0 );
  
  // Compute the force Jacboian here!
   hessE.diagonal().array() += m_b;
}