#include <geometrycentral/surface/manifold_surface_mesh.h>
#include <geometrycentral/surface/meshio.h>
#include <geometrycentral/surface/vertex_position_geometry.h>

#include <polyscope/polyscope.h>
#include <polyscope/surface_mesh.h>

using namespace geometrycentral;
using namespace geometrycentral::surface;

// == Geometry-central data
std::unique_ptr<ManifoldSurfaceMesh> mesh;
std::unique_ptr<VertexPositionGeometry> geometry;

/** Gradient of fvals, a scalar quantity on vertices. Outputs a vector quantity on one face */
Vector3 face_gradient_of(Face f, std::vector<double> fvals) {
  Vector3 normal = geometry->faceNormal(f);
  double area = geometry->faceArea(f);

  std::vector<Vertex> v_indices;
  for(Vertex v: f.adjacentVertices()) {
    v_indices.push_back(v);
  }
  
  Vector3 loc_face_gradient = {0, 0, 0};
  for (int i=0; i < 3; i++) {
    Vertex v0 = v_indices.at(i);
    Vertex v1 = v_indices.at((i+1)%3);
    Vertex v2 = v_indices.at((i+2)%3);

    loc_face_gradient += (
      fvals[v0.getIndex()]/(2*area)*
      cross(normal, geometry->vertexPositions[v2]-geometry->vertexPositions[v1])
    );
  }
  return loc_face_gradient;
}

/** Compute the divergence of U (vector on faces), outputs as result a scalar vertices */
std::vector<double> divergence(std::vector<Vector3> U) {
  std::vector<double> out;

  for (Vertex v : mesh->vertices()) {
    out.push_back(0);
  }

  for (Face f : mesh->faces()) {
    Vector3 normal = geometry->faceNormal(f);
    std::vector<Vertex> v_indices;
    for(Vertex v: f.adjacentVertices()) {
      v_indices.push_back(v);
    }

    for (int i=0; i < 3; i++) {
      Vertex v0 = v_indices.at(i);
      Vertex v1 = v_indices.at((i+1)%3);
      Vertex v2 = v_indices.at((i+2)%3);

      out[v0.getIndex()] -= dot(
        U[f.getIndex()],
        cross(normal, geometry->vertexPositions[v2]-geometry->vertexPositions[v1])
      );
    }
  }
  return out;
}

int main(int argc, char **argv) {
  // Initialize polyscope
  polyscope::init();

  // Load mesh
  std::tie(mesh, geometry) = readManifoldSurfaceMesh(argv[1]);

  // Register the mesh with polyscope
  auto sfm = polyscope::registerSurfaceMesh("Input obj",
                                 geometry->inputVertexPositions,
                                 mesh->getFaceVertexList(),
                                 polyscopePermutations(*mesh));
  geometry->requireEdgeCotanWeights();
  geometry->requireVertexPositions();
  geometry->requireFaceNormals();
  geometry->requireFaceAreas();

  auto vertCount = geometry->inputVertexPositions.size();

  std::vector<double> sin_coords;
  {//* First question : showing basic quantity
    std::vector<double> snd_coordinate;
    for (auto i=0; i < vertCount; i++) {
      auto pos = geometry->inputVertexPositions[i];
      snd_coordinate.push_back((double)pos[1]);
    }
    sfm->addVertexScalarQuantity("Second coordinate", snd_coordinate);

    float lambda = 18.88;
    for (auto i=0; i < vertCount; i++) {
      auto pos = geometry->inputVertexPositions[i];
      double posX = cos(PI / 4) * pos.x - sin(PI / 4) * pos.y;
      double posY = sin(PI / 4) * pos.x + cos(PI / 4) * pos.y;
      sin_coords.push_back(sin(lambda * (posX + 0.5)) + sin(lambda * posY));
    }
    sfm->addVertexScalarQuantity("Sinusoid data from coordinates", sin_coords);
  }

  {//* Second question : compute the area and the normal vector for each triangle
    // After doing geometry->requireFaceNormals(),
    // we can access normals with geometry->faceNormal(Face f)

    sfm->addFaceVectorQuantity("Faces normals", geometry->faceNormals);

    // After doing geometry->requireFaceAreas(),
    // we can access normals with geometry->faceArea(Face f)
    sfm->addFaceScalarQuantity("Faces area", geometry->faceAreas);
  }
  
  std::vector<Vector3> face_gradients;
  {//* Third question : compute the per face gradient of f
    for (Face f: mesh->faces()) {
      face_gradients.push_back(face_gradient_of(f, sin_coords));
    }
    sfm->addFaceVectorQuantity("Faces gradients", face_gradients);
  }
  
  {//* Fourth question : Compute and display the divergence of grad(f)
    sfm->addVertexScalarQuantity("Faces div of gradients", divergence(face_gradients));
  }

  //* Laplace-Beltrami
  {//* 6th question : implement the Laplace-Beltrami operator on triangular meshes
    //* Create Laplace-Beltrami in a standard matrix
    Eigen::SparseMatrix<double> laplaceBeltrami(vertCount, vertCount);

    {//* Create LaplaceBeltrami Lij
      // We create the Lii and Lij (i!=j) values in different ways : Lii are added to the triplet list at the end
      std::vector<Eigen::Triplet<double>> LtripletList;
      std::vector<double> Lii (vertCount, 0); // for Lii values

      for (Vertex v1 : mesh->vertices()) {
        for (Edge e : v1.adjacentEdges()) {
          auto wij = geometry->edgeCotanWeight(e);

          Lii.at(v1.getIndex()) += wij;

          auto v2 = e.firstVertex().getIndex() == v1.getIndex() ? e.secondVertex() : e.firstVertex();
          LtripletList.push_back(Eigen::Triplet<double>(v1.getIndex(), v2.getIndex(), -wij));

          assert(v1.getIndex() != v2.getIndex());
        }
      }
      
      for (Vertex v : mesh->vertices()) { // add Lii triplets
        LtripletList.push_back(Eigen::Triplet<double>(v.getIndex(), v.getIndex(), Lii.at(v.getIndex())));
      }
      laplaceBeltrami.setFromTriplets(LtripletList.begin(), LtripletList.end());
      laplaceBeltrami.finalize();

      //TODO: check that row and columns sum to 0
      //double sumCol, sumRow;
      //for (Vertex v1: mesh->vertices()) {
      //  sumCol = sumRow = 0;
      //  for (Vertex v2: mesh->vertices()) {
      //    sumCol += laplaceBeltrami((size_t)v1.getIndex(), (size_t)v2.getIndex()); // does not work
      //    sumRow += laplaceBeltrami((size_t)v2.getIndex(), (size_t)v1.getIndex());
      //  }
      //  std::cout << sumCol << " " << sumRow << std::endl;
      //}
    }

    Eigen::SparseMatrix<double> M_inv(vertCount, vertCount);
    {//* Create M in a sparse matrix
      std::vector<Eigen::Triplet<double>> MtripletList;

      for (Vertex v : mesh->vertices()) {
        double val = 0;
        for (auto f : v.adjacentFaces()) {
          val += 1/3*geometry->faceArea(f);
        }
        assert(val != 0);
        MtripletList.push_back(Eigen::Triplet<double>(v.getIndex(), v.getIndex(), 1/val));
      }
      M_inv.setFromTriplets(MtripletList.begin(), MtripletList.end());
      M_inv.finalize();
    }
    
    //* Show the result on some function of q1
    Vector<double> eigen_sin_coords(vertCount); // We need to convert this to an Eigen::Vector
    for (auto i=0; i < vertCount; i++) {
      eigen_sin_coords[i] = sin_coords.at(i);
    }

    Eigen::SparseMatrix<double> DeltaCotan = M_inv * laplaceBeltrami;
    geometrycentral::Vector<double> lb_sin_coords = DeltaCotan * eigen_sin_coords;

    sfm->addVertexScalarQuantity("Laplace-Beltrami of Sinusoid data", lb_sin_coords);
  }


  // Give control to the polyscope gui
  polyscope::show();

  return EXIT_SUCCESS;
}