#include <igl/opengl/gl.h>
#include <igl/arap.h>
#include <igl/biharmonic_coordinates.h>
#include <igl/cat.h>
#include <igl/cotmatrix.h>
#include <igl/massmatrix.h>
#include <igl/matrix_to_list.h>
#include <igl/parula.h>
#include <igl/point_mesh_squared_distance.h>
#include <igl/readDMAT.h>
#include <igl/readMESH.h>
#include <igl/remove_unreferenced.h>
#include <igl/writeDMAT.h>
#include <igl/opengl/glfw/Viewer.h>
#include <Eigen/Sparse>
#include <iostream>
#include <queue>


struct Mesh
{
  Eigen::MatrixXd V,U;
  Eigen::MatrixXi T,F;
} low,high,scene;

Eigen::MatrixXd W;
igl::ARAPData arap_data;

int main(int argc, char * argv[])
{
  using namespace Eigen;
  using namespace std;
  using namespace igl;
  
  // read the mesh, if the code is prepared outside of tutorial, the TUTORIAL_SHARED_PATH
  // should be the data folder
  if(!readMESH(TUTORIAL_SHARED_PATH "/octopus-low.mesh",low.V,low.T,low.F))
  {
    cout<<"failed to load mesh"<<endl;
  }
  if(!readMESH(TUTORIAL_SHARED_PATH "/octopus-high.mesh",high.V,high.T,high.F))
  {
    cout<<"failed to load mesh"<<endl;
  }

  // Precomputation
  {
    Eigen::VectorXi b;
    {
      // this will create a vector from 0 to V.rows()-1 where the gap is 1
      Eigen::VectorXi J = Eigen::VectorXi::LinSpaced(high.V.rows(),0,high.V.rows()-1);
      Eigen::VectorXd sqrD;
      Eigen::MatrixXd _2;
      cout<<"Finding closest points..."<<endl;
      // using J which is N by 1 instead of a matrix that represents faces of N by 3
      // so that we will find the closest vertices istead of closest point on the face
      // so far the two meshes are not seperated. So what we are really doing here
      // is computing handles from low resolution and use that for the high resolution one
      igl::point_mesh_squared_distance(low.V,high.V,J,sqrD,b,_2);
      assert(sqrD.minCoeff() < 1e-7 && "low.V should exist in high.V");
    }
    // force perfect positioning, rather have popping in low-res than high-res.
    // The correct/elaborate thing to do is express original low.V in terms of
    // linear interpolation (or extrapolation) via elements in (high.V,high.F)

    // this is to replace the vertices on low resolution
    // with the vertices in high resolution. b is the list of vertices
    // corresponding to the indices in high resolution which has closest
    // distance to the points in low resolution
    low.V = high.V(b,Eigen::all);

    // list of points --> list of singleton lists
    std::vector<std::vector<int> > S;
    // S will hav size of low.V.rows() and each list inside will have 1 element
    igl::matrix_to_list(b,S);
    cout<<"Computing weights for "<<b.size()<<
      " handles at "<<high.V.rows()<<" vertices..."<<endl;
    // Technically k should equal 3 for smooth interpolation in 3d, but 2 is
    // faster and looks OK
    const int k = 2;

    // using all the points in low resolution as handles for the region
    // it will be too expansive to use all the points in high reolution as handles
    // but since low and high resembles the same thing, using points in low reesolution
    // will give you similar performance
    igl::biharmonic_coordinates(high.V,high.T,S,k,W);
    cout<<"Reindexing..."<<endl;
    // Throw away interior tet-vertices, keep weights and indices of boundary
    VectorXi I,J;
    igl::remove_unreferenced(high.V.rows(),high.F,I,J);
    for_each(high.F.data(),high.F.data()+high.F.size(),[&I](int & a){a=I(a);});
    for_each(b.data(),b.data()+b.size(),[&I](int & a){a=I(a);});
    high.V = high.V(J,Eigen::all).eval();
    W = W(J,Eigen::all).eval();
  }

  // Resize low res (high res will also be resized by affine precision of W)
  low.V.rowwise() -= low.V.colwise().mean();
  low.V /= (low.V.maxCoeff()-low.V.minCoeff());
  low.V.rowwise() += RowVector3d(0,1,0);
  low.U = low.V;
  high.U = high.V;

  arap_data.with_dynamics = true;
  arap_data.max_iter = 10;
  arap_data.energy = ARAP_ENERGY_TYPE_DEFAULT;
  arap_data.h = 0.01;
  arap_data.ym = 0.001;
  if(!arap_precomputation(low.V,low.T,3,VectorXi(),arap_data))
  {
    cerr<<"arap_precomputation failed."<<endl;
    return EXIT_FAILURE;
  }
  // Constant gravitational force
  Eigen::SparseMatrix<double> M;
  igl::massmatrix(low.V,low.T,igl::MASSMATRIX_TYPE_DEFAULT,M);
  const size_t n = low.V.rows();
  // f = ma
  arap_data.f_ext =  M * RowVector3d(0,-9.8,0).replicate(n,1);
  // Random initial velocities to wiggle things
  arap_data.vel = MatrixXd::Random(n,3);
  
  igl::opengl::glfw::Viewer viewer;
  // Create one huge mesh containing both meshes
  igl::cat(1,low.U,high.U,scene.U);
  // need to remap the indices since we cat the V matrices
  igl::cat(1,low.F,MatrixXi(high.F.array()+low.V.rows()),scene.F);
  // Color each mesh
  viewer.data().set_mesh(scene.U,scene.F);
  MatrixXd C(scene.F.rows(),3);
  C<<
    RowVector3d(0.8,0.5,0.2).replicate(low.F.rows(),1),
    RowVector3d(0.3,0.4,1.0).replicate(high.F.rows(),1);
  viewer.data().set_colors(C);

  viewer.callback_key_pressed = 
    [&](igl::opengl::glfw::Viewer & viewer,unsigned int key,int mods)->bool
  {
    switch(key)
    {
      default: 
        return false;
      case ' ':
        viewer.core().is_animating = !viewer.core().is_animating;
        return true;
      case 'r':
        low.U = low.V;
        return true;
    }
  };
  viewer.callback_pre_draw = [&](igl::opengl::glfw::Viewer & viewer)->bool
  {
    glEnable(GL_CULL_FACE);
    if(viewer.core().is_animating)
    {
      arap_solve(MatrixXd(0,3),arap_data,low.U);
      for(int v = 0;v<low.U.rows();v++)
      {
        // collide with y=0 plane
        const int y = 1;
        if(low.U(v,y) < 0)
        {
          low.U(v,y) = -low.U(v,y);
          // ~ coefficient of restitution
          const double cr = 1.1;
          arap_data.vel(v,y) = - arap_data.vel(v,y) / cr;
        }
      }

      scene.U.block(0,0,low.U.rows(),low.U.cols()) = low.U;
      high.U = W * (low.U.rowwise() + RowVector3d(1,0,0));
      scene.U.block(low.U.rows(),0,high.U.rows(),high.U.cols()) = high.U;

      viewer.data().set_vertices(scene.U);
      viewer.data().compute_normals();
    }
    return false;
  };
  viewer.data().show_lines = false;
  viewer.core().is_animating = true;
  viewer.core().animation_max_fps = 30.;
  viewer.data().set_face_based(true);
  cout<<R"(
[space] to toggle animation
'r'     to reset positions 
      )";
  viewer.core().rotation_type = 
    igl::opengl::ViewerCore::ROTATION_TYPE_TWO_AXIS_VALUATOR_FIXED_UP;
  viewer.launch();
}