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

// Local libraries
#include "Frames/FrameField.hpp"
#include "PDS/PDS.hpp"
#include "Pointcloud/Pointcloud.hpp"
#include "Utilities/math_operations.hpp"
#include "Utilities/string_operations.hpp"

// External libraries
#include <GEL/CGLA/CGLA.h>
#include <GEL/HMesh/HMesh.h>

// Standard libraries
#include <filesystem>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>

// Parallelisation
#include <omp.h>
#if USE_MPI == 1
    #include <mpi.h>
#endif

// Macros
#ifndef MPI_VERSION
    #define MPI_VERSION 0
#endif

namespace fs = std::filesystem;

// ========================================================================== //
// Data import functions

void StreamSurfaces::save() {
    export_adjacency();
    export_surf();
    export_probes();
    export_inputs();
    frame->export_densities(export_location + "intermediate/");

    // Singularity Handling
    if (!singularity_toggle) return;

    Utilities::debug_info(__FILE__, __LINE__, "Exporting singularities", debug);
    frame->export_singularity(export_location + "intermediate/");
}
// Read the information file to find search_radius and domain sizes
void StreamSurfaces::read_information() {

    // Current information file format.
    // [nelx, nely, nelz, nF, n_pp ]
    std::string   file_name = data_location + data_name + "information.csv";
    std::ifstream information(file_name);
    if (!information.good())
        Utilities::error(__FILE__, __LINE__,
                         "Information file could not be read\n\t" + file_name);

    std::string line;
    getline(information, line);
    std::stringstream        lineStream(line);
    std::string              cell;
    std::vector<std::string> values;
    CGLA::Vec3d              p;

    int no_information = 0;
    while (getline(lineStream, cell, ',')) {
        values.push_back(cell);
        no_information++;
    }
    information.close();
    if (no_information < 4) {
        Utilities::error(__FILE__, __LINE__,
                         "Information file does not contain enough elements.");
    }

    // Read the values from the file.
    int nelx = stoi(values[0]); // Number of elements in X
    int nely = stoi(values[1]); // Number of elements in Y
    int nelz = stoi(values[2]); // Number of elements in Z
    int nels = stoi(values[3]); // Number of frame directions

    // Setup dimensions and bounding box
    dimX = nelx, dimY = nely, dimZ = nelz, dimS = nels;
    std::vector<int> dim = {nelx, nely, nelz};
}
// Read inputs from input file
void StreamSurfaces::read_input() {

    // Prepare the RNG.
    std::mt19937 rand_dev;
    rand_dev.seed(0);
    std::uniform_int_distribution<int>     randi_0_dimS(0, dimS - 1);
    std::uniform_real_distribution<double> rand_0_1(0.0, 1.0);

    surfaces = std::vector<StreamSurface>();

    if (seed_method.compare("grid") == 0) {
        // Generating the inputs from scratch

        // Compute the number of inputs per dimension and spacing between points
        CGLA::Vec3i dims = {seed_grid_amount[0], seed_grid_amount[1],
                            seed_grid_amount[2]};
        CGLA::Vec3d sp   = bbox.diagonal() / CGLA::Vec3d(dims - CGLA::Vec3i(1));

        int i    = 0;
        surfaces = std::vector<StreamSurface>(dims[0] * dims[1] * dims[2]);

        // Populate the grid
        for (CGLA::Vec3i pi : Geometry::Range3D(CGLA::Vec3i(0, 0, 0), dims)) {
            int         d_idx = randi_0_dimS(rand_dev);
            CGLA::Vec3d p     = bbox.min_corner + CGLA::Vec3d(pi) * sp;
            CGLA::Vec3d n     = frame->lookup_frame(p)[d_idx];
            double      l     = frame->lookup_layer(p)[d_idx];

            surfaces[i++].pcl = PointCloud(p, n, l);
        }

    } else if (seed_method.compare("file") == 0) {
        // Reading inputs from file
        std::string seed_file = data_location + data_name + seed_file_name;

        // Open the file
        std::ifstream dataFile(seed_file);
        if (!dataFile.good())
            Utilities::error(__FILE__, __LINE__,
                             "Seed-point file not found: \n\t" + seed_file +
                                 ".");

        // Read the file and fill up the vectors
        std::string line;
        while (getline(dataFile, line)) {

            // Allocation of variables
            std::stringstream   lineStream(line);
            std::string         cell;
            std::vector<double> values;
            CGLA::Vec3d         p, n;

            // Read a line into values
            while (getline(lineStream, cell, ' ')) values.push_back(stod(cell));

            // if we dont have enough values continue
            if (values.size() < 6) continue;

            // Read position of point end check validity
            p = {values[0], values[1], values[2]};
            n = {values[3], values[4], values[5]};

            if (!inside_domain(p))
                Utilities::error(__FILE__, __LINE__,
                                 "seed-point outside domain.");

            // Determine how the values should be interpreted
            if (seed_file_type.compare("index") == 0)
                n = frame->lookup_frame(p)[(int)values[3] - 1];
            else if (seed_file_type.compare("vector") == 0)
                n = frame->lookup_frame_nearest(p, n);
            else
                Utilities::error(__FILE__, __LINE__,
                                 "seed_file_type: " + seed_file_type +
                                     " not defined");
            double l = frame->lookup_layer_nearest(p, n);

            // Save the input
            surfaces.push_back(PointCloud(p, n, l));
        }
        dataFile.close();

    } else if (seed_method.compare("pds") == 0) {

        // Generate a random origin
        CGLA::Vec3d p0 = bbox.min_corner;
        p0 += CGLA::Vec3d(rand_0_1(rand_dev), rand_0_1(rand_dev),
                          rand_0_1(rand_dev)) *
              (bbox.max_corner - bbox.min_corner);

        // Setup a sampler and fill the domain
        Matrix<double> bbox_mat(3, 2);

        bbox_mat.set(bbox.min_corner[0], 0, 0);
        bbox_mat.set(bbox.min_corner[1], 1, 0);
        bbox_mat.set(bbox.min_corner[2], 2, 0);

        bbox_mat.set(bbox.max_corner[0], 0, 1);
        bbox_mat.set(bbox.max_corner[1], 1, 1);
        bbox_mat.set(bbox.max_corner[2], 2, 1);

        PDS pds(seed_pds_r * bbox_diag, bbox_mat);
        pds.generate(p0, rand_dev);

        // Extract a list of seeds from the sampler
        std::vector<CGLA::Vec3d> seed_position = pds.get_points();
        shuffle(seed_position.begin(), seed_position.end(), rand_dev);

        // Construct a list of point clouds from the point list
        for (size_t i = 0; i < seed_position.size(); i++) {
            CGLA::Vec3d p     = seed_position[i];
            int         d_idx = randi_0_dimS(rand_dev);
            CGLA::Vec3d n     = frame->lookup_frame(p)[d_idx];
            double      l     = frame->lookup_layer(p)[d_idx];

            if (CGLA::length(n) < 0.5 || l < threshold_void) continue;

            // Save the input
            surfaces.push_back(PointCloud(p, n, l));
        }

    } else if (seed_method.compare("boundary") == 0) {

        if (!fs::exists(data_location + data_name + "boundary.obj"))
            Utilities::error(__FILE__, __LINE__,
                             "\n\tInput method: " + seed_method +
                                 " no boundary file found.");

        HMesh::Manifold mesh;
        HMesh::obj_load(data_location + data_name + "boundary.obj", mesh);

        HMesh::quadric_simplify(mesh, 1000.0 / (double)mesh.no_vertices());

        surfaces = std::vector<StreamSurface>(2 * mesh.no_vertices());
        int i    = 0;
        for (HMesh::VertexID vid : mesh.vertices()) {
            CGLA::Vec3d p = mesh.pos(vid);
            CGLA::Vec3d n = HMesh::normal(mesh, vid);

            std::vector<CGLA::Vec3d> f  = frame->lookup_frame(p);
            size_t                   id = frame->nearest_index(f, n);

            CGLA::Vec3d n1 = f[(id + 1) % dimS];
            CGLA::Vec3d n2 = f[(id + 2) % dimS];

            surfaces[i++] = StreamSurface(PointCloud(p, n1, 1.0));
            surfaces[i++] = StreamSurface(PointCloud(p, n2, 1.0));
        }

        shuffle(surfaces.begin(), surfaces.end(), rand_dev);

    } else {
        Utilities::error(__FILE__, __LINE__,
                         "\n\tInput method: " + seed_method + " not defined.");
    }

    PointCloud pcl;
    for (StreamSurface surf : surfaces) pcl.push_back(surf.pcl);

    Utilities::valid_directory(export_location + "surfaces_full");
    ply_writer(export_location + "surfaces_full/SeedPoints.ply", pcl);

    if ((int)surfaces.size() < desired_surfaces)
        Utilities::warning(
            __FILE__, __LINE__,
            "Not enough seed points, please modify settings to resolve this");
}
// Read or generate the list of probe points
void StreamSurfaces::read_probe() {

    if (probe_method.compare("grid") == 0) {
        // Compute the spacing required for each dimension to obtain the desired
        // number of points.

        CGLA::Vec3i dims = {probe_grid_max[0], probe_grid_max[1],
                            probe_grid_max[2]};
        dims =
            CGLA::v_min(dims, CGLA::Vec3i(bbox.diagonal() /
                                          (tolerated_error * desired_spacing)));

        probe_points   = std::vector<CGLA::Vec3d>(dims[0] * dims[1] * dims[2]);
        CGLA::Vec3d sp = (bbox.diagonal()) / CGLA::Vec3d(dims - CGLA::Vec3i(1));

        int i = 0;
        // Populate the grid
        for (CGLA::Vec3i pi : Geometry::Range3D(dims))
            probe_points[i++] = bbox.min_corner + CGLA::Vec3d(pi) * sp;

    } else if (probe_method.compare("file") == 0) {

        std::string file_name = data_location + data_name + probe_file_name;
        // Open the file
        std::ifstream dataFile(file_name);
        if (!dataFile.good())
            Utilities::error(__FILE__, __LINE__,
                             "Probe-point file not found: \n\t" + file_name +
                                 ".");

        // Read the file and fill up the vectors
        std::string line;
        while (getline(dataFile, line)) {

            std::stringstream        lineStream(line);
            std::string              cell;
            std::vector<std::string> values;
            CGLA::Vec3d              p;

            while (getline(lineStream, cell, ',')) values.push_back(cell);

            p = {stod(values[0]), stod(values[1]), stod(values[2])};
            if (!inside_domain(p))
                Utilities::error(__FILE__, __LINE__,
                                 "Probe-point outside domain.");

            probe_points.push_back(p);
        }
        dataFile.close();

    } else {
        Utilities::error(__FILE__, __LINE__,
                         "Probe-point Utilities::error:\n\tProbe method: " +
                             probe_method + " not defined.");
    }
}

// ========================================================================== //
// Data import functions

void StreamSurfaces::import_surf(std::string folder) {

    std::string location = export_location + folder;

    if (!std::filesystem::is_directory(location))
        Utilities::error(__FILE__, __LINE__, "Import location do not exist.");

    int idx = 0;
    for (const auto &entry :
         std::filesystem::directory_iterator(export_location)) {
        std::string file = entry.path();
        if (file.find(".ply") == std::string::npos) continue;
        ply_reader(file, surfaces[idx++].pcl);
    }
}
// Remove unselected inputs from the list
void StreamSurfaces::import_selected() {

    std::string intermediate_location = export_location + "intermediate/";
    Utilities::valid_directory(intermediate_location);

    // Read file with optimised weights
    std::ifstream selection_file(intermediate_location +
                                 "weights_computed.csv");
    if (!selection_file.good())
        Utilities::error(__FILE__, __LINE__, "Computed weights not found");

    // Read the file and fill up the vectors
    std::string      line;
    std::vector<int> is_selected;
    while (getline(selection_file, line)) {
        std::stringstream lineStream(line);
        std::string       cell;
        std::string       value;
        getline(lineStream, value);
        is_selected.push_back(stoi(value));
    }
    selection_file.close();

    // Read in the surfaces selected
    std::string surface_path = export_location + "surfaces_full/";

    // Sorted list of files.
    std::set<std::filesystem::path> sorted_by_name;
    for (auto &entry : std::filesystem::directory_iterator(surface_path))
        sorted_by_name.insert(entry.path());

    // Determine the list of existing file indices
    std::vector<std::string> index_list;
    for (auto &file : sorted_by_name) {
        std::string filename = file.filename();

        if (filename.find("point_cloud_") == std::string::npos) continue;

        std::string fileIndex = filename.substr(
            filename.find_last_of('_') + 1,
            filename.find_last_of('.') - filename.find_last_of('_') - 1);

        size_t idx = stoi(fileIndex);
        index_list.resize(std::max(idx + 1, index_list.size()));
        index_list[idx] = fileIndex;
    }

    // Allocate the list of surfaces
    surfaces = std::vector<StreamSurface>(index_list.size());

    // Load the selected surfaces
    for (int idx : is_selected) {
        std::string fileIndex = index_list[idx];

        StreamSurface surf;

        // Point cloud loading
        std::string pclName =
            surface_path + "point_cloud_" + fileIndex + ".ply";
        if (std::filesystem::exists(pclName)) ply_reader(pclName, surf.pcl);

        // Triangle mesh loading
        std::string meshName =
            surface_path + "triangle_mesh_" + fileIndex + ".obj";
        if (std::filesystem::exists(meshName))
            HMesh::obj_load(meshName, surf.mesh);

        // If no points are loaded, continue
        if (surf.pcl.size() < 1) continue;

        // Save the surface to the list
        surfaces[stoi(fileIndex)] = surf;
    }
    desired_surfaces = surfaces.size();
}

// ========================================================================== //
// Data export functions

void StreamSurfaces::export_surf() {
    // This function takes care of the data exporting.
    // Supported extensions:
    //  .ply

    export_surf("surfaces_full");
}
void StreamSurfaces::export_surf(std::string folder) {
    // This function takes care of the data exporting.
    // Supported extensions:
    //  .ply
    int N    = surfaces.size();
    int id_0 = 0, stride = 1;
    Utilities::valid_directory(export_location + folder);

#if MPI_VERSION > 0
    // Setup MPI options
    MPI_Comm_rank(MPI_COMM_WORLD, &id_0);   // Assign first loop index as rank
    MPI_Comm_size(MPI_COMM_WORLD, &stride); // Use size as stride
#else
    #ifdef _OPENMP
        // OpenMP settings.
        #pragma omp parallel for
    #endif /* _OPENMP */
#endif     /* MPI_VERSION */
    for (int id = id_0; id < N; id += stride) export_surf(id, folder);
}
void StreamSurfaces::export_surf(int idx, std::string folder) {
    // This function takes care of the data exporting.
    // Supported extensions:
    //  .ply

    std::string location = export_location + folder + "/";
    Utilities::valid_directory(location);

    if (surfaces[idx].pcl.size() < 2) return;

    // Prepend 0's to ensure valid list
    std::ostringstream str_idx;
    str_idx << std::setw(1 + log10(surfaces.size())) << std::setfill('0')
            << idx;

    std::string fileName = location + "point_cloud_" + str_idx.str() + ".ply";
    ply_writer(fileName, surfaces[idx].pcl);

    if (surfaces[idx].mesh.no_faces() > 0) {
        fileName = location + "triangle_mesh_" + str_idx.str() + ".obj";
        this->obj_save(fileName, surfaces[idx].mesh);
    }
}

void StreamSurfaces::export_adjacency() {

    // Construct output file name and location
    std::string location = export_location + "intermediate/";
    Utilities::valid_directory(location);

#if MPI_VERSION > 0
    // Setup MPI Communication
    int mpi_rank, mpi_size;
    MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
    MPI_Comm_size(MPI_COMM_WORLD, &mpi_size);

    // Determine which processes have any information
    std::vector<size_t> adj_size(mpi_size, 0);
    adj_size[mpi_rank] = indicator_response.size();

    // Let Rank 0 know how many elements each processor have
    std::vector<size_t> send_buff = adj_size;
    MPI_Reduce(send_buff.data(), adj_size.data(), adj_size.size(),
               MPI_UNSIGNED_LONG, MPI_MAX, 0, MPI_COMM_WORLD);

    // Debug printing
    Utilities::debug_info(__FILE__, __LINE__,
                          "Recieving all indicator response information",
                          debug);

    // MPI Recieve information
    if (mpi_rank == 0) {

        // Do our best to preallocate the indicator response list
        indicator_response.reserve(Utilities::sum(adj_size));

        for (int rank = 1; rank < mpi_size; rank++) {

            // If there is no information to recieve, continue
            if (adj_size[rank] == 0) continue;

            std::vector<int> recv(adj_size[rank], 0);
            MPI_Recv(recv.data(), adj_size[rank], MPI_INT, MPI_ANY_SOURCE, rank,
                     MPI_COMM_WORLD, MPI_STATUS_IGNORE);

            for (int v : recv) indicator_response.push_back(v);
        }
    } else { // MPI Send information

        // If this core have the information, send it.
        if (adj_size[mpi_rank] != 0)
            MPI_Send(indicator_response.data(), indicator_response.size(),
                     MPI_INT, 0, mpi_rank, MPI_COMM_WORLD);
    }
    MPI_Barrier(MPI_COMM_WORLD);

    // Only rank 0 continue when all information have been shared
    if (mpi_rank != 0) return;

#endif /* MPI_VERSION */

    // Write indicator response to file
    std::string   fileName = location + "indicator_response.csv";
    std::ofstream file;
    file.open(fileName, std::ios::out);

    for (size_t i = 2; i < indicator_response.size(); i += 3) {
        if (indicator_response[i] == 0) continue;

        // Save the surface ID, probe ID and indicator value
        file << indicator_response[i - 2] << ",";
        file << indicator_response[i - 1] << ",";
        file << indicator_response[i - 0] << std::endl;
    }

    file.close();
}
void StreamSurfaces::export_probes() {

    // Construct output file name and location
    if (probe_points.size() == 0)
        Utilities::error(
            __FILE__, __LINE__,
            "Error in export_probes\n\tAdjancency matrix not computed.");

    std::string location = export_location + "intermediate/";
    Utilities::valid_directory(location);

#if MPI_VERSION > 0
    int mpi_rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
    if (mpi_rank != 0) return;
#endif

    std::string fileName = location + "probes.csv";

    // Write probe list to file
    std::ofstream file;
    file.open(fileName, std::ios::out);

    for (CGLA::Vec3d prb : probe_points)
        file << prb[0] << "," << prb[1] << "," << prb[2] << "\n";

    file.close();
}
void StreamSurfaces::export_inputs() {

    // Construct output file name and location
    std::string location = export_location + "intermediate/";
    Utilities::valid_directory(location);

#if MPI_VERSION > 0
    int mpi_rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
    if (mpi_rank != 0) return;
#endif

    std::string fileName = location + "inputs.csv";

    // Write input list to file
    std::ofstream file;
    file.open(fileName, std::ios::out);

    for (size_t idx = 0; idx < surfaces.size(); idx++) {
        if (surfaces[idx].pcl.size() < 2) continue;

        CGLA::Vec3d p = surfaces[idx].pcl.points(0);
        CGLA::Vec3d n = surfaces[idx].pcl.normals(0);

        file << p[0] << "," << p[1] << "," << p[2] << ",";
        file << n[0] << "," << n[1] << "," << n[2] << "\n";
    }

    file.close();
}

// ========================================================================== //
// Raw type functions

// Function used to export meshes as .obj files
void StreamSurfaces::obj_save(std::string fileName, HMesh::Manifold &mesh) {

    // File and std::string setup
    std::string number =
        fileName.substr(fileName.find_last_of("_"),
                        fileName.find(".obj") - fileName.find_last_of("_"));
    std::string   folder = fileName.substr(0, fileName.find_last_of("/") + 1);
    std::string   stat_file = folder + "stats" + number + ".csv";
    std::ofstream file(stat_file, std::ios::out);

    HMesh::VertexAttributeVector<CGLA::Vec3d> color;
    double                                    err_max = 0.0;

    for (HMesh::VertexID vid : mesh.vertices()) {
        CGLA::Vec3d pos   = mesh.pos(vid);
        CGLA::Vec3d nrm   = cond_normalize(normal(mesh, vid));
        CGLA::Vec3d F     = frame->interpolate_frame_nearest(pos, nrm);
        double      layer = frame->interpolate_layer_nearest(pos, nrm);

        double error = 0.0;
        if (!(layer < threshold_void || length(F) < 1.0 || length(nrm) < 1.0))
            error = 180.0 / M_PI * acos(abs(dot(nrm, F)));

        err_max = err_max < error ? error : err_max;

        color[vid][0] = layer;
        color[vid][1] = layer;
        color[vid][2] = error;

        // Save to statistics file
        file << error << "\n";
    }

    // Scale the error highlight to ensure visibility
    if (debug)
        for (HMesh::VertexID vid : mesh.vertices())
            color[vid][2] = color[vid][2] > 10 ? 1.0 : color[vid][0];
    else
        for (HMesh::VertexID vid : mesh.vertices())
            color[vid][2] = color[vid][0];

    // Export the file
    HMesh::obj_save(fileName, mesh, color);
}

// ========================================================================== //
// Printing functions

void StreamSurfaces::print_surfaces() {

#if MPI_VERSION > 0
    int rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    std::vector<int> recv_sizes(surfaces.size()), send_sizes(surfaces.size());
    for (size_t i = 0; i < surfaces.size(); i++)
        send_sizes[i] = surfaces[i].pcl.size();

    MPI_Reduce(send_sizes.data(), recv_sizes.data(), surfaces.size(), MPI_INT,
               MPI_MAX, 0, MPI_COMM_WORLD);
    if (rank != 0) return;

    for (size_t i = 0; i < surfaces.size(); i++)
        surfaces[i] =
            PointCloud(std::max((int)surfaces[i].pcl.size(), recv_sizes[i]));

#endif

    int N_surf = 0;
    for (size_t id = 0; id < surfaces.size(); id++)
        N_surf += surfaces[id].pcl.size() < 2 ? 0 : 1;

    printf("%s: %d\n", "Surfaces generated", N_surf);
    if (!debug) return;

    // Determine the amount of digits needed
    int Ns_surf = 1;
    int Ns_id   = (int)log10(surfaces.size()) + 2;
    for (auto surf : surfaces)
        Ns_surf = std::max(Ns_surf, 1 + (int)log10(surf.pcl.size()));

    // Print the table of surface sizes
    printf("  %*s   %s\n", (int)log10(surfaces.size()) + 1, "ID", "Points");
    for (size_t id = 0; id < surfaces.size(); id++) {
        if (surfaces[id].pcl.size() < 2) continue;
        printf("  %*d | %*d\n",             // String format
               Ns_id,                       // String ID length
               (int)id,                     // PointCloud index
               Ns_surf,                     // String surf length
               (int)surfaces[id].pcl.size() // Number of points
        );
    }
    printf("\n");
}

// End of File
// ========================================================================== //