//== INCLUDES =================================================================

#include "Smoother.h"

#include <cfloat>

#include <Eigen/Sparse>

#include <graphene/surface_mesh/algorithms/surface_mesh_tools/diffgeo.h>

//== NAMESPACES ===============================================================

namespace graphene {
namespace surface_mesh {


//== IMPLEMENTATION ==========================================================


Smoother::
Smoother(Surface_mesh& mesh)
: mesh_(mesh), initialized_(false)
{
    vweight_ = mesh_.vertex_property<double>("smoother:vweight");
    eweight_ = mesh_.edge_property<double>("smoother:eweight");
}


//-----------------------------------------------------------------------------


Smoother::
~Smoother()
{
    mesh_.remove_vertex_property(vweight_);
    mesh_.remove_edge_property(eweight_);
}


//-----------------------------------------------------------------------------


void
Smoother::
initialize(bool use_uniform_laplace)
{
    auto points  = mesh_.vertex_property<Point>("v:point");
    auto vweight = mesh_.vertex_property<double>("smoother:vweight");
    auto eweight = mesh_.edge_property<double>("smoother:eweight");

    // compute weights
    for (auto v : mesh_.vertices())
    {
        if (use_uniform_laplace)
        {
            vweight[v] = 1.0 / mesh_.valence(v);
        }
        else
        {
            vweight[v] = 0.5 / voronoi_area(mesh_, v);
        }
    }

    for (auto e : mesh_.edges())
    {
        if (use_uniform_laplace)
        {
            eweight[e] = 1.0;
        }
        else
        {
            eweight[e] = std::max(0.0, cotan_weight(mesh_, e));
        }
    }

    initialized_ = true;
}


//-----------------------------------------------------------------------------


void
Smoother::
smooth(unsigned int iterations, bool use_uniform_laplace)
{
    // make sure the cotan weights have been computed
    if (!initialized_) initialize(use_uniform_laplace);

    // get properties
    Surface_mesh::Vertex_property<Point>   points  = mesh_.get_vertex_property<Point>("v:point");
    Surface_mesh::Vertex_property<double>  vweight = mesh_.get_vertex_property<double>("smoother:vweight");
    Surface_mesh::Edge_property<double>    eweight = mesh_.get_edge_property<double>("smoother:eweight");

    Surface_mesh::Vertex_property<bool>  vselected = mesh_.get_vertex_property<bool>("v:selected");

    // add properties
    Surface_mesh::Vertex_property<Point>   laplace = mesh_.add_vertex_property<Point>("smoother:laplace");

    // check whether some vertices are selected
    bool no_selection = true;
    if (vselected)
    {
        for (auto v : mesh_.vertices())
        {
            if (vselected[v])
            {
                no_selection = false;
                break;
            }
        }
    }


    // here come the smoothing iterations...
    for (unsigned int i=0; i<iterations; ++i)
    {

        // 1st pass: compute laplace
        for (auto v : mesh_.vertices())
        {
            if (!mesh_.is_isolated(v))
            {
                if (no_selection || vselected[v])
                {
                    Point l(0);
                    for (auto he : mesh_.halfedges(v))
                    {
                        auto e  = mesh_.edge(he);
                        auto vv = mesh_.to_vertex(he);
                        l += eweight[e] * (points[vv] - points[v]);
                    }

                    l *= vweight[v];

                    laplace[v] = l;
                }
            }
        }


        // 2nd pass: update vertex positions
        for (auto v : mesh_.vertices())
        {
            if (!mesh_.is_isolated(v))
            {
                if (no_selection || vselected[v])
                {
                    points[v] += double(0.5) * laplace[v];
                }
            }
        }
    }


    // remove properties
    mesh_.remove_vertex_property(laplace);
}

//-----------------------------------------------------------------------------

void Smoother::smooth_implicit(double timestep, bool use_uniform_laplace)
{
    auto points = mesh_.get_vertex_property<Point>("v:point");

    // scale mesh into unit sphere
    Point  bbmin( FLT_MAX,  FLT_MAX,  FLT_MAX);
    Point  bbmax(-FLT_MAX, -FLT_MAX, -FLT_MAX);
    for (auto v : mesh_.vertices())
    {
        bbmin.minimize(points[v]);
        bbmax.maximize(points[v]);
    }

    Point  center = 0.5 * (bbmin + bbmax);
    double radius = 0.5 * distance(bbmax, bbmin);
    for (auto v : mesh_.vertices())
    {
        points[v] -= center;
        points[v] /= radius;
    }

    if (!initialized_) initialize(use_uniform_laplace);

    // properties
    auto vweight   = mesh_.get_vertex_property<double>("smoother:vweight");
    auto eweight   = mesh_.get_edge_property<double>("smoother:eweight");
    auto idx       = mesh_.add_vertex_property<int>("v:idx", -1);
    auto vselected = mesh_.get_vertex_property<bool>("v:selected");

    // collect free (non-boundary) vertices in array free_vertices[]
    // assign indices such that idx[ free_vertices[j] ] == j
    unsigned j = 0;
    std::vector<Surface_mesh::Vertex> free_vertices;
    free_vertices.reserve(mesh_.n_vertices());
    for (auto v : mesh_.vertices())
    {
        if (!mesh_.is_boundary(v) && vselected[v])
        {
            idx[v] = j++;
            free_vertices.push_back(v);
        }
    }
    const unsigned int n = free_vertices.size();

    // A*X = B
    Eigen::SparseMatrix<double> A(n, n);
    Eigen::MatrixXd B(n, 3);

    // nonzero elements of A as triplets: (row, column, value)
    std::vector<Eigen::Triplet<double>> triplets;

    // setup matrix A and rhs B
    for (unsigned int i = 0; i < n; ++i)
    {
        auto v = free_vertices[i];

        // rhs row
        B(i, 0) = points[v][0] / vweight[v];
        B(i, 1) = points[v][1] / vweight[v];
        B(i, 2) = points[v][2] / vweight[v];

        // lhs row
        auto ww = 0.0;
        for (auto h : mesh_.halfedges(v))
        {
            auto vv = mesh_.to_vertex(h);
            auto e = mesh_.edge(h);
            ww += eweight[e];

            // fixed boundary vertex -> right hand side
            if (mesh_.is_boundary(vv) || !vselected[vv])
            {
                B(i, 0) -= -timestep * eweight[e] * points[vv][0];
                B(i, 1) -= -timestep * eweight[e] * points[vv][1];
                B(i, 2) -= -timestep * eweight[e] * points[vv][2];
            }
            // free interior vertex -> matrix
            else
            {
                triplets.emplace_back(i, idx[vv], -timestep * eweight[e]);
            }
        }

        // center vertex -> matrix
        triplets.emplace_back(i, i, 1.0 / vweight[v] + timestep * ww);
    }

    // build sparse matrix from triplets
    A.setFromTriplets(triplets.begin(), triplets.end());

    // solve A*X = B
    Eigen::SimplicialLDLT<Eigen::SparseMatrix<double>> solver(A);
    Eigen::MatrixXd X = solver.solve(B);
    if (solver.info() != Eigen::Success)
    {
        // clean-up
        mesh_.remove_vertex_property(idx);
        printf("[ERROR] SurfaceSmoothing: Failed to solve linear system.");
    }
    else
    {
        // copy solution
        for (unsigned int i = 0; i < n; ++i)
        {
            Eigen::Vector3d p = X.row(i);
            points[free_vertices[i]] = Point(p[0], p[1], p[2]);
        }
    }


    // undo unit sphere scaling
    for (auto v : mesh_.vertices())
    {
        points[v] *= radius;
        points[v] += center;
    }


    // if (rescale)
    // {
    //    // restore original surface area
    //    double area_after = surface_area(mesh_);
    //    double scale = sqrt(area_before / area_after);
    //    for (auto v : mesh_.vertices())
    //        mesh_.position(v) *= scale;

    //    // restore original center
    //    Point center_after = centroid(mesh_);
    //    Point trans = center_before - center_after;
    //    for (auto v : mesh_.vertices())
    //        mesh_.position(v) += trans;
    // }

    // clean-up
    mesh_.remove_vertex_property(idx);
}

//-----------------------------------------------------------------------------

void smooth(Surface_mesh& mesh, unsigned int iterations)
{
    Smoother s(mesh);
    s.smooth(iterations);
}


//=============================================================================
} // namespace surface_mesh
} // namespace graphene
//=============================================================================
