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

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

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

// Standard libraries
#include <iostream>
#include <queue>
#include <random>
#include <string>
#include <vector>

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

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

// ========================================================================== //
// Smoothing operations

// Combined function to determine if a point is valid.
bool StreamSurfaces::point_is_valid(CGLA::Vec3d &p, CGLA::Vec3d &norm, PDS &pds,
                                    StreamSurface &surf, double dist) {

    // Setup the restrictions
    const double chk_normal = 0.9;
    const double chk_proj   = 0.85;
    const size_t chk_neigh_size =
        std::min((size_t)(2.5 * (dist - pds.get_radius()) / pds.get_radius()),
                 surf.size());

    // If either point or normal are invalid return false
    if (CGLA::length(norm) < 1.0) return false;

    // Ensure position is valid
    if (!inside_domain(p) || pds.get(p) >= 0) return false;

    // Check for thickness constraints
    double p_layer   = frame->interpolate_layer_nearest(p, norm);
    double p_density = frame->interpolate_density(p);
    if (p_layer < threshold_void || p_density > threshold_solid) return false;

    const std::vector<int> neighbour_list = pds.neighbours(p, dist);
    if (!super_sampling && neighbour_list.size() < chk_neigh_size) return false;

    // Compute distance to points in the nearest region.
    for (const int &ni : neighbour_list) {
        const CGLA::Vec3d pp = surf.pcl.points(ni) - p;
        if (pp.length() < pds.get_radius()) return false;
    }

    // Check for normal angles and projection angles
    bool chk = true;
    for (const int &ni : neighbour_list) {
        const CGLA::Vec3d pn = surf.pcl.points(ni);
        const CGLA::Vec3d nn = surf.pcl.normals(ni);

        const CGLA::Vec3d pp   = pn - p;
        const CGLA::Vec3d proj = pn - CGLA::dot(pp, norm) * norm;

        chk = chk                                             //
              && CGLA::dot(nn, norm) > chk_normal             //
              && (proj - p).length() / pp.length() > chk_proj //
            ;
    }

    return chk;
}

void StreamSurfaces::create_surfaces() {

    Utilities::debug_info(__FILE__, __LINE__, "Create surfaces", debug);

    // Setup variables
    std::mt19937 rand_dev;
    int          id_0 = 0, stride = 1;
    bool         stop  = false;
    const double pds_r = super_sampling
                             ? frame->element_size * super_sample_pds_r
                             : frame->element_size * surf_pds_r;

    // Indicator of which surfaces have been generated
    std::vector<int> generated(surfaces.size(), 0);

#if MPI_VERSION > 0

    // Setup MPI options
    int mpi_rank, mpi_size;
    MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
    MPI_Comm_size(MPI_COMM_WORLD, &mpi_size);
    id_0 = mpi_rank, stride = mpi_size;

    surfaces.resize(std::max(mpi_size, (int)surfaces.size()));

    std::vector<MPI_Request> rreq; // Recieving request
    std::vector<MPI_Request> sreq; // Sending request

    for (int i = 0; i < (int)surfaces.size(); i++) {
        // Skip local index and do not communicate on super sampling
        if (i % mpi_size == mpi_rank) continue;

        rreq.push_back(MPI_Request());
        MPI_Irecv(&generated[i], 1, MPI_INT, MPI_ANY_SOURCE, i, MPI_COMM_WORLD,
                  &rreq.back());
    }

    // Prepare probing of recieval
    std::vector<int> flag(rreq.size());

#else
    #ifdef _OPENMP
        // OpenMP settings.
        #pragma omp parallel for private(rand_dev), schedule(dynamic), \
            shared(stop, total_surfaces)
    #endif /* _OPENMP */
#endif     /* MPI_VERSION */
    // Create each grid in turn.
    for (size_t id = id_0; id < surfaces.size(); id += stride) {

        // ================================================================== //
        // Check the stopping criterion

#if MPI_VERSION > 0 /* Probe the MPI retrival */
        MPI_Testall(rreq.size(), rreq.data(), flag.data(), MPI_STATUSES_IGNORE);
#endif

        const int total_surfaces = Utilities::sum(generated);
        stop                     = stop || total_surfaces >= desired_surfaces;

        if (stop || surfaces[id].mesh.no_faces() > 0 ||
            surfaces[id].size() < 1) {
#if MPI_VERSION > 0 /* Share the currently generated surface size */
            for (int rank = 0; rank < mpi_size; rank++) {

                // Skip self and do not communicate on super sampling
                if (rank == mpi_rank) continue;

                sreq.push_back(MPI_Request());
                MPI_Isend(&generated[id], 1, MPI_INT, rank, id, MPI_COMM_WORLD,
                          &sreq.back());
            }
#endif /* MPI_VERSION */
            continue;
        }

        Utilities::debug_info(std::to_string(total_surfaces) + " / " +
                                  std::to_string(desired_surfaces),
                              debug);
        // ================================================================== //
        // Point generation and smoothing

        rand_dev.seed(id); // Seed the random device

        // Run surface creation
        single_surface(surfaces[id], pds_r, rand_dev);

        // Smooth the surface
        smooth_single(surfaces[id], surf_pds_r * frame->element_size);

        // Export checkpoint
        if (debug) export_surf(id, "check_points");

        // ================================================================== //
        // Update the size list and stopping criterions

        generated[id] = surfaces[id].size() > 2 ? 1 : 0;

#if MPI_VERSION > 0 /* Share the currently generated surface size */
        for (int rank = 0; rank < mpi_size; rank++) {
            if (rank == mpi_rank) continue; // Skip self

            sreq.push_back(MPI_Request());
            MPI_Isend(&generated[id], 1, MPI_INT, rank, id, MPI_COMM_WORLD,
                      &sreq.back());
        }
#endif /* MPI_VERSION */
    }

#if MPI_VERSION > 0 /* Ensure all mpi sending is completed */
    MPI_Waitall(sreq.size(), sreq.data(), MPI_STATUSES_IGNORE);
    MPI_Waitall(rreq.size(), rreq.data(), MPI_STATUSES_IGNORE);
#endif

    Utilities::debug_info(__FILE__, __LINE__, "Creation complete", debug);
}

// ========================================================================== //
// Boolean operations

// Return true if point is inside the data domain.
bool StreamSurfaces::inside_domain(CGLA::Vec3d &p) { return bbox.contains(p); }
// Is the current index fixed
bool StreamSurfaces::is_fixed(int idx) {
    if (fixed.size() == 0) return false;
    return std::any_of(fixed.begin(), fixed.end(),
                       [idx](int i) { return i == idx; });
}

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