/*
 *  Copyright (c) 2000-2022 Inria
 *  All rights reserved.
 *
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions are met:
 *
 *  * Redistributions of source code must retain the above copyright notice,
 *  this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright notice,
 *  this list of conditions and the following disclaimer in the documentation
 *  and/or other materials provided with the distribution.
 *  * Neither the name of the ALICE Project-Team nor the names of its
 *  contributors may be used to endorse or promote products derived from this
 *  software without specific prior written permission.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 *  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 *  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 *  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 *  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 *  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 *  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 *  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 *  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 *  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *  POSSIBILITY OF SUCH DAMAGE.
 *
 *  Contact: Bruno Levy
 *
 *     https://www.inria.fr/fr/bruno-levy
 *
 *     Inria,
 *     Domaine de Voluceau,
 *     78150 Le Chesnay - Rocquencourt
 *     FRANCE
 *
 */

#include <geogram/parameterization/mesh_LSCM.h>
#include <geogram/mesh/mesh.h>
#include <geogram/mesh/mesh_geometry.h>
#include <geogram/bibliography/bibliography.h>
#include <geogram/NL/nl.h>
#include <geogram/NL/nl_ext.h>
#include <algorithm>

namespace {
    using namespace GEO;


    /**
     * \brief Computes Least Squares Conformal Maps in least squares or
     *  spectral mode.
     * \details The method is described in the following references:
     *  - Least Squares Conformal Maps, Levy, Petitjean, Ray, Maillot, ACM
     *   SIGGRAPH, 2002
     *  - Spectral Conformal Parameterization, Mullen, Tong, Alliez, Desbrun,
     *   Computer Graphics Forum (SGP conf. proc.), 2008
     */
    class LSCM {
    public:

        /**
         * \brief LSCM constructor
         * \param[in] M a reference to the mesh. It needs to correspond to a
         *   topological disk (open surface with one border and no handle).
         * \param[in] tex_coord a vector property of dimension 2 where to
         *   store the texture coordinates.
         */
        LSCM(Mesh& M, Attribute<double>& tex_coord, Attribute<double>& angle) :
            mesh_(M), tex_coord_(tex_coord), angle_(angle), eigen_(0) {
            geo_assert(tex_coord.dimension() == 2);
            locked_1_ = index_t(-1);
            locked_2_ = index_t(-1);
            verbose_ = false;
        }

        /**
         * \brief Enables or disables messages.
         * \param[in] x if true, messages are displayed on the console
         *  with statistics. Default is non-verbose.
         */
        void set_verbose(bool x) {
            verbose_ = x;
        }

        /**
         * \brief Sets whether spectral mode is used.
         * \details In default mode, the trivial solution (all vertices to zero)
         *  is avoided by locking two vertices (that are as "extremal"
         *  as possible). In spectral mode, the trivial solution is avoided by
         *  finding the first minimizer that is orthogonal to it (more elegant,
         *  but more costly).
         */
        void set_spectral(bool x) {
            spectral_ = x;
        }

        /**
         * \brief Computes the least squares conformal map and stores it in
         *  the texture coordinates of the mesh.
         * \details Outline of the algorithm (steps 1,2,3 are not used
         *   in spectral mode):
         *   - 1) Find an initial solution by projecting on a plane
         *   - 2) Lock two vertices of the mesh
         *   - 3) Copy the initial u,v coordinates to OpenNL
         *   - 4) Construct the LSCM equation with OpenNL
         *   - 5) Solve the equation with OpenNL
         *   - 6) Copy OpenNL solution to the u,v coordinates
         */

        void apply() {

            geo_cite("DBLP:journals/tog/LevyPRM02");
            if(spectral_) {
                geo_cite("DBLP:journals/cgf/MullenTAD08");
            }

            const int nb_eigens = 10;
            nlNewContext();
            NLuint nb_vertices = NLuint(mesh_.vertices.nb());

            if(spectral_) {
                if(nlInitExtension("ARPACK")) {
                    if(verbose_) {
                        Logger::out("LSCM") << "ARPACK extension initialized"
                                            << std::endl;
                    }
                    nlEigenSolverParameteri(NL_EIGEN_SOLVER, NL_ARPACK_EXT);
                    nlEigenSolverParameteri(NL_NB_EIGENS, nb_eigens);
                    if(verbose_) {
                        nlEnable(NL_VERBOSE);
                    }
                } else {
                    if(verbose_) {
                        Logger::out("LSCM")
                            << "Could not initialize ARPACK extension"
                            << std::endl;
                        Logger::out("LSCM")
                            << "Falling back to least squares mode"
                            << std::endl;
                    }
                    spectral_ = false;
                }
            } else {
#ifndef GEO_OS_ANDROID
                if(
                    nb_vertices <= 200000 && (
                        nlExtensionIsInitialized("SUPERLU") ||
                        nlInitExtension("SUPERLU")
                    )
                ) {
                    if(verbose_) {
                        Logger::out("LSCM") << "using SUPERLU"
                                            << std::endl;
                    }
                    nlSolverParameteri(NL_SOLVER, NL_PERM_SUPERLU_EXT);
                } else if(
                    nlExtensionIsInitialized("AMGCL") ||
                    nlInitExtension("AMGCL")
                ) {
                    if(verbose_) {
                        nlEnable(NL_VERBOSE);
                        Logger::out("LSCM") << "using AMGCL"
                                            << std::endl;
                    }
                    nlSolverParameteri(NL_SOLVER, NL_AMGCL_EXT);
                } else
#endif
                {
                    if(verbose_) {
                        nlEnable(NL_VERBOSE);
                        Logger::out("LSCM") << "using JacobiCG"
                                            << std::endl;
                    }
                }
            }
            if(!spectral_) {
                project();
            }
            nlSolverParameteri(NL_NB_VARIABLES, NLint(2*nb_vertices));
            nlSolverParameteri(NL_LEAST_SQUARES, NL_TRUE);
            nlSolverParameteri(NL_MAX_ITERATIONS, NLint(5*nb_vertices));
            if(spectral_) {
                nlSolverParameterd(NL_THRESHOLD, 0.0);
            } else {
                nlSolverParameterd(NL_THRESHOLD, 1e-10);
            }
            nlBegin(NL_SYSTEM);
            mesh_to_solver();
            nlBegin(NL_MATRIX);
            setup_lscm();
            nlEnd(NL_MATRIX);
            nlEnd(NL_SYSTEM);
            if(verbose_) {
                Logger::out("LSCM") << "Solving ..." << std::endl;
            }

            if(spectral_) {
                nlEigenSolve();
                if(verbose_) {
                    for(NLuint i=0; i<nb_eigens; ++i) {
                        Logger::out("LSCM") << "[" << i << "] "
                                            << nlGetEigenValue(i) << std::endl;
                    }
                }

                // Find first "non-zero" eigenvalue
                double small_eigen = ::fabs(nlGetEigenValue(0)) ;
                eigen_ = 1;
                for(NLuint i=1; i<nb_eigens; ++i) {
                    if(::fabs(nlGetEigenValue(i)) / small_eigen > 1e3) {
                        eigen_ = i ;
                        break ;
                    }
                }
            } else{
                nlSolve();
            }

            solver_to_mesh();
            normalize_uv();

            if(!spectral_) {
                if(verbose_) {
                    double time;
                    NLint iterations;
                    nlGetDoublev(NL_ELAPSED_TIME, &time);
                    nlGetIntegerv(NL_USED_ITERATIONS, &iterations);
                    Logger::out("LSCM") << "Solver time: " << time << std::endl;
                    Logger::out("LSCM") << "Used iterations: "
                                        << iterations << std::endl;
                }
            }
            nlDeleteContext(nlGetCurrent());
        }

    protected:

        /**
         * \brief Creates the LSCM equations in OpenNL.
         */
        void setup_lscm() {
            for(NLuint f=0; f<mesh_.facets.nb(); ++f) {
                setup_lscm(f);
            }
        }

        /**
         * \brief Creates the LSCM equations in OpenNL, related
         *  with a given facet.
         * \param[in] f the index of the facet.
         * \details no-need to triangulate the facet,
         *   we do that "virtually", by creating triangles
         *   radiating around vertex 0 of the facet.
         *   (however, this may be invalid for concave facets)
         */
        void setup_lscm(NLuint f) {
            NLuint nv = NLuint(mesh_.facets.nb_vertices(f));
            if(angle_.is_bound()) {
                index_t c0 = mesh_.facets.corners_begin(f);
                double a0 = angle_[c0];
                for(NLuint i=1; i<nv-1; ++i) {
                    double ai = angle_[c0+i];
                    double aip1 = angle_[c0+1+1];
                    setup_conformal_map_relations(
                        NLuint(mesh_.facets.vertex(f,0)),
                        NLuint(mesh_.facets.vertex(f,i)),
                        NLuint(mesh_.facets.vertex(f,i+1)),
                        a0, ai, aip1
                    );
                }
            } else {
                for(NLuint i=1; i<nv-1; ++i) {
                    setup_conformal_map_relations(
                        NLuint(mesh_.facets.vertex(f,0)),
                        NLuint(mesh_.facets.vertex(f,i)),
                        NLuint(mesh_.facets.vertex(f,i+1))
                    );
                }
            }
        }

        /**
         * \brief Computes the coordinates of the vertices of a triangle
         * in a local 2D orthonormal basis of the triangle's plane.
         * \param[in] p0 , p1 , p2 the 3D coordinates of the vertices of
         *   the triangle
         * \param[out] z0 , z1 , z2 the 2D coordinates of the vertices of
         *   the triangle
         */
        static void project_triangle(
            const vec3& p0,
            const vec3& p1,
            const vec3& p2,
            vec2& z0,
            vec2& z1,
            vec2& z2
        ) {
            vec3 X = p1 - p0;
            X = normalize(X);
            vec3 Z = cross(X,(p2 - p0));
            Z = normalize(Z);
            vec3 Y = cross(Z,X);
            const vec3& O = p0;

            double x0 = 0;
            double y0 = 0;
            double x1 = (p1 - O).length();
            double y1 = 0;
            double x2 = dot((p2 - O),X);
            double y2 = dot((p2 - O),Y);

            z0 = vec2(x0,y0);
            z1 = vec2(x1,y1);
            z2 = vec2(x2,y2);
        }

        /**
         * \brief Creates the LSCM equation in OpenNL, related with
         *   a given triangle, specified by vertex indices.
         * \param[in] v0 , v1 , v2 the indices of the three vertices of
         *   the triangle.
         * \details Uses the geometric form of LSCM equation:
         *  (Z1 - Z0)(U2 - U0) = (Z2 - Z0)(U1 - U0)
         *  Where Uk = uk + i.vk is the complex number
         *                       corresponding to (u,v) coords
         *       Zk = xk + i.yk is the complex number
         *                       corresponding to local (x,y) coords
         * There is no divide with this expression,
         *  this makes it more numerically stable in
         * the presence of degenerate triangles.
         */
        void setup_conformal_map_relations(
            NLuint v0, NLuint v1, NLuint v2
        ) {

            const vec3& p0 = Geom::mesh_vertex(mesh_, v0);
            const vec3& p1 = Geom::mesh_vertex(mesh_, v1);
            const vec3& p2 = Geom::mesh_vertex(mesh_, v2);

            vec2 z0,z1,z2;
            project_triangle(p0,p1,p2,z0,z1,z2);
            vec2 z01 = z1 - z0;
            vec2 z02 = z2 - z0;
            double a = z01.x;
            double b = z01.y;
            double c = z02.x;
            double d = z02.y;
            geo_assert(b == 0.0);

            // Note  : 2*id + 0 --> u
            //         2*id + 1 --> v
            NLuint u0_id = 2*v0    ;
            NLuint v0_id = 2*v0 + 1;
            NLuint u1_id = 2*v1    ;
            NLuint v1_id = 2*v1 + 1;
            NLuint u2_id = 2*v2    ;
            NLuint v2_id = 2*v2 + 1;

            // Note : rhs = 0

            // Real part
            nlBegin(NL_ROW);
            nlCoefficient(u0_id, -a+c) ;
            nlCoefficient(v0_id,  b-d) ;
            nlCoefficient(u1_id,   -c) ;
            nlCoefficient(v1_id,    d) ;
            nlCoefficient(u2_id,    a);
            nlEnd(NL_ROW);

            // Imaginary part
            nlBegin(NL_ROW);
            nlCoefficient(u0_id, -b+d);
            nlCoefficient(v0_id, -a+c);
            nlCoefficient(u1_id,   -d);
            nlCoefficient(v1_id,   -c);
            nlCoefficient(v2_id,    a);
            nlEnd(NL_ROW);
        }

        /**
         * \brief Creates the LSCM equation in OpenNL, related with
         *   a given triangle, specified by vertex indices, and with
         *   specified desired angles.
         * \details This version is used to recoved the u,v coordinates
         *   from the angles computed by ABF++.
         * \param[in] v0 , v1 , v2 the indices of the three vertices of
         *   the triangle.
         * \param[in] alpha0 , alpha1 , alpha2 the desired angles at the
         *   three vertices of the triangle
         */
        void setup_conformal_map_relations(
            NLuint v0, NLuint v1, NLuint v2,
            double alpha0, double alpha1, double alpha2
        ) {
            const vec3& p0 = Geom::mesh_vertex(mesh_, v0);
            const vec3& p1 = Geom::mesh_vertex(mesh_, v1);
            const vec3& p2 = Geom::mesh_vertex(mesh_, v2);

            double scaling = ::sin(alpha1) / ::sin(alpha2) ;
            double a = scaling * ::cos(alpha0);
            double b = scaling * ::sin(alpha0);

            double area = Geom::triangle_area(p0,p1,p2) ;
            double s = ::sqrt(area) ;

            // Note  : 2*id + 0 --> u
            //         2*id + 1 --> v
            NLuint u0_id = 2*v0    ;
            NLuint v0_id = 2*v0 + 1;
            NLuint u1_id = 2*v1    ;
            NLuint v1_id = 2*v1 + 1;
            NLuint u2_id = 2*v2    ;
            NLuint v2_id = 2*v2 + 1;

            // Note : rhs = 0

            // Real part
            nlRowScaling(s);
            nlBegin(NL_ROW);
            nlCoefficient(u0_id, 1.0 - a) ;
            nlCoefficient(v0_id, b) ;
            nlCoefficient(u1_id, a) ;
            nlCoefficient(v1_id, -b) ;
            nlCoefficient(u2_id, -1.0);
            nlEnd(NL_ROW);

            // Imaginary part
            nlRowScaling(s);
            nlBegin(NL_ROW);
            nlCoefficient(u0_id, -b);
            nlCoefficient(v0_id, 1.0-a);
            nlCoefficient(u1_id, b);
            nlCoefficient(v1_id, a);
            nlCoefficient(v2_id, -1.0);
            nlEnd(NL_ROW);
        }

        /**
         * \brief Copies u,v coordinates from OpenNL solver to the mesh.
         */
        void solver_to_mesh() {
            for(index_t i: mesh_.vertices) {
                double u = spectral_ ? nlMultiGetVariable(NLuint(2*i),eigen_)
                    : nlGetVariable(2*i);
                double v = spectral_ ? nlMultiGetVariable(NLuint(2*i+1),eigen_)
                    : nlGetVariable(2*i+1);
                tex_coord_[2*i] = u;
                tex_coord_[2*i+1] = v;
            }
        }

        /**
         * \brief Translates and scales tex coords in such a way that they fit
         * within the unit square.
         */
        void normalize_uv() {
            double u_min=  Numeric::max_float64();
            double v_min=  Numeric::max_float64();
            double u_max= -Numeric::max_float64();
            double v_max= -Numeric::max_float64();
            for(NLuint i=0; i<mesh_.vertices.nb(); ++i) {
                double u = tex_coord_[2*i];
                double v = tex_coord_[2*i+1];
                u_min = std::min(u_min, u);
                v_min = std::min(v_min, v);
                u_max = std::max(u_max, u);
                v_max = std::max(v_max, v);
            }
            double l = std::max(u_max-u_min,v_max-v_min);
            for(NLuint i=0; i<mesh_.vertices.nb(); ++i) {
                tex_coord_[2*i] -= u_min;
                tex_coord_[2*i] /= l;
                tex_coord_[2*i+1] -= v_min;
                tex_coord_[2*i+1] /= l;
            }
        }

        /**
         * \brief Tests whether a vertex is locked.
         * \param[in] v the index of the vertex
         * \retval true if the vertex is locked
         * \retval false otherwise
         */
        bool is_locked(index_t v) {
            return (v==locked_1_ || v==locked_2_);
        }

        /**
         * \brief Copies u,v coordinates from the mesh to OpenNL solver.
         */
        void mesh_to_solver() {
            for(NLuint i=0; i<mesh_.vertices.nb(); ++i) {
                double u = tex_coord_[2*i];
                double v = tex_coord_[2*i+1];
                nlSetVariable(2 * i    , u);
                nlSetVariable(2 * i + 1, v);
                if(!spectral_ && is_locked(i)) {
                    nlLockVariable(2 * i    );
                    nlLockVariable(2 * i + 1);
                }
            }
        }

        /**
         * \brief Chooses an initial solution, and locks two vertices.
         */
        void project() {
            // Get bbox
            double xmin =  1e30;
            double ymin =  1e30;
            double zmin =  1e30;
            double xmax = -1e30;
            double ymax = -1e30;
            double zmax = -1e30;

            for(index_t i: mesh_.vertices) {
                const vec3& p = Geom::mesh_vertex(mesh_,i);
                xmin = std::min(p.x, xmin);
                ymin = std::min(p.y, ymin);
                zmin = std::min(p.z, zmin);

                xmax = std::max(p.x, xmax);
                ymax = std::max(p.y, ymax);
                zmax = std::max(p.z, zmax);
            }

            double dx = xmax - xmin;
            double dy = ymax - ymin;
            double dz = zmax - zmin;

            vec3 V1,V2;

            // Find shortest bbox axis
            if(dx <= dy && dx <= dz) {
                if(dy > dz) {
                    V1 = vec3(0,1,0);
                    V2 = vec3(0,0,1);
                } else {
                    V2 = vec3(0,1,0);
                    V1 = vec3(0,0,1);
                }
            } else if(dy <= dx && dy <= dz) {
                if(dx > dz) {
                    V1 = vec3(1,0,0);
                    V2 = vec3(0,0,1);
                } else {
                    V2 = vec3(1,0,0);
                    V1 = vec3(0,0,1);
                }
            } else if(dz <= dx && dz <= dy) {
                if(dx > dy) {
                    V1 = vec3(1,0,0);
                    V2 = vec3(0,1,0);
                } else {
                    V2 = vec3(1,0,0);
                    V1 = vec3(0,1,0);
                }
            }

            // Project onto shortest bbox axis,
            // and lock extrema vertices

            double  umin = 1e30;
            double  umax = -1e30;

            for(index_t i: mesh_.vertices) {
                const vec3& p = Geom::mesh_vertex(mesh_,i);
                double u = dot(p,V1);
                double v = dot(p,V2);
                tex_coord_[2*i]   = u;
                tex_coord_[2*i+1] = v;
                if(u < umin) {
                    locked_1_ = i;
                    umin = u;
                }
                if(u > umax) {
                    locked_2_ = i;
                    umax = u;
                }
            }
        }

        Mesh& mesh_;

        Attribute<double>& tex_coord_;

        Attribute<double>& angle_;
        /**
         * \brief true if spectral mode is used,
         *  false if locked least squares mode is used.
         */
        bool spectral_;

        /**
         * \brief In spectral mode, the index of the first
         *  non-zero eigenvalue.
         */
        NLuint eigen_;

        /**
         * \brief The indices of the two locked vertices.
         */
        index_t locked_1_, locked_2_;

        bool verbose_;
    };

}

namespace GEO {

    void mesh_compute_LSCM(
        Mesh& M, const std::string& attribute_name, bool spectral,
        const std::string& angle_attribute_name,
        bool verbose
    ) {
        Attribute<double> tex_coord;
        tex_coord.bind_if_is_defined(M.vertices.attributes(), attribute_name);
        if(tex_coord.is_bound() && tex_coord.dimension() != 2) {
            Logger::err("LSCM") << "Attribute " << attribute_name
                                << " already exists in mesh with dimension "
                                << tex_coord.dimension()
                                << " (expected 2)"
                                << std::endl;
            return;
        }
        if(!tex_coord.is_bound()) {
            tex_coord.create_vector_attribute(
                M.vertices.attributes(),attribute_name,2
            );
        }
        Attribute<double> angle;
        angle.bind_if_is_defined(
            M.facet_corners.attributes(), angle_attribute_name
        );
        LSCM lscm(M,tex_coord,angle);
        lscm.set_spectral(spectral);
        lscm.set_verbose(verbose);
        lscm.apply();
    }
}
