// ========================================================================== //
// Class inclusion
#include "Splatter/Splatter.hpp"

// Local libraries
#include "Options/options.hpp"
#include "Timer/Timer.hpp"
#include "Utilities/string_operations.hpp"

// External libraries
#include <GEL/CGLA/CGLA.h>
#include <GEL/Geometry/GridAlgorithm.h>
#include <GEL/Geometry/KDTree.h>
#include <GEL/Geometry/RGrid.h>
#include <GEL/Geometry/rply.h>
#include <GEL/HMesh/HMesh.h>
#include <GEL/HMesh/cleanup.h>
#include <GEL/HMesh/polygonize.h>
#include <GEL/Util/ArgExtracter.h>

// Standard libraries
#include <filesystem>
#include <fstream>
#include <future>
#include <iostream>
#include <sstream>
#include <thread>

// Namespace usage and simplification
namespace fs = std::filesystem;

// ========================================================================== //
// Public API

// Default construction
Splatter::Splatter(std::string opt_name) {
    Options opt(opt_name);
    constructor(opt);
}
// Option based construction
Splatter::Splatter(Options &opt) { constructor(opt); }
// De-construction
Splatter::~Splatter(){};

// Surface import function
void Splatter::import_surfaces(std::string folder) {

    if (folder.size() == 0) folder = export_location + "surfaces_optimal/";

    for (auto file : fs::directory_iterator(folder)) {
        if (file.path().extension().compare(".obj") == 0)
            files.push_back(file.path());
    }

    // Load in the meshes
    Utilities::debug_info("Loading files ...", debug);
    for (auto f : files) {
        Utilities::debug_info("Opening " + f, debug);

        mesh_vec.push_back(MeshVCol());
        HMesh::Manifold &m = mesh_vec.back().m;
        vcol_obj_load(f, m, mesh_vec.back().vcol);
    }

    if (mesh_vec.size() < 1)
        Utilities::error(__FILE__, __LINE__, "No surfaces loaded.");
}
// Main surface splatting
void Splatter::splat_surfaces() {
    prepare_splatter();

    // Construct lambda function used in splatting
    const double      global_rad = desired_spacing * xform.inv_scale();
    const CGLA::Vec3i dim        = global_grid.get_dims();
    auto              splat      = [&](int m_no) -> std::vector<MicroGrid> {
        std::vector<MicroGrid> micro_grids;
        const auto &[m, vcol] = mesh_vec[m_no];
        for (auto f : m.faces()) {
            HMesh::Walker   w  = m.walker(f);
            HMesh::VertexID v0 = w.vertex();
            HMesh::VertexID v1 = w.next().vertex();
            HMesh::VertexID v2 = w.next().next().vertex();
            Triangle        tri(m.pos(v0), m.pos(v1), m.pos(v2));

            const auto [bbmin, bbmax] = tri.bbox(global_rad);
            const CGLA::Vec3i bbmin_i =
                CGLA::v_min(dim, CGLA::v_max(CGLA::Vec3i(0),
                                             CGLA::Vec3i(xform.apply(bbmin))));
            const CGLA::Vec3i bbmax_i =
                CGLA::v_min(dim, CGLA::v_max(CGLA::Vec3i(0),
                                             CGLA::Vec3i(xform.apply(bbmax)) +
                                                 CGLA::Vec3i(1)));

            const CGLA::Vec3d t_vec =
                CGLA::Vec3d(vcol[v0][0], vcol[v1][0], vcol[v2][0]);

            MicroGrid micro_grid(bbmin_i, bbmax_i);
            for (CGLA::Vec3i pi : Geometry::Range3D(bbmin_i, bbmax_i)) {

                // Compute barycentric coordinates for the point
                const CGLA::Vec3d p_voxel = xform.inverse(CGLA::Vec3d(pi));
                const CGLA::Vec4d b_d = tri.bary_dist(p_voxel);

                double a =
                    desired_spacing *
                    std::max(
                        threshold_void,
                        CGLA::dot(CGLA::Vec3d(b_d[0], b_d[1], b_d[2]), t_vec));
                double b = 0.0;
                double t = -b_d[3];

                const double d = smooth_step(-a, b, t);
                micro_grid.push(d);
            }
            micro_grids.push_back(micro_grid);
        }
        return micro_grids;
    };

    // Run the splatting locally for each mesh
    Utilities::debug_info("Splatting: ...", debug);
    std::vector<std::future<std::vector<MicroGrid>>> micro_grid_futures(
        mesh_vec.size());
    for (size_t i = 0; i < mesh_vec.size(); ++i)
        micro_grid_futures[i] = std::async(std::launch::async, splat, i);

    // Combine all the local volumes
    Utilities::debug_info("\nAdding Micro Grids ...", debug);
    for (size_t i = 0; i < mesh_vec.size(); ++i) {
        Geometry::RGrid<float> grid(dim, 0.0);
        Geometry::RGrid<float> wgrid(dim, 0.0);
        auto                   mgv = micro_grid_futures[i].get();
        for (const auto &g : mgv) g.store(grid, wgrid);

        for (CGLA::Vec3i pi : Geometry::Range3D(dim))
            global_grid[pi] = std::max(global_grid[pi], grid[pi]);
    }
}
// Add solid regions
void Splatter::splat_solid() {
    // Add the solid regions
    Utilities::debug_info("Inserting blocks...", debug);
    std::ifstream ifs(density_file.c_str());
    std::string   line_buffer;
    while (getline(ifs, line_buffer)) {
        float i, j, k;
        float d;
        sscanf(line_buffer.c_str(), "%f,%f,%f,%f", &i, &j, &k, &d);
        if (d > threshold_solid) {
            const CGLA::Vec3d pijk(i, j, k);
            const CGLA::Vec3i bbmin_i = CGLA::v_min(
                global_grid.get_dims(),
                CGLA::v_max(CGLA::Vec3i(0),
                            CGLA::Vec3i(xform.apply(pijk - CGLA::Vec3d(0.5)))));
            const CGLA::Vec3i bbmax_i = CGLA::v_min(
                global_grid.get_dims(),
                CGLA::v_max(CGLA::Vec3i(0),
                            CGLA::Vec3i(xform.apply(pijk + CGLA::Vec3d(0.5))) +
                                CGLA::Vec3i(2)));
            for (CGLA::Vec3i pi : Geometry::Range3D(bbmin_i, bbmax_i)) {
                global_grid[pi] = 1.0;
            }
        }
    }
}
// Extract the interface mesh
void Splatter::extract_iso_contour(double iso) {
    estiamte_volume_fraction(iso);

    // Construct the ISO-contour mesh
    Utilities::debug_info("Meshing", debug);
    HMesh::volume_polygonize(xform, global_grid, mesh, iso, true, true, 0);

    // Post processing
    Utilities::debug_info("Post processing mesh", debug);
    // HMesh::remove_needles(mesh);
}
// Save the extracted mesh
void Splatter::save_mesh(double simplify_to) {

    // Save the mesh and finalize the computation
    Utilities::debug_info("Saving ...", debug);
    fs::create_directory("result/splatter/");
    obj_save("result/splatter/mesh.obj", mesh);

    // Simplified version
    simplify_to = std::min(splatter_simplify_to, simplify_to);
    if (mesh.no_vertices() > simplify_to) {
        Utilities::debug_info("Simplifying ...", debug);
        std::string simpName =
            "mesh_" +
            std::to_string((int)(simplify_to / pow(10, log10(simplify_to)))) +
            "e" + std::to_string((int)log10(simplify_to)) + ".obj";
        quadric_simplify(mesh, simplify_to / mesh.no_vertices());
        obj_save("result/splatter/" + simpName, mesh);
    }
}

// ========================================================================== //
// Private IO functions

// Option reading
void Splatter::read_options(Options &opt) {

    opt.get("debug", debug);
    opt.get("splatter_toggle", splatter_toggle);
    opt.get("data_name", data_name);
    opt.get("data_location", data_location);
    opt.get("export_location", export_location);
    opt.get("threshold_void", threshold_void);
    opt.get("threshold_solid", threshold_solid);
    opt.get("desired_spacing", desired_spacing);
    opt.get("splatter_min_scaling", splatter_min_scaling);
    opt.get("splatter_max_scaling", splatter_max_scaling);
    opt.get("splatter_simplify_to", splatter_simplify_to);
    opt.get("splatter_solid_overlap", splatter_solid_overlap);

    opt.write_file();

    // Ensure options are correctly setup
    data_name += data_name.back() != '/' ? "/" : "";
    data_location += data_location.back() != '/' ? "/" : "";
    export_location += export_location.back() != '/' ? "/" : "";
    info_file    = data_location + data_name + "information.csv";
    density_file = export_location + "intermediate/densities.csv";

    threshold_solid -= splatter_solid_overlap;

    // Read the info file if there
    if (fs::exists(info_file)) {
        CGLA::Vec3d   _pmin, _pmax;
        std::ifstream f(info_file);
        char          dummy;
        _pmin = CGLA::Vec3d(0);
        f >> _pmax[0] >> dummy >> _pmax[1] >> dummy >> _pmax[2];

        bb.extend(_pmin);
        bb.extend(_pmax);
    }
}
// Mesh import function using colors
bool Splatter::vcol_obj_load(const std::string &filename, HMesh::Manifold &m,
                             HMesh::VertexAttributeVector<CGLA::Vec3d> &vcol) {
    std::ifstream obj_file(filename.data());

    if (obj_file) {
        std::string              buf;
        std::vector<CGLA::Vec3d> vertices;
        std::vector<CGLA::Vec3d> colors;
        std::vector<int>         faces;
        std::vector<int>         indices;
        while (get_multi_line(obj_file, buf)) {
            std::istringstream iss(std::move(buf));
            std::string        code;
            if (iss >> code) {
                if (code == "v") {
                    CGLA::Vec3d v, c;
                    iss >> v >> c;
                    vertices.push_back(v);
                    colors.push_back(c);
                } else if (code == "f") {
                    std::string str;
                    int         ctr = 0;
                    while (iss >> str) {
                        int v, t, n;
                        if (sscanf(str.c_str(), "%d/%d/%d", &v, &t, &n) == 3 ||
                            sscanf(str.c_str(), "%d//%d", &v, &n) == 2 ||
                            sscanf(str.c_str(), "%d/%d", &v, &t) == 2 ||
                            sscanf(str.c_str(), "%d", &v) == 1) {
                            indices.push_back(v > 0 ? v - 1
                                                    : vertices.size() + v);
                            ++ctr;
                        }
                    }
                    faces.push_back(ctr);
                }
            }
        }
        m.clear();

        auto orig_vertex_indices =
            build(m, vertices.size(), reinterpret_cast<double *>(&vertices[0]),
                  faces.size(), &faces[0], &indices[0]);

        for (auto v : m.vertices()) {
            vcol[v]    = colors[orig_vertex_indices[v]];
            vcol[v][0] = std::max(std::min(vcol[v][0], 1.0), 0.05);
        }

        for (int iter = 0; iter < 5; ++iter) {
            auto vcol_filtered = vcol;
            for (auto v : m.vertices()) {
                float w_sum = 1.0;
                for (auto n : m.incident_vertices(v)) {
                    double w = 1.0 - std::abs(vcol[v][0] - vcol[n][0]);
                    vcol_filtered[v] += w * vcol[n];
                    w_sum += w;
                }
                vcol_filtered[v] /= w_sum;
            }
            vcol = vcol_filtered;
        }
        return true;
    }
    return false;
}

// ========================================================================== //
// Private tools

// Constructor tool
void Splatter::constructor(Options &opt) {

    // Setup the options and exit if we are told to do so.
    read_options(opt);
    if (!splatter_toggle) exit(EXIT_SUCCESS);
}
// Prepare grids and systems for splatting
void Splatter::prepare_splatter() {

    // Setup the bounding box
    for (const auto &[m, vcol] : mesh_vec) {
        CGLA::Vec3d _pmin, _pmax;
        bbox(m, _pmin, _pmax);

        bb.extend(_pmin);
        bb.extend(_pmax);
    }

    // Setup the voxel grid max scaling always overule.
    CGLA::Vec3d diagonal = bb.diagonal();
    desired_spacing      = diagonal.min_coord() / desired_spacing;
    double scaling       = 1.5 / (threshold_void * desired_spacing);
    scaling =
        std::min(std::max(scaling, splatter_min_scaling), splatter_max_scaling);

    // Setting up the Boundinng boxes and XForm object
    CGLA::Vec3i dim = CGLA::Vec3i(scaling * diagonal) + CGLA::Vec3i(1);
    xform           = Geometry::XForm(bb.min_corner, bb.max_corner, dim, 0.0);
    global_grid     = Geometry::RGrid<float>(dim, 0.0);

    if (debug) {
        std::cout // Cout
            << "Bounding box: " << bb.min_corner << bb.max_corner << "\n"
            << "Volume dimensions: " << dim << "\n"
            << "XForm scaling factor: " << xform.get_scale() << "\n"
            << "Bounding box corners to voxel grid coordinates: \n"
            << bb.min_corner << " -> " << xform.apply(bb.min_corner) << "\n"
            << bb.max_corner << " -> " << xform.apply(bb.max_corner) << "\n"
            << "Voxel grid corners to bounding box coordinates: \n"
            << CGLA::Vec3d(0) << " -> " << xform.inverse(CGLA::Vec3d(0)) << "\n"
            << CGLA::Vec3d(dim - CGLA::Vec3i(1)) << " -> "
            << xform.inverse(CGLA::Vec3d(dim - CGLA::Vec3i(1))) << "\n";
    }
}
// Estimate number volume fraction inside the interface mesh
void Splatter::estiamte_volume_fraction(double iso) {

    int N = global_grid.get_size();

    double n_s = 0;
    for (CGLA::Vec3i pi : Geometry::Range3D(global_grid.get_dims()))
        n_s += global_grid[pi] > iso ? 1.0 : 0.0;

    printf("Volume: %.2e / %.2e = %.2f\n", n_s, (double)N, n_s / (double)N);
}

// ========================================================================== //
// Helper Class and Struct implementations

// Triangle class implementations
std::pair<CGLA::Vec3d, CGLA::Vec3d> Splatter::Triangle::bbox(double rad) {
    CGLA::Vec3d pmin =
        CGLA::v_min(CGLA::v_min(vert[0], vert[1]), vert[2]) - CGLA::Vec3d(rad);
    CGLA::Vec3d pmax =
        CGLA::v_max(CGLA::v_max(vert[0], vert[1]), vert[2]) + CGLA::Vec3d(rad);
    return std::make_pair(pmin, pmax);
}
CGLA::Vec4d Splatter::Triangle::bary_dist(const CGLA::Vec3d &p) {
    CGLA::Vec3d normal =
        CGLA::normalize(CGLA::cross(vert[1] - vert[0], vert[2] - vert[0]));
    double      plane_dist = CGLA::dot(p - vert[0], normal);
    CGLA::Vec3d p_plane    = p - normal * plane_dist;

    CGLA::Vec3d b = {
        CGLA::dot(normal, CGLA::cross(vert[1] - p_plane, vert[2] - p_plane)),
        CGLA::dot(normal, CGLA::cross(vert[2] - p_plane, vert[0] - p_plane)),
        CGLA::dot(normal, CGLA::cross(vert[0] - p_plane, vert[1] - p_plane))};

    if (b[0] >= 0 && b[1] >= 0 && b[2] >= 0) {
        CGLA::Vec4d b_d(b, 0);
        b_d /= b_d[0] + b_d[1] + b_d[2];
        b_d[3] = std::abs(plane_dist);
        return b_d;
    }

    CGLA::Vec4d b_d(0);
    double      min_edge_dist = DBL_MAX;
    for (int i = 0; i < 3; ++i) {
        CGLA::Vec3d edge = vert[(i + 1) % 3] - vert[i];
        double t = std::min(1.0, std::max(0.0, CGLA::dot(edge, p - vert[i]) /
                                                   CGLA::dot(edge, edge)));
        double edge_dist = CGLA::length(p - (vert[i] + t * edge));
        if (edge_dist < min_edge_dist) {
            b_d              = CGLA::Vec4d(0);
            b_d[(i + 1) % 3] = t;
            b_d[i]           = 1.0 - t;
            min_edge_dist    = edge_dist;
        }
    }
    b_d[3] = min_edge_dist;
    return b_d;
}

// Micro grid implementation
void Splatter::MicroGrid::push(float d) { dist[idx++] = d; }
void Splatter::MicroGrid::store(Geometry::RGrid<float> &grid,
                                Geometry::RGrid<float> &wgrid) const {
    int i = 0;
    for (CGLA::Vec3i pi : Geometry::Range3D(lower, upper)) {
        grid[pi] = std::max(grid[pi], dist[i]);
        ++i;
    }
}

// ====================================================================== //
// Functions

inline double Splatter::smooth_step(double a, double b, double t) {
    if (t < a) return 0.0;
    if (t > b) return 1.0;
    double x = (t - a) / (b - a);
    return 3 * x * x - 2 * x * x * x;
}

void Splatter::right_trim(std::string &str) {
    while (isspace(str.back())) str.pop_back();
}

std::istream &Splatter::get_multi_line(std::istream &is, std::string &buf) {
    std::getline(is, buf);
    right_trim(buf);
    while (buf.back() == '\\') {
        buf.pop_back();
        std::string continuation;
        std::getline(is, continuation);
        right_trim(continuation);
        buf += continuation;
    }
    return is;
}
