/*
 *  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/mesh/mesh_baking.h>
#include <geogram/mesh/mesh.h>
#include <geogram/mesh/mesh_geometry.h>
#include <geogram/mesh/mesh_AABB.h>
#include <geogram/image/image_rasterizer.h>
#include <geogram/points/kd_tree.h>

namespace {
    using namespace GEO;

    /**
     * \brief Reads a (facet or vertex or facet corner) vector attribute
     *  as a color.
     * \details If the vector attribute has less components than 4, then
     *  the additional color components are set to (0 + \p bias)* \p scale.
     *  If the vector attribute has more components than 4, then the additional
     *  attribute components are ignored.
     * \param[out] C the color
     * \param[in] attribute a reference to the attribute.
     * \param[in] attrib_loc one of MESH_FACETS,MESH_VERTICES,MESH_FACET_CORNERS
     * \param[in] f the facet
     * \param[in] c the facet corner
     * \param[in] v the vertex
     * \param[in] bias optional value to be added to the attribute values
     *  before storing them into the color components.
     * \param[in] scale optional value that scales the attribute values
     *  after \p bias is added and before storing them into the color
     *  components.
     */
    inline void get_attribute_as_color(
        Color& C,
        Attribute<double>& attribute,
        MeshElementsFlags attrib_loc,
        index_t f, index_t c, index_t v,
        double bias = 0.0,
        double scale = 1.0
    ) {
        index_t dim = attribute.dimension();
        index_t base = NO_INDEX;
        if(attrib_loc == MESH_FACETS) {
            base = f*dim;
        } else if(attrib_loc == MESH_VERTICES) {
            base = v*dim;
        } else if(attrib_loc == MESH_FACET_CORNERS) {
            base = c*dim;
        } else {
            geo_assert_not_reached;
        }
        C.set_r(attribute[base]);
        C.set_g((dim >= 2) ? attribute[base+1] : 0.0);
        C.set_b((dim >= 3) ? attribute[base+2] : 0.0);
        C.set_a((dim >= 4) ? attribute[base+3] : 0.0);

        C.set_r(scale*(C.r() + bias));
        C.set_g(scale*(C.g() + bias));
        C.set_b(scale*(C.b() + bias));
        C.set_a(scale*(C.a() + bias));
    }
}

namespace GEO {

    /**************************************************************************/

    void bake_mesh_facet_normals(Mesh* mesh, Image* target) {
        Attribute<double> tex_coord;
        tex_coord.bind_if_is_defined(
            mesh->facet_corners.attributes(), "tex_coord"
        );
        geo_assert(tex_coord.is_bound() && tex_coord.dimension() == 2);
        ImageRasterizer rasterizer(target);
        for(index_t f: mesh->facets) {
            vec3 N = normalize(Geom::mesh_facet_normal(*mesh, f));
            Color C = 0.5*Color(N.x+1.0, N.y+1.0, N.z+1.0, 2.0);
            index_t c1 = mesh->facets.corners_begin(f);
            vec2 p1(tex_coord[2*c1], tex_coord[2*c1+1]);
            for(index_t c2 = c1+1;
                c2+1 < mesh->facets.corners_end(f); ++c2) {
                index_t c3 = c2+1;
                vec2 p2(tex_coord[2*c2], tex_coord[2*c2+1]);
                vec2 p3(tex_coord[2*c3], tex_coord[2*c3+1]);
                rasterizer.triangle(
                    p1, C,
                    p2, C,
                    p3, C
                );
            }
        }
    }

    /**************************************************************************/

    void bake_mesh_vertex_normals(Mesh* mesh, Image* normal_map) {

        // Step 1: compute vertex normals.
        Attribute<double> N;
        N.create_vector_attribute(mesh->vertices.attributes(), "N", 3);
        for(index_t v: mesh->vertices) {
            N[3*v]   = 0.0;
            N[3*v+1] = 0.0;
            N[3*v+2] = 0.0;
        }
        for(index_t f: mesh->facets) {
            vec3 Nf = GEO::Geom::mesh_facet_normal(*mesh, f);
            for(index_t corner : mesh->facets.corners(f)) {
                index_t v = mesh->facet_corners.vertex(corner);
                N[3*v]   += Nf.x;
                N[3*v+1] += Nf.y ;
                N[3*v+2] += Nf.z;
            }
        }
        for(index_t v: mesh->vertices) {
            vec3 Nf(N[3*v], N[3*v+1], N[3*v+2]);
            Nf = normalize(Nf);
            N[3*v]   = Nf.x;
            N[3*v+1] = Nf.y;
            N[3*v+2] = Nf.z;
        }

        // Step 2: bake interpolated vertex normals.
        bake_mesh_attribute(mesh, normal_map, N);

        // Step 3: normalize interpolated normals.
        FOR(y, normal_map->height()) {
            FOR(x, normal_map->width()) {
                double* pix = normal_map->pixel_base_float64_ptr(x,y);
                vec3 Npix(pix);
                Npix = normalize(Npix);
                pix[0] = Npix.x;
                pix[1] = Npix.y;
                pix[2] = Npix.z;
            }
        }
        N.unbind();
        mesh->vertices.attributes().delete_attribute_store("N");
    }

    /**************************************************************************/

    void bake_mesh_attribute(
        Mesh* mesh, Image* target, Attribute<double>& attribute,
        double bias, double scale
    ) {
        Attribute<double> tex_coord;
        tex_coord.bind_if_is_defined(
            mesh->facet_corners.attributes(), "tex_coord"
        );
        geo_assert(tex_coord.is_bound() && tex_coord.dimension() == 2);
        geo_assert(attribute.is_bound());

        MeshElementsFlags attrib_loc = MESH_NONE;

        if(attribute.manager() == &mesh->vertices.attributes()) {
            attrib_loc = MESH_VERTICES;
        } else if(attribute.manager() == &mesh->facets.attributes()) {
            attrib_loc = MESH_FACETS;
        } else if(attribute.manager() == &mesh->facet_corners.attributes()) {
            attrib_loc = MESH_FACET_CORNERS;
        } else {
            // Attribute is not a vertex/facet/facet_corner attribute
            // or it is bound to a different mesh.
            geo_assert_not_reached;
        }

        ImageRasterizer rasterizer(target);
        for(index_t f: mesh->facets) {
            Color C1,C2,C3;
            index_t c1 = mesh->facets.corners_begin(f);
            vec2 p1(tex_coord[2*c1], tex_coord[2*c1+1]);
            index_t v1 = mesh->facet_corners.vertex(c1);
            get_attribute_as_color(
                C1, attribute, attrib_loc, f, c1, v1, bias, scale
            );
            for(index_t c2 = c1+1;
                c2+1 < mesh->facets.corners_end(f); ++c2) {
                index_t c3 = c2+1;
                vec2 p2(tex_coord[2*c2], tex_coord[2*c2+1]);
                vec2 p3(tex_coord[2*c3], tex_coord[2*c3+1]);
                index_t v2 = mesh->facet_corners.vertex(c2);
                index_t v3 = mesh->facet_corners.vertex(c3);
                get_attribute_as_color(
                    C2, attribute, attrib_loc, f, c2, v2, bias, scale
                );
                get_attribute_as_color(
                    C3, attribute, attrib_loc, f, c3, v3, bias, scale
                );
                rasterizer.triangle(
                    p1, C1,
                    p2, C2,
                    p3, C3
                );
            }
        }
    }

    /**************************************************************************/

    void bake_mesh_geometry(Mesh* mesh, Image* target, bool clear) {
        geo_assert(target->component_encoding() == Image::FLOAT64);
        Attribute<double> point;
        point.bind_if_is_defined(mesh->vertices.attributes(), "point");
        geo_assert(point.is_bound());

        if(clear) {
            double* base_mem = target->base_mem_float64_ptr();
            size_t nb = target->nb_pixels()*target->components_per_pixel();
            for(size_t i=0; i<nb; ++i) {
                base_mem[i] = Numeric::max_float64();
            }
        }
        bake_mesh_attribute(mesh, target, point);
    }

    /**************************************************************************/

    void bake_mesh_facet_normals_indirect(
        Image* geometry, Image* target, Mesh* highres
    ) {
        geo_assert(geometry->dimension() == 2);
        geo_assert(target->dimension() == 2);
        geo_assert(
            geometry->width() == target->width() &&
            geometry->height() == target->height()
        );
        geo_assert(geometry->component_encoding() == Image::FLOAT64);

        MeshFacetsAABB AABB(*highres);
        vector<Color> normal(highres->facets.nb());
        for(index_t f: highres->facets) {
            vec3 N = normalize(Geom::mesh_facet_normal(*highres,f));
            normal[f] = Color(
                0.5*(N.x+1.0), 0.5*(N.y+1.0), 0.5*(N.z+1.0)
            );
        }
        ImageRasterizer rasterizer(target);

        parallel_for(
            0, target->height(),
            [geometry, target, &AABB, &rasterizer, &normal](index_t y) {
                index_t nearest_facet = NO_FACET;
                vec3 nearest_point;
                double sq_dist;
                for(index_t x=0; x<target->width(); ++x) {
                    vec3 p((double*)(void*)geometry->pixel_base(x,y));
                    if(
                        p[0] == Numeric::max_float64() &&
                        p[1] == Numeric::max_float64() &&
                        p[2] == Numeric::max_float64()
                    ) {
                        continue;
                    }
                    if(nearest_facet == NO_FACET) {
                        nearest_facet =
                            AABB.nearest_facet(p, nearest_point, sq_dist
                                              );
                    } else {
                        sq_dist = length2(p - nearest_point);
                        AABB.nearest_facet_with_hint(
                            p,nearest_facet,nearest_point,sq_dist
                        );
                    }
                    rasterizer.set_pixel(
                        int(x),int(y),normal[nearest_facet]
                    );
                }
            }
        );
    }

    /**************************************************************************/

    void bake_mesh_points_attribute_indirect(
        Image* geometry, Image* target,
        Mesh* highres, Attribute<double>& attribute,
        double bias, double scale
    ) {
        geo_assert(geometry->dimension() == 2);
        geo_assert(target->dimension() == 2);
        geo_assert(
            geometry->width() == target->width() &&
            geometry->height() == target->height()
        );
        geo_assert(geometry->component_encoding() == Image::FLOAT64);
        geo_assert(highres->vertices.dimension() >= 3);

        NearestNeighborSearch_var kd_tree = new BalancedKdTree(3);

        kd_tree->set_points(
            highres->vertices.nb(), highres->vertices.point_ptr(0),
            highres->vertices.dimension()
        );

        ImageRasterizer rasterizer(target);

        parallel_for(
            0, target->height(),
            [geometry, target, kd_tree,
             &rasterizer, &attribute, bias, scale
            ](index_t y) {

                index_t nearest_vertex = NO_VERTEX;
                double sq_dist;
                Color C;
                for(index_t x=0; x<target->width(); ++x) {
                    vec3 p((double*)(void*)geometry->pixel_base(x,y));
                    if(
                        p[0] == Numeric::max_float64() &&
                        p[1] == Numeric::max_float64() &&
                        p[2] == Numeric::max_float64()
                    ) {
                        continue;
                    }
                    kd_tree->get_nearest_neighbors(
                        1, p.data(), &nearest_vertex, &sq_dist
                    );
                    get_attribute_as_color(
                        C, attribute, MESH_VERTICES,
                        NO_INDEX, NO_INDEX, nearest_vertex,
                        bias, scale
                    );
                    rasterizer.set_pixel(int(x),int(y),C);
                }
            }
        );
    }

    /**************************************************************************/

}
