#include "Flatten.h"
#include <algorithm>
#include <array>
#include <cmath>
#include <memory>
#include <numbers>
#include <utility>
#include <vector>
#include <Eigen/Dense>

// This is an emulation, in spirit at least, of the effect of of the
// Fortran flat4() subroutine. While our implementation differs from
// that of the original, the results should be as good or better.
//
// One key difference other than what's obvious below is that this
// isn't responsible for converting a power-scaled spectrum to dB.
// If you need to do that before sending data here, that's on you;
// std::transform and std::log10 are going to be your friends. We
// do flattening, and only flattening.
//
// Note that this is a functor; it's serially reusable, but it's not
// reentrant. Call it from one thread only. In practical use, that's
// not expected to be a problem, and it allows us to reuse allocated
// memory in a serial manner, rather than requesting it and freeing
// it constantly.

/******************************************************************************/
// Flatten Constants
/******************************************************************************/

namespace {
// Tunable settings; degree of the polynomial used for the baseline
// curve fit, and the percentile of the span at which to sample. In
// general, a 5th degree polynomial and the 10th percentile should
// be optimal.

constexpr auto FLATTEN_DEGREE = 5;
constexpr auto FLATTEN_SAMPLE = 10;

// We're going to do a pairwise Estrin's evaluation of the polynomial
// coefficients, so it's critical that the degree of the polynomial is
// odd, resulting in an even number of coefficients.

static_assert(FLATTEN_DEGREE & 1, "Degree must be odd");
static_assert(FLATTEN_SAMPLE >= 0 && FLATTEN_SAMPLE <= 100,
              "Sample must be a percentage");

// Since we know the degree of the polynomial, and thus the number of
// nodes that we're going to use, we can do all the trigonometry work
// required to calculate the Chebyshev nodes in advance, by computing
// them over the range [0, 1]; we can then scale these at runtime to
// a span of any size by simple multiplication.
//
// Downside to this with C++17 is that std::cos() is not yet constexpr,
// as it is in C++23, so we must provide our own implementation until
// then.

constexpr auto FLATTEN_NODES = []() {
    // Full-range cosine function using symmetries of cos(x).

    constexpr auto cos = [](double x) {
        constexpr auto RAD_360 = std::numbers::pi * 2;
        constexpr auto RAD_180 = std::numbers::pi;
        constexpr auto RAD_90 = std::numbers::pi / 2;

        // Polynomial approximation of cos(x) for x in [0, RAD_90],
        // Accuracy here in theory is 1e-18, but double precision
        // itself is only 1-e16, so within the domain of doubles,
        // this should be extremely accurate.

        constexpr auto cos = [](double x) {
            constexpr std::array coefficients = {
                1.0,                           // Coefficient for x^0
                -0.49999999999999994,          // Coefficient for x^2
                0.041666666666666664,          // Coefficient for x^4
                -0.001388888888888889,         // Coefficient for x^6
                0.000024801587301587,          // Coefficient for x^8
                -0.00000027557319223986,       // Coefficient for x^10
                0.00000000208767569878681,     // Coefficient for x^12
                -0.00000000001147074513875176, // Coefficient for x^14
                0.0000000000000477947733238733 // Coefficient for x^16
            };

            auto const x2 = x * x;
            auto const x4 = x2 * x2;
            auto const x6 = x4 * x2;
            auto const x8 = x4 * x4;
            auto const x10 = x8 * x2;
            auto const x12 = x8 * x4;
            auto const x14 = x12 * x2;
            auto const x16 = x8 * x8;

            return coefficients[0] + coefficients[1] * x2 +
                   coefficients[2] * x4 + coefficients[3] * x6 +
                   coefficients[4] * x8 + coefficients[5] * x10 +
                   coefficients[6] * x12 + coefficients[7] * x14 +
                   coefficients[8] * x16;
        };

        // Reduce x to [0, RAD_360)

        x -= static_cast<long long>(x / RAD_360) * RAD_360;

        // Map x to [0, RAD_180]

        if (x > RAD_180)
            x = RAD_360 - x;

        // Map x to [0, RAD_90] and evaluate the polynomial;
        // flip the sign for angles in the second quadrant.

        return x > RAD_90 ? -cos(RAD_180 - x) : cos(x);
    };

    // Down to the actual business of generating Chebyshev nodes
    // suitable for scaling; once we move to C++20 as the minimum
    // compiler, we can remove the cos() function above and instead
    // call std::cos() here, as it's required to be constexpr in
    // C++20 and above, and presumably it'll be of high quality.

    auto nodes = std::array<double, FLATTEN_DEGREE + 1>{};
    constexpr auto slice = std::numbers::pi / (2.0 * nodes.size());

    for (std::size_t i = 0; i < nodes.size(); ++i) {
        nodes[i] = 0.5 * (1.0 - cos(slice * (2.0 * i + 1)));
    }

    return nodes;
}();
} // namespace

/******************************************************************************/
// Private Implementation
/******************************************************************************/

class Flatten::Impl {
    using Points = Eigen::Matrix<double, FLATTEN_NODES.size(), 2>;
    using Vandermonde =
        Eigen::Matrix<double, FLATTEN_NODES.size(), FLATTEN_NODES.size()>;
    using Coefficients = Eigen::Vector<double, FLATTEN_NODES.size()>;

    Points p;
    Vandermonde V;
    Coefficients c;

    // Polynomial evaluation using Estrin's method, loop is unrolled at
    // compile time. A compiler should emit SIMD instructions from what
    // it sees here when the optimizer is involved, but even without it,
    // we'll likely see fused multiply-add instructions.

    inline auto evaluate(float const x) const {
        return [this]<Eigen::Index... I>(
                   std::size_t const i,
                   std::integer_sequence<Eigen::Index, I...>) {
            auto baseline = 0.0;
            auto exponent = 1.0;

            ((baseline += (c[I * 2] + c[I * 2 + 1] * i) * exponent,
              exponent *= i * i),
             ...);

            return static_cast<float>(baseline);
        }(x, std::make_integer_sequence<Eigen::Index,
                                        Coefficients::SizeAtCompileTime / 2>{});
    }

  public:
    void operator()(float *const data, std::size_t const size) {
        // Loop invariants; sentinel one past the end of the range, and
        // the number of points in each of the arms on either side of a
        // node.

        auto const end = data + size;
        auto const arm = size / (2 * FLATTEN_NODES.size());

        // Collect lower envelope points; use Chebyshev node interpolants
        // to reduce Runge's phenomenon oscillations.

        for (std::size_t i = 0; i < FLATTEN_NODES.size(); ++i) {
            auto const node = size * FLATTEN_NODES[i];
            auto const base = data + static_cast<int>(std::round(node));
            auto span = std::vector<float>(std::clamp(base - arm, data, end),
                                           std::clamp(base + arm, data, end));

            auto const n = span.size() * FLATTEN_SAMPLE / 100;

            std::nth_element(span.begin(), span.begin() + n, span.end());

            p.row(i) << node, span[n];
        }

        // Extract x and y values from points and prepare the Vandermonde
        // matrix, initializing the first column with 1 (x^0); remaining
        // columns are filled with the Schur product.

        Eigen::VectorXd x = p.col(0);
        Eigen::VectorXd y = p.col(1);

        V.col(0).setOnes();
        for (Eigen::Index i = 1; i < V.cols(); ++i) {
            V.col(i) = V.col(i - 1).cwiseProduct(x);
        }

        // Solve the least squares problem for polynomial coefficients;
        // evaluate the polynomial and subtract the baseline.

        c = V.colPivHouseholderQr().solve(y);

        for (std::size_t i = 0; i < size; ++i)
            data[i] -= evaluate(i);
    }
};

/******************************************************************************/
// Public Implementation
/******************************************************************************/

Flatten::Flatten(bool const flatten)
    : m_impl(flatten ? std::make_unique<Impl>() : nullptr) {}

Flatten::~Flatten() = default;

void Flatten::operator()(bool const flatten) {
    m_impl.reset(flatten ? new Impl() : nullptr);
}

void Flatten::operator()(float *const data, std::size_t const size) {
    if (m_impl && size)
        (*m_impl)(data, size);
}

/******************************************************************************/