/**
 * ------------------------------------
 * @author: Weipeng Kong
 * @date: 2021/11/17
 * @email: yjxkwp@foxmail.com
 * @site: https://donot.fit
 * @description: 
 * ------------------------------------
**/

#include <iostream>
#include "Octree/sdf/SDFOctree.h"
#include "Octree/OctreeBuilder.h"
#include "Octree/OctreeTraverser.h"
#include "Octree/sdf/SDFTraversalSampler.h"
#include <igl/marching_cubes.h>
#include <igl/voxel_grid.h>
#include <igl/writeOBJ.h>
#include <pMesh/mesh/TriangleMesh.h>
#include <pMesh/mesh/HexahedronMesh.h>
#include <pMesh/io/reader/OBJReader.h>
#include <pMesh/io/reader/VTKReader.h>
#include <pMesh/io/writer/VTKWriter.h>
#include <pMesh/io/adapter/DefaultReadAdapter.h>
#include <pMesh/io/adapter/DefaultWriteAdapter.h>
#include <boost/timer.hpp>
#include <boost/log/trivial.hpp>
#include <pMesh/io/reader/BaseReader.h>
#include <pMesh/io/adapter/DefaultReadAdapter.h>
#include "test-path.h"


class VTKTraverser : public pMesh::io::BaseReader, public Octree::OctreeTraverser {
    int node_cnt = 0;
public:
    explicit VTKTraverser(const Octree::SDFOctree &octree, const TraverseStrategy strategy = DFS);

    void visit_node(Octree::OctreeNode &node) override;

    pMesh::io::ReadAdapter *reader_adapter;

    bool operator>>(pMesh::io::ReadAdapter &adapter) override;
};

int main() {
    pMesh::io::fs_path data_base_path = TEST_DATA_BASE_PATH;
    BOOST_LOG_TRIVIAL(debug) << "base data path is " << boost::filesystem::absolute(data_base_path);

    auto mesh_path = data_base_path / "stanford-bunny.obj";
    auto out_octree_path = data_base_path / "stanford-bunny-octree.vtk";
    auto out_mc_path = data_base_path / "stanford-bunny-mc.obj";
    auto out_error_path = data_base_path / "stanford-bunny-error_points.vtk";

    pMesh::Triangle3dMesh<> mesh;
    pMesh::io::OBJReader(mesh_path)
            >> pMesh::io::DefaultSurfaceReadAdapter<>(mesh, false)();

    BOOST_LOG_TRIVIAL(info) << "Load Face #" << mesh.f_size();

    const int level = 4;
    const double sample_step = 0.003;

    boost::timer t;
    Octree::AABB aabb(mesh.aabb());
    Octree::SDFOctree octree(level, aabb);

    aabb.extend(2);
    Octree::OctreeBuilder builder(mesh, octree);
    builder.build();
    BOOST_LOG_TRIVIAL(debug) << "time elapse " << t.elapsed();

    auto aabb_size = octree.aabb().size();
    BOOST_LOG_TRIVIAL(debug) << "Global AABB " << aabb_size.transpose();

    auto final_aabb_size = aabb_size / pow(2, level);
    BOOST_LOG_TRIVIAL(debug) << "Local AABB " << final_aabb_size.transpose();

    // origin points of model
    int opm = std::ceil(aabb_size.x() / sample_step) *
              std::ceil(aabb_size.y() / sample_step) *
              std::ceil(aabb_size.z() / sample_step);

    BOOST_LOG_TRIVIAL(debug) << "Origin Points of model = " << opm;

    // points per voxel
    int ppv = std::ceil(final_aabb_size.x() / sample_step) *
              std::ceil(final_aabb_size.y() / sample_step) *
              std::ceil(final_aabb_size.z() / sample_step);

    BOOST_LOG_TRIVIAL(debug) << "Points per voxel = " << ppv;

    pMesh::HexahedronMesh<> hexahedron_mesh;
    VTKTraverser(octree) >> pMesh::io::DefaultVolumeReadAdapter<>(hexahedron_mesh)();

    pMesh::io::VTKWriter(12, out_octree_path)
            << pMesh::io::DefaultVolumeWriteAdapter<>(hexahedron_mesh)();

    BOOST_LOG_TRIVIAL(debug) << "Total Points = " << hexahedron_mesh.c_size() * ppv;

    BOOST_LOG_TRIVIAL(debug) << "Compression ratio = " << (double) hexahedron_mesh.c_size() * ppv / opm;

    BOOST_LOG_TRIVIAL(debug) << "Calculating SDF";
    Octree::SDFTraversalBuilder sdf_builder(octree, mesh, sample_step);
    sdf_builder.build();
    BOOST_LOG_TRIVIAL(debug) << "End computing SDF";

    if (1) {
        Eigen::MatrixXd V = Eigen::MatrixXd(mesh.v_size(), 3);
        Eigen::MatrixXi F = Eigen::MatrixXi(mesh.f_size(), 3);

        for (int i = 0; i < mesh.v_size(); ++i) {
            V.row(i) = mesh.vertices[i].attr.coordinate;
        }

        for (int i = 0; i < mesh.f_size(); ++i) {
            const auto &face = mesh.faces[i].attr.vertices;
            F.row(i) << face[0].id(), face[1].id(), face[2].id();
        }

        Eigen::MatrixXd GV;
        Eigen::RowVector3i res;
        const int s = 100;
        igl::voxel_grid(V, 0, s, 1, GV, res);

        // compute values
        std::cout << "Computing distances..." << std::endl;
        Eigen::VectorXd S = Eigen::VectorXd(GV.rows()), B;

#if 0
        pMesh::Triangle3dMesh<> error_points;
#endif

        {
            for (int i = 0; i < GV.rows(); ++i) {
                auto node = octree.map_node(GV.row(i));
                if (node == nullptr) {
                    S[i] = 1000;
//                    BOOST_LOG_TRIVIAL(debug) << "dd";
#if 0
                    error_points.vertices.emplace_back(pMesh::Surface::Vertex(0, GV.row(i)));
                    error_points.faces.emplace_back(pMesh::Surface::Face(
                            {.vertices={pMesh::Triangle3dMesh<>::VertexHandle(error_points.vertices.size() - 1)}}));
#endif
                } else {
                    S[i] = node->get_sdf(GV.row(i));
                }
            }
            // Convert distances to binary inside-outside data --> aliasing artifacts
//            B = S;
//            for_each(B.data(),B.data()+B.size(),[](double& b){b=(b>0?1:(b<0?-1:0));});
        }

#if 0
        pMesh::io::VTKWriter(1, out_error_path) << pMesh::io::DefaultSurfaceWriteAdapter(error_points)();
#endif

        std::cout << "Marching cubes..." << std::endl;
        Eigen::MatrixXd SV, BV;
        Eigen::MatrixXi SF, BF;

        igl::marching_cubes(S, GV, res(0), res(1), res(2), 0, SV, SF);


        igl::writeOBJ(out_mc_path.string(), SV, SF);
//        igl::marching_cubes(B,GV,res(0),res(1),res(2),0,BV,BF);
    }

    return 0;
}


void VTKTraverser::visit_node(Octree::OctreeNode &node) {
    if (node.get_level() != octree.get_max_level()) {
//        assert(!node.is_leaf());
        assert(node.tri_ids.empty());
        return;
    }

    auto min = node.get_aabb().min();
    auto max = node.get_aabb().max();
    reader_adapter->feed_vertex({min.x(), min.y(), min.z()});
    reader_adapter->feed_vertex({max.x(), min.y(), min.z()});
    reader_adapter->feed_vertex({max.x(), max.y(), min.z()});
    reader_adapter->feed_vertex({min.x(), max.y(), min.z()});
    reader_adapter->feed_vertex({min.x(), min.y(), max.z()});
    reader_adapter->feed_vertex({max.x(), min.y(), max.z()});
    reader_adapter->feed_vertex({max.x(), max.y(), max.z()});
    reader_adapter->feed_vertex({min.x(), max.y(), max.z()});

    reader_adapter->feed_collection(
            {
                    0 + node_cnt, 1 + node_cnt, 2 + node_cnt, 3 + node_cnt, 4 + node_cnt,
                    5 + node_cnt, 6 + node_cnt, 7 + node_cnt
            });
    node_cnt += 8;
}

bool VTKTraverser::operator>>(pMesh::io::ReadAdapter &adapter) {
    reader_adapter = &adapter;
    node_cnt = 0;
    adapter.start();

    this->traverse();

    adapter.end();
    return true;
}


VTKTraverser::VTKTraverser(const Octree::SDFOctree &octree, const OctreeTraverser::TraverseStrategy strategy)
        : OctreeTraverser(
        octree, strategy) {
}