#include <chrono>
#include <iostream>
#include <string>
#include <utility>
#include <vector>

#include <Eigen/Core>
#include <Eigen/Sparse>
#include <Eigen/SparseLU>

#include <igl/AABB.h>
#include <igl/cotmatrix.h>
#include <igl/massmatrix.h>
#include <igl/per_face_normals.h>
#include <igl/per_vertex_normals.h>
#include <igl/read_triangle_mesh.h>
#include <igl/write_triangle_mesh.h>

// polyscope has a lot warnings
_Pragma("GCC diagnostic push")
_Pragma("GCC diagnostic ignored \"-Wunused-function\"")
_Pragma("GCC diagnostic ignored \"-Wunused-parameter\"")
#include <polyscope/surface_mesh.h>
_Pragma("GCC diagnostic pop")

#define INPUT_COLOR  {  28.0 / 255.0,  99.0 / 255.0, 227.0 / 255.0 }
#define OUTPUT_COLOR { 227.0 / 255.0,  28.0 / 255.0,  49.0 / 255.0 }
#define STYLE_COLOR  { 227.0 / 255.0, 156.0 / 255.0,  28.0 / 255.0 }
#define CONVEX_COLOR {  46.0 / 255.0, 227.0 / 255.0,  28.0 / 255.0 }

template <class Duration = std::chrono::duration<double>,
          class Clock    = std::chrono::steady_clock,
          class Expr>
Duration duration_of(Expr expr)
{
    std::chrono::time_point<Clock> start;
    std::chrono::time_point<Clock> end;

    start = Clock::now();
    expr();
    end = Clock::now();

    return Duration{end - start};
}

struct Surface_mesh
{
    Eigen::MatrixX3d vertices;
    Eigen::MatrixX3i faces;
};

std::pair<double, Eigen::Vector3d> normalize(Eigen::MatrixX3d& vertices)
{
    Eigen::Vector3d bb_min = vertices.colwise().minCoeff();
    Eigen::Vector3d bb_max = vertices.colwise().maxCoeff();

    Eigen::Vector3d center = (bb_min + bb_max) / 2.0;
    vertices.rowwise() -= center.transpose();

    bb_min -= center;
    bb_max -= center;

    double scaling_factor = 1.0 / (bb_max - bb_min).norm();
    vertices.array() *= scaling_factor;

    std::pair<double, Eigen::Vector3d> scaling_translation;
    scaling_translation.first = scaling_factor;
    scaling_translation.second = -center * scaling_factor;

    return scaling_translation;
}

void unormalize(const std::pair<double, Eigen::Vector3d>& scaling_translation,
                Eigen::MatrixX3d& vertices)
{
    vertices.rowwise() -= scaling_translation.second.transpose();
    vertices.array() /= scaling_translation.first;
}

// Not the best but coded in few lines.
class Style_mapper
{
public:

    Style_mapper(Surface_mesh const* style_mesh_ptr) :
        style_mesh_ptr_(style_mesh_ptr)
    {}

    void precompute(bool convex_flag,
                    int nb_of_iterations,
                    double delta)
    {
        vertices_ = style_mesh_ptr_->vertices;

        igl::per_face_normals(vertices_, style_mesh_ptr_->faces, face_normals_);

        // mean curvature flow here

        // ensure that all normals are oriented to the outter side
        // mesh should be convex
        if (vertices_.row(0).dot(face_normals_.row(0)) < 0.0)
        {
            face_normals_.array() *= -1.0;
        }

        if (!convex_flag) apply_mean_curvature_flow(nb_of_iterations, delta);

        aabb_.init(vertices_, style_mesh_ptr_->faces);
    }

    Eigen::Vector3d generate_normal(const Eigen::Vector3d& normal) const
    {
        Eigen::Vector3d origin = Eigen::Vector3d::Zero(3);
        igl::Hit hit;
        aabb_.intersect_ray(vertices_, style_mesh_ptr_->faces, origin, normal, hit);
        return face_normals_.row(hit.id);
    }

    const Eigen::MatrixX3d& vertices() const { return vertices_; }
    const Eigen::MatrixX3i& faces() const { return style_mesh_ptr_->faces; }
    const Eigen::MatrixX3d& face_normals() const { return face_normals_; }

private:

    using AABB = igl::AABB<Eigen::MatrixX3d, 3>;

    void apply_mean_curvature_flow(int nb_of_iterations,
                                   double delta)
    {
        Eigen::SparseMatrix<double> L;
        igl::cotmatrix(vertices_, style_mesh_ptr_->faces, L);

        for (int i = 0; i < nb_of_iterations; ++i)
        {
            Eigen::SparseMatrix<double> M;
            igl::massmatrix(vertices_, style_mesh_ptr_->faces, igl::MASSMATRIX_TYPE_BARYCENTRIC, M);

            auto A = M - delta * L;
            auto b = M * vertices_;

            Eigen::SparseLU<Eigen::SparseMatrix<double>> solver(A);
            vertices_ = solver.solve(b).eval();
        }

        normalize(vertices_);
    }

    Surface_mesh const* style_mesh_ptr_;

    Eigen::MatrixX3d    vertices_;
    Eigen::MatrixX3d    face_normals_;
    AABB                aabb_;

};

Eigen::Matrix3d normal_rotation(const Eigen::Vector3d& u,
                                const Eigen::Vector3d& v)
{
    Eigen::Vector3d w = u.cross(v);
    Eigen::Matrix3d K(3, 3);
    K <<   0.0, -w(2),  w(1),
          w(2),   0.0, -w(0),
         -w(1),  w(0),   0.0;
    return Eigen::Matrix3d::Identity(3, 3) + K + 1.0 / (1.0 + u.dot(v)) * K * K;
}

Eigen::Matrix3d fit_rotation(const Eigen::Matrix3d& matrix)
{
    Eigen::JacobiSVD<Eigen::Matrix3d> SVD(matrix, Eigen::ComputeFullU | Eigen::ComputeFullV);
    Eigen::Matrix3d R = SVD.matrixV() * SVD.matrixU().transpose();
    if (R.determinant() < 0) {
        Eigen::Matrix3d V = SVD.matrixV();
        V.col(V.cols() - 1).array() *= -1.0;
        R = V * SVD.matrixU().transpose();
    }
    return R;
}

enum Method
{
    NAIVE_ROTATIONS,
    STATIC_NORMALS,
    DYNAMIC_NORMALS
};

enum Regularization
{
    ARAP,
    ACAP
};

class Surface_mesh_deformation
{
public:

    Surface_mesh_deformation(Surface_mesh const* input_mesh_ptr,
                             Style_mapper const* style_mapper_ptr) :
        input_mesh_ptr_(input_mesh_ptr),
        style_mapper_ptr_(style_mapper_ptr)
    {
        const int nb_of_vertices = input_mesh_ptr_->vertices.rows();
        input_areas_.resize(nb_of_vertices);
        input_normals_.resize(nb_of_vertices, 3);
        source_normals_.resize(nb_of_vertices, 3);
        target_normals_.resize(nb_of_vertices, 3);
        tmp_scalings_.resize(nb_of_vertices);
        scalings_.resize(nb_of_vertices);
        rotations_.resize(nb_of_vertices);
        A_.resize(nb_of_vertices, nb_of_vertices);
        x_.resize(nb_of_vertices, 3);
        b_.resize(nb_of_vertices, 3);
    }

    void precompute()
    {
        Eigen::SparseMatrix<double> M;
        igl::massmatrix(input_mesh_ptr_->vertices, input_mesh_ptr_->faces, igl::MASSMATRIX_TYPE_VORONOI, M);
        input_areas_ = M.diagonal();

        igl::per_vertex_normals(input_mesh_ptr_->vertices, input_mesh_ptr_->faces, input_normals_);

        input_bary_ = input_mesh_ptr_->vertices.colwise().mean();

        igl::cotmatrix(input_mesh_ptr_->vertices, input_mesh_ptr_->faces, A_);
        A_ *= -1.0;

        solver_.compute(A_);
    }

    void solve(Method method,
               Regularization regularization,
               int nb_of_iterations,
               double lambda)
    {
        switch (method)
        {
        case NAIVE_ROTATIONS: solve_naive_rotations();                         break;
        case STATIC_NORMALS:  solve_static_normals(regularization, nb_of_iterations, lambda);  break;
        case DYNAMIC_NORMALS: solve_dynamic_normals(regularization, nb_of_iterations, lambda); break;
        }
    }

    const Eigen::MatrixX3d& output_vertices() const { return x_; }
    const Eigen::MatrixX3i& output_faces() const { return input_mesh_ptr_->faces; }

    const Eigen::MatrixX3d& input_normals() const { return input_normals_; }
    const Eigen::MatrixX3d& source_normals() const { return source_normals_; }
    const Eigen::MatrixX3d& target_normals() const { return target_normals_; }

    const Eigen::VectorXd& scalings() const { return scalings_; }

    Eigen::MatrixX3d rotated_normals() const
    {
        const int nb_of_vertices = input_mesh_ptr_->vertices.rows();
        Eigen::MatrixX3d normals(nb_of_vertices, 3);
        #pragma omp parallel for
        for (int vi = 0; vi < nb_of_vertices; ++vi)
        {
            normals.row(vi) = (rotations_[vi] * input_normals_.row(vi).transpose()).transpose();
        }
        return normals;
    }

    Eigen::MatrixX3d output_normals() const
    {
        Eigen::MatrixX3d normals;
        igl::per_vertex_normals(x_, input_mesh_ptr_->faces, normals);
        return normals;
    }

    Eigen::VectorXd energy(double lambda) const
    {
        const int nb_of_vertices = input_mesh_ptr_->vertices.rows();

        Eigen::VectorXd energy(nb_of_vertices);
        energy.array() = 0.0;

        for (int k = 0; k < A_.outerSize(); ++k)
        {
            for (Eigen::SparseMatrix<double>::InnerIterator it(A_, k); it; ++it)
            {
                const int vi = it.row();
                const int vj = it.col();

                if (vi == vj) continue;

                const double wij = -it.value();
                const auto&  si  = input_mesh_ptr_->vertices.row(vi).transpose();
                const auto&  sj  = input_mesh_ptr_->vertices.row(vj).transpose();
                const auto&  xi  = x_.row(vi).transpose();
                const auto&  xj  = x_.row(vj).transpose();
                const auto&  Ri  = rotations_[vi];
                const double ki  = scalings_(vi);

                energy(vi) += wij * ((xi - xj) - ki * Ri * (si - sj)).squaredNorm();
            }
        }

        #pragma omp parallel for
        for (int vi = 0; vi < nb_of_vertices; ++vi)
        {
            const auto& ni = input_normals_.row(vi).transpose();
            const auto& ti = target_normals_.row(vi).transpose();
            const auto& Ri  = rotations_[vi];

            energy(vi) += lambda * input_areas_(vi) * (Ri * ni - ti).squaredNorm();
        }

        return energy;
    }

private:

    void initialize()
    {
        x_ = input_mesh_ptr_->vertices;
        source_normals_ = input_normals_;
        scalings_.array() = 1.0;
    }

    void compute_source_normals()
    {
        igl::per_vertex_normals(x_, input_mesh_ptr_->faces, source_normals_);
    }

    void compute_target_normals()
    {
        const int nb_of_vertices = input_mesh_ptr_->vertices.rows();
        #pragma omp parallel for
        for (int vi = 0; vi < nb_of_vertices; ++vi)
        {
            target_normals_.row(vi) = style_mapper_ptr_->generate_normal(source_normals_.row(vi));
        }
    }

    void optimize_rotations_from_normals()
    {
        const int nb_of_vertices = input_mesh_ptr_->vertices.rows();
        #pragma omp parallel for
        for (int vi = 0; vi < nb_of_vertices; ++vi)
        {
            rotations_[vi] = normal_rotation(source_normals_.row(vi), target_normals_.row(vi));
        }
    }

    void optimize_rotations_from_energy(double lambda)
    {
        const int nb_of_vertices = input_mesh_ptr_->vertices.rows();

        #pragma omp parallel for
        for (int vi = 0; vi < nb_of_vertices; ++vi)
        {
            rotations_[vi].array() = 0.0;
        }

        //#pragma omp parallel for collapse(2)
        for (int k = 0; k < A_.outerSize(); ++k)
        {
            for (Eigen::SparseMatrix<double>::InnerIterator it(A_, k); it; ++it)
            {
                const int vi = it.row();
                const int vj = it.col();

                if (vi == vj) continue;

                const double wij = -it.value();
                const auto&  si  = input_mesh_ptr_->vertices.row(vi).transpose();
                const auto&  sj  = input_mesh_ptr_->vertices.row(vj).transpose();
                const auto&  xi  = x_.row(vi).transpose();
                const auto&  xj  = x_.row(vj).transpose();

                const Eigen::Matrix3d R = wij * (si - sj) * (xi - xj).transpose();
                //#pragma omp atomic
                rotations_[vi] += R;
            }
        }

        #pragma omp parallel for
        for (int vi = 0; vi < nb_of_vertices; ++vi)
        {
            const auto& ni = input_normals_.row(vi).transpose();
            const auto& ti = target_normals_.row(vi).transpose();

            rotations_[vi] += lambda * input_areas_(vi) * ni * ti.transpose();

            rotations_[vi] = fit_rotation(rotations_[vi]);
        }
    }

    void optimize_scalings(double lambda)
    {
        const int nb_of_vertices = input_mesh_ptr_->vertices.rows();

        scalings_.array() = 0.0;
        tmp_scalings_.array() = 0.0;

        for (int k = 0; k < A_.outerSize(); ++k)
        {
            for (Eigen::SparseMatrix<double>::InnerIterator it(A_, k); it; ++it)
            {
                const int vi = it.row();
                const int vj = it.col();

                if (vi == vj) continue;

                const double wij = -it.value();
                const auto&  si  = input_mesh_ptr_->vertices.row(vi).transpose();
                const auto&  sj  = input_mesh_ptr_->vertices.row(vj).transpose();
                const auto&  xi  = x_.row(vi).transpose();
                const auto&  xj  = x_.row(vj).transpose();
                const auto&  Ri  = rotations_[vi];

                scalings_(vi) += wij * (xi - xj).transpose() * Ri * (si - sj);
                tmp_scalings_(vi) += wij * (si - sj).transpose() * (si - sj);
            }
        }

        #pragma omp parallel for
        for (int vi = 0; vi < nb_of_vertices; ++vi)
        {
            const auto& ni = input_normals_.row(vi).transpose();
            const auto& ti = target_normals_.row(vi).transpose();

            const double nTt = lambda * input_areas_(vi) * ni.transpose().dot(ti);

            scalings_(vi) = (scalings_(vi) + nTt) / (tmp_scalings_(vi) + nTt);
        }
    }

    void optimize_positions()
    {
        // Spokes only instead spokes and rims because igl dont make it easy
        b_.array() = 0.0;
        for (int k = 0; k < A_.outerSize(); ++k)
        {
            for (Eigen::SparseMatrix<double>::InnerIterator it(A_, k); it; ++it)
            {
                const int vi = it.row();
                const int vj = it.col();

                if (vi == vj) continue;

                const double wij = -it.value();
                const auto&  si  = input_mesh_ptr_->vertices.row(vi).transpose();
                const auto&  sj  = input_mesh_ptr_->vertices.row(vj).transpose();
                const auto&  Ri  = rotations_[vi];
                const auto&  Rj  = rotations_[vj];
                const double ki  = scalings_(vi);
                const double kj  = scalings_(vj);

                 b_.row(vi) += (0.5 * wij * (ki * Ri + kj * Rj) * (si - sj)).transpose();
            }
        }

        x_ = solver_.solve(b_);

        // solution is up to a translation
        x_.rowwise() += (input_bary_.transpose() - x_.colwise().mean()).eval();
    }

    void solve_naive_rotations()
    {
        initialize();
        compute_target_normals();
        optimize_rotations_from_normals();
        optimize_positions();
    }

    void solve_static_normals(Regularization regularization,
                              int nb_of_iterations,
                              double lambda)
    {
        initialize();
        compute_target_normals();
        for (int i = 0; i < nb_of_iterations; ++i)
        {
            optimize_rotations_from_energy(lambda);
            if (regularization == ACAP) optimize_scalings(lambda);
            optimize_positions();
        }
    }

    void solve_dynamic_normals(Regularization regularization,
                               int nb_of_iterations,
                               double lambda)
    {
        initialize();
        for (int i = 0; i < nb_of_iterations; ++i)
        {
            compute_source_normals();
            compute_target_normals();
            optimize_rotations_from_energy(lambda);
            if (regularization == ACAP) optimize_scalings(lambda);
            optimize_positions();
        }
    }

    Surface_mesh const*                          input_mesh_ptr_;
    Style_mapper const*                          style_mapper_ptr_;

    Eigen::VectorXd                              input_areas_;

    Eigen::Vector3d                              input_bary_;

    Eigen::MatrixX3d                             input_normals_;
    Eigen::MatrixX3d                             source_normals_;
    Eigen::MatrixX3d                             target_normals_;

    Eigen::VectorXd                              tmp_scalings_;
    Eigen::VectorXd                              scalings_;
    std::vector<Eigen::Matrix3d>                 rotations_;

    Eigen::SparseMatrix<double>                  A_;
    Eigen::MatrixX3d                             x_;
    Eigen::MatrixX3d                             b_;

    Eigen::SparseLU<Eigen::SparseMatrix<double>> solver_;

};

struct Args
{
    std::string    input_filename;
    std::string    style_filename;

    bool           convex_flag;

    Method         method;
    Regularization regularization;
    int            ndssa_nb_of_iterations;
    double         lambda;

    int            mcf_nb_of_iterations;
    double         delta;

    std::string    output_filename;

    bool           gui_flag;
};

bool parse_args(int argc,
                char** argv,
                Args& args)
{
    args.convex_flag            = true;

    args.method                 = STATIC_NORMALS;
    args.regularization         = ARAP;
    args.ndssa_nb_of_iterations = 10;
    args.lambda                 = 1.0;

    args.mcf_nb_of_iterations   = 10;
    args.delta                  = 0.001;

    args.output_filename        = "output.obj";

    args.gui_flag               = false;

    bool mesh_seen = false;
    bool style_seen = false;

    for (int i = 1; i < argc; ++i)
    {
        std::string argi = argv[i];
             if ( argi == "--help")                                       { return false;                                          }
        else if ( argi == "--convex"  || argi == "-c")                    { args.convex_flag            = false;                   }
        else if ( argi == "--naive"   || argi == "-n")                    { args.method                 = NAIVE_ROTATIONS;         }
        else if ( argi == "--static"  || argi == "-s")                    { args.method                 = STATIC_NORMALS;          }
        else if ( argi == "--dynamic" || argi == "-d")                    { args.method                 = DYNAMIC_NORMALS;         }
        else if ( argi == "--arap"    || argi == "-R")                    { args.regularization         = ARAP;                    }
        else if ( argi == "--acap"    || argi == "-C")                    { args.regularization         = ACAP;                    }
        else if ((argi == "--iter"    || argi == "-i") && (i + 1) < argc) { args.ndssa_nb_of_iterations = std::stoi(argv[++i]);    }
        else if ((argi == "--lambda"  || argi == "-l") && (i + 1) < argc) { args.lambda                 = std::stod(argv[++i]);    }
        else if ((argi == "--ITER"    || argi == "-I") && (i + 1) < argc) { args.mcf_nb_of_iterations   = std::stoi(argv[++i]);    }
        else if ((argi == "--delta "  || argi == "-D") && (i + 1) < argc) { args.delta                  = std::stod(argv[++i]);    }
        else if ((argi == "--output"  || argi == "-o") && (i + 1) < argc) { args.output_filename        = argv[++i];               }
        else if ( argi == "--gui")                                        { args.gui_flag               = true;                    }
        else if (!mesh_seen)                                              { args.input_filename         = argi; mesh_seen = true;  }
        else if (!style_seen)                                             { args.style_filename         = argi; style_seen = true; }
        else                                                              { return false;                                          }
    }

    if (!mesh_seen)  return false;
    if (!style_seen) return false;

    return true;
}

char const* to_string(Method method)
{
    switch (method)
    {
    case NAIVE_ROTATIONS: return "NAIVE_ROTATIONS";
    case STATIC_NORMALS:  return "STATIC_NORMALS";
    case DYNAMIC_NORMALS: return "DYNAMIC_NORMALS";
    }
    return "";
}

char const* to_string(Regularization regularization)
{
    switch (regularization)
    {
    case ARAP: return "ARAP";
    case ACAP:  return "ACAP";
    }
    return "";
}

void print_args(const Args& args)
{
    std::cout << "# args                 :"                                    << std::endl
              << "input_filename         : " << args.input_filename            << std::endl
              << "style_filename         : " << args.style_filename            << std::endl
              << "convex_flag            : " << args.convex_flag               << std::endl
              << "method                 : " << to_string(args.method)         << std::endl
              << "regularization         : " << to_string(args.regularization) << std::endl
              << "ndssa_nb_of_iterations : " << args.ndssa_nb_of_iterations    << std::endl
              << "lambda                 : " << args.lambda                    << std::endl
              << "mcf_nb_of_iterations   : " << args.mcf_nb_of_iterations      << std::endl
              << "delta                  : " << args.delta                     << std::endl
              << "output_filename        : " << args.output_filename           << std::endl
              << "gui_flag               : " << args.gui_flag                  << std::endl
              << "------------------------------------------------"            << std::endl;
}

void print_help()
{
    std::cout << "Usage : ndssa [option] INPUT STYLE\n"
                 "\n"
                 "Options:\n"
                 "\n"
                 "  --help                print this help\n"
                 "  --convex  [ -c ]      specify STYLE is not convex\n"
                 "  --naive   [ -n ]      use naive rotation method\n"
                 "  --static  [ -t ]      use static normal method (default method)\n"
                 "  --dynamic [ -d ]      use dynamic normal method\n"
                 "  --arap    [ -R ]      use arap regularization (default regularization)\n"
                 "  --acap    [ -C ]      use acap regularization\n"
                 "  --iter    [ -i ] ARG  set number of iterations (default is 10)\n"
                 "  --lambda  [ -l ] ARG  set lambda (default is 1.0)\n"
                 "  --ITER    [ -I ] ARG  set number of mcf iterations (default is 10)\n"
                 "  --delta   [ -D ] ARG  set mcf delta (default is 0.001)\n"
                 "  --output  [ -o ] ARG  set output mesh filename (default is output.obj)\n"
                 "  --gui                 run in gui mode\n" << std::endl;
}

int main(int argc, char** argv)
{
    Args args;

    if (!parse_args(argc, argv, args))
    {
        print_help();
        return EXIT_FAILURE;
    }

    print_args(args);

    Surface_mesh input_mesh;
    Surface_mesh style_mesh;

    if (!igl::read_triangle_mesh(args.input_filename, input_mesh.vertices, input_mesh.faces))
    {
        std::cerr << "could not read '" << args.input_filename << "'" << std::endl;
        return EXIT_FAILURE;
    }
    if (!igl::read_triangle_mesh(args.style_filename, style_mesh.vertices, style_mesh.faces))
    {
        std::cerr << "could not read '" << args.style_filename << "'" << std::endl;
        return EXIT_FAILURE;
    }

    const auto scaling_translation = normalize(input_mesh.vertices);
    normalize(style_mesh.vertices);

    Style_mapper style_mapper(&style_mesh);
    Surface_mesh_deformation deformation(&input_mesh, &style_mapper);

    auto t0 = duration_of([&]() { style_mapper.precompute(args.convex_flag, args.mcf_nb_of_iterations, args.delta); });
    auto t1 = duration_of([&]() { deformation.precompute(); });

    std::cout << "# precompute           :"                              << std::endl;
    std::cout << "style_mapper           : " << t0.count() << " seconds" << std::endl;
    std::cout << "deformation            : " << t1.count() << " seconds" << std::endl;
    std::cout << "------------------------------------------------"      << std::endl;

    if (args.gui_flag)
    {
        polyscope::options::programName     = "ndssa";
        polyscope::options::verbosity       = 0;
        polyscope::options::usePrefsFile    = false;
        polyscope::options::groundPlaneMode = polyscope::GroundPlaneMode::ShadowOnly;

        polyscope::init();

        ImGui::GetIO().IniFilename = nullptr;

        auto input_ptr = polyscope::registerSurfaceMesh("input", input_mesh.vertices, input_mesh.faces);
        input_ptr->setEnabled(false);
        input_ptr->setSurfaceColor(INPUT_COLOR);
        input_ptr->setEdgeColor({ 1.0, 1.0, 1.0 });
        input_ptr->setEdgeWidth(1.0);

        auto style_ptr = polyscope::registerSurfaceMesh("style", style_mesh.vertices, style_mesh.faces);
        style_ptr->setEnabled(false);
        style_ptr->setSurfaceColor(STYLE_COLOR);
        style_ptr->setEdgeColor({ 1.0, 1.0, 1.0 });
        style_ptr->setEdgeWidth(1.0);
        style_ptr->addFaceVectorQuantity("normals", style_mapper.face_normals());

        auto convex_ptr = polyscope::registerSurfaceMesh("style.convex", style_mapper.vertices(), style_mapper.faces());
        convex_ptr->setEnabled(false);
        convex_ptr->setSurfaceColor(CONVEX_COLOR);
        convex_ptr->setEdgeColor({ 1.0, 1.0, 1.0 });
        convex_ptr->setEdgeWidth(1.0);
        convex_ptr->addFaceVectorQuantity("normals", style_mapper.face_normals());

        auto output_ptr = polyscope::registerSurfaceMesh("output", input_mesh.vertices, input_mesh.faces);
        output_ptr->setSurfaceColor(OUTPUT_COLOR);
        output_ptr->setEdgeColor({ 1.0, 1.0, 1.0 });
        output_ptr->setEdgeWidth(1.0);
        output_ptr->addVertexVectorQuantity("normals.input", deformation.input_normals());
        output_ptr->addVertexVectorQuantity("normals.rotated", Eigen::MatrixX3d::Zero(input_mesh.vertices.rows(), 3));
        output_ptr->addVertexVectorQuantity("normals.source", Eigen::MatrixX3d::Zero(input_mesh.vertices.rows(), 3));
        output_ptr->addVertexVectorQuantity("normals.target", Eigen::MatrixX3d::Zero(input_mesh.vertices.rows(), 3));
        output_ptr->addVertexVectorQuantity("normals.output", Eigen::MatrixX3d::Zero(input_mesh.vertices.rows(), 3));
        output_ptr->addVertexScalarQuantity("energy", Eigen::VectorXd::Zero(input_mesh.vertices.rows()));
        output_ptr->addVertexScalarQuantity("scaling", Eigen::VectorXd::Ones(input_mesh.vertices.rows()));
        
        char const* method_choices[] =
        {
            "naive rotations",
            "static normals",
            "dynamic normals"
        };
        char const* regularization_choices[] =
        {
            "arap",
            "acap"
        };

        polyscope::state::userCallback = [&]()
        {
            ImGui::Text("Deformation");
            ImGui::Combo("method", (int*) &args.method, method_choices, IM_ARRAYSIZE(method_choices));
            ImGui::Combo("regularization", (int*) &args.regularization, regularization_choices, IM_ARRAYSIZE(regularization_choices));
            ImGui::InputInt("nb of iterations##ndssa", &args.ndssa_nb_of_iterations);
            ImGui::InputDouble("lambda", &args.lambda);

            if (ImGui::Button("compute##ndssa"))
            {
                auto t2 = duration_of([&]() { deformation.solve(args.method, args.regularization, args.ndssa_nb_of_iterations, args.lambda); });

                const auto& energy = deformation.energy(args.lambda);

                std::cout << "# deformation          :"                                 << std::endl;
                std::cout << "method                 : " << to_string(args.method)      << std::endl;
                std::cout << "nb_of_iterations       : " << args.ndssa_nb_of_iterations << std::endl;
                std::cout << "lambda                 : " << args.lambda                 << std::endl;
                std::cout << "energy                 : " << energy.sum()                << std::endl;
                std::cout << "duration               : " << t2.count() << " seconds"    << std::endl;
                std::cout << "------------------------------------------------"         << std::endl;

                output_ptr->updateVertexPositions(deformation.output_vertices());
                output_ptr->addVertexVectorQuantity("normals.rotated", deformation.rotated_normals());
                output_ptr->addVertexVectorQuantity("normals.source", deformation.source_normals());
                output_ptr->addVertexVectorQuantity("normals.target", deformation.target_normals());
                output_ptr->addVertexVectorQuantity("normals.output", deformation.output_normals());
                output_ptr->addVertexScalarQuantity("energy", energy);
                output_ptr->addVertexScalarQuantity("scaling", deformation.scalings());
            }

            ImGui::Dummy(ImVec2(0.0, 20.0));

            ImGui::Text("Mean Curvature Flow");
            ImGui::InputInt("nb of iterations##mcf", &args.mcf_nb_of_iterations);
            ImGui::InputDouble("delta", &args.delta);

            if (ImGui::Button("compute##mcf"))
            {
                t0 = duration_of([&]() { style_mapper.precompute(false, args.mcf_nb_of_iterations, args.delta); });

                std::cout << "# precompute           :"                              << std::endl;
                std::cout << "style_mapper           : " << t0.count() << " seconds" << std::endl;
                std::cout << "------------------------------------------------"      << std::endl;

                convex_ptr->updateVertexPositions(style_mapper.vertices());
            }
        };

        polyscope::show();
    }
    else
    {
        auto t2 = duration_of([&]() { deformation.solve(args.method, args.regularization, args.ndssa_nb_of_iterations, args.lambda); });

        const auto& energy = deformation.energy(args.lambda);

        std::cout << "# deformation          :"                                 << std::endl;
        std::cout << "method                 : " << to_string(args.method)      << std::endl;
        std::cout << "nb_of_iterations       : " << args.ndssa_nb_of_iterations << std::endl;
        std::cout << "lambda                 : " << args.lambda                 << std::endl;
        std::cout << "energy                 : " << energy.sum()                << std::endl;
        std::cout << "duration               : " << t2.count() << " seconds"    << std::endl;
        std::cout << "------------------------------------------------"         << std::endl;
    }

    Eigen::MatrixX3d output_vertices = deformation.output_vertices();
    unormalize(scaling_translation, output_vertices);

    if (!igl::write_triangle_mesh(args.output_filename, output_vertices, input_mesh.faces))
    {
        std::cerr << "could not write '" << args.output_filename << "'" << std::endl;
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}
