// ========================================================================== //
// 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>

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

CGLA::Vec3d StreamSurfaces::runge_kutta_4(CGLA::Vec3d &p0, double step_size,
                                          CGLA::Vec3d &d0, CGLA::Vec3d &norm0) {
    auto v = [&](const CGLA::Vec3d &p, const CGLA::Vec3d &d0,
                 const CGLA::Vec3d &norm0) {
        // This lambda function describes how the direction vector is
        // affected moving through the domain. The initial direction
        // (d0) is projected onto the estimated plane for the next
        // step.
        CGLA::Vec3d n = frame->interpolate_frame_nearest(p, norm0);
        if (CGLA::length(n) < 0.5) n = norm0;
        CGLA::Vec3d d = d0 - CGLA::dot(d0, n) * n;

        return CGLA::cond_normalize(d);
    };

    // Compute Runge-Kutta approximation of the new point
    CGLA::Vec3d k1 = step_size * v(p0, d0, norm0);
    CGLA::Vec3d k2 = step_size * v(p0 + 0.5 * k1, d0, norm0);
    CGLA::Vec3d k3 = step_size * v(p0 + 0.5 * k2, d0, norm0);
    CGLA::Vec3d k4 = step_size * v(p0 + 1.0 * k3, d0, norm0);

    return p0 + (k1 + 2.0 * k2 + 2.0 * k3 + k4) / 6.0;
}

CGLA::Vec3d StreamSurfaces::estimate_position(CGLA::Vec3d &p0, PDS &pds,
                                              StreamSurface &surf,
                                              const double   R,
                                              const double   step_size) {

    // Get list of neighbours close by the original point.
    std::vector<int> p_neigh_idx = pds.neighbours(p0, R);

    std::vector<CGLA::Vec3d> p_est_list(p_neigh_idx.size());
    double                   w_total = 1.0;
    for (size_t idx = 0; idx < p_neigh_idx.size(); idx++) {
        int         neigh_idx  = p_neigh_idx[idx];
        CGLA::Vec3d p_neigh    = surf.pcl.points(neigh_idx);
        CGLA::Vec3d N_neigh    = surf.pcl.normals(neigh_idx);
        CGLA::Vec3d d_neigh    = CGLA::cond_normalize(p0 - p_neigh);
        double      step_neigh = CGLA::length(p0 - p_neigh);

        double w = step_size / step_neigh;
        w_total += w;

        p_est_list[idx] =
            w * runge_kutta_4(p_neigh, step_neigh, d_neigh, N_neigh);
    }
    for (CGLA::Vec3d p_est_i : p_est_list) p0 += p_est_i;
    p0 /= w_total;

    return p0;
}

// ========================================================================== //
// Surface Creation function for the numerical Pointcloud method

void StreamSurfaces::single_surface(StreamSurface &surf, const double pds_r,
                                    std::mt19937 &rand_dev) {

    // If we do not have a pointcloud.
    if (surf.pcl.size() < 1 || surf.mesh.no_faces() > 1) return;

    // Allocate variables needed.
    std::queue<int>                        Q;
    PDS                                    pds(pds_r, bbox);
    std::uniform_real_distribution<double> random(0.0, 1.0);

    // Insert original points
    for (size_t i = 0; i < surf.pcl.size(); i++) {
        CGLA::Vec3d p0 = surf.pcl.points(i);
        CGLA::Vec3d N0 = surf.pcl.normals(i);

        // If we have an invalid interpolation, skip this.
        if (frame->interpolate_frame_index(p0, N0) == -1) continue;

        Q.push(i);
        pds.add(i, p0);
    }

    // Start generating points.
    while (!Q.empty()) {

        // Extract the front of the queue
        size_t      idx   = Q.front();             // Extract index
        CGLA::Vec3d p0    = surf.pcl.points(idx);  // Origin of sample disk
        CGLA::Vec3d norm0 = surf.pcl.normals(idx); // Normal at sample origin
        Q.pop();                                   // Remove the front

        // Setup the 0-rotational direction
        std::vector<CGLA::Vec3d> fr    = frame->interpolate_frame(p0);
        int                      d_idx = frame->nearest_index(fr, norm0);
        CGLA::Vec3d              dir0  = fr[(d_idx + 1) % dimS];
        if (d_idx == -1) continue;

        // Attempt to place points in the disk 30 times.
        for (int pds_k = 0; pds_k < 30; pds_k++) {

            // ============================================================== //
            // Determine the next point.

            // Generate random distance and angle
            double rng_r = random(rand_dev);
            double rng_a = 2.0 * M_PI * random(rand_dev);

            // Compute step size and initial direction
            double      step_size = (1.0 + rng_r) * pds_r;
            CGLA::Vec3d d0        = Utilities::rotate(dir0, norm0, rng_a);
            CGLA::Vec3d p         = p0 + step_size * d0;

            // Estimate position
            p = estimate_position(p, pds, surf, 5.0 * pds_r, step_size);

            // Update the normal and rotational origin
            CGLA::Vec3d n = frame->interpolate_frame_nearest(p, norm0);

            // ============================================================== //
            // Criterions for keeping the point

            if (!point_is_valid(p, n, pds, surf, 5.0 * pds_r)) continue;

            // ============================================================== //
            // Finalize the iteration
            Q.push(surf.pcl.size());
            pds.add(surf.pcl.size(), p);

            // Save the generated point
            double p_layer = frame->interpolate_layer_nearest(p, n);
            surf.pcl.push_back(p, n, p_layer);
        }
    }
}

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