// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// BSpline.hpp — B-spline basis functions, interpolation, regression
//
// Evaluates B-spline basis using De Boor's recursive algorithm.
// Cubic spline interpolation (natural / clamped / not-a-knot).
// B-spline regression (least-squares fitting).

#ifndef SANGI_MATH_INTERPOLATION_BSPLINE_HPP
#define SANGI_MATH_INTERPOLATION_BSPLINE_HPP

#include <math/core/matrix.hpp>
#include <algorithm>
#include <cassert>
#include <cmath>
#include <stdexcept>
#include <vector>

namespace sangi {

    // ================================================================
    // B-spline basis functions
    // ================================================================

    /// Evaluates the B-spline basis function N_{i,p}(t) via De Boor's recursion
    /// @param i      index of the basis function
    /// @param p      degree
    /// @param t      parameter value
    /// @param knots  knot vector
    template<typename T>
    T bsplineBasis(size_t i, int p, T t, const std::vector<T>& knots) {
        if (p == 0) {
            return (t >= knots[i] && t < knots[i + 1]) ? T{1} : T{0};
        }

        T left = T{0}, right = T{0};
        T denom1 = knots[i + p] - knots[i];
        if (denom1 > T{1e-15})
            left = (t - knots[i]) / denom1 * bsplineBasis(i, p - 1, t, knots);

        T denom2 = knots[i + p + 1] - knots[i + 1];
        if (denom2 > T{1e-15})
            right = (knots[i + p + 1] - t) / denom2 * bsplineBasis(i + 1, p - 1, t, knots);

        return left + right;
    }

    /// Evaluates all basis functions at once (n bases: N_{0,p} ... N_{n-1,p})
    template<typename T>
    std::vector<T> bsplineBasisAll(int p, T t, const std::vector<T>& knots) {
        int n = static_cast<int>(knots.size()) - p - 1;
        std::vector<T> result(n);
        for (int i = 0; i < n; ++i)
            result[i] = bsplineBasis(static_cast<size_t>(i), p, t, knots);

        // Right-endpoint handling: when t == knots.back(), the last basis is 1
        if (t >= knots.back() - T{1e-15}) {
            std::fill(result.begin(), result.end(), T{0});
            result.back() = T{1};
        }
        return result;
    }

    // ================================================================
    // Uniform knot-vector generation
    // ================================================================

    /// clamped uniform knot vector (multiplicity p+1 at both ends)
    template<typename T>
    std::vector<T> uniformKnots(int n, int p, T a = T{0}, T b = T{1}) {
        int m = n + p + 1;
        std::vector<T> knots(m);
        for (int i = 0; i <= p; ++i) knots[i] = a;
        for (int i = m - p - 1; i < m; ++i) knots[i] = b;

        int interior = n - p;
        for (int i = 1; i < interior; ++i)
            knots[p + i] = a + (b - a) * static_cast<T>(i) / static_cast<T>(interior);
        return knots;
    }

    // ================================================================
    // De Boor algorithm (curve evaluation)
    // ================================================================

    /// Evaluates a point on a B-spline curve: C(t) = Σ N_{i,p}(t) * P_i
    /// @param controlPoints control points (n of them, each a d-dimensional vector)
    /// @param knots          knot vector
    /// @param p              degree
    /// @param t              parameter value
    template<typename T>
    std::vector<T> deBoor(const std::vector<std::vector<T>>& controlPoints,
                            const std::vector<T>& knots, int p, T t)
    {
        int n = static_cast<int>(controlPoints.size());
        int dim = static_cast<int>(controlPoints[0].size());

        // Find the knot span
        int k = p;
        for (int i = p; i < n; ++i) {
            if (t >= knots[i] && t < knots[i + 1]) { k = i; break; }
        }
        if (t >= knots[n]) k = n - 1;

        // De Boor's triangular table
        std::vector<std::vector<T>> d(p + 1, std::vector<T>(dim));
        for (int j = 0; j <= p; ++j)
            d[j] = controlPoints[k - p + j];

        for (int r = 1; r <= p; ++r) {
            for (int j = p; j >= r; --j) {
                T denom = knots[k + 1 + j - r] - knots[k + 1 + j - p - 1];
                T alpha = (denom > T{1e-15}) ?
                    (t - knots[k + 1 + j - p - 1]) / denom : T{0};
                for (int dd = 0; dd < dim; ++dd)
                    d[j][dd] = (T{1} - alpha) * d[j - 1][dd] + alpha * d[j][dd];
            }
        }
        return d[p];
    }

    // ================================================================
    // Cubic spline interpolation
    // ================================================================

    /// Result of cubic spline interpolation
    template<typename T>
    struct CubicSplineResult {
        std::vector<T> a, b, c, d;  // per-interval coefficients: S_i(x) = a + b(x-x_i) + c(x-x_i)² + d(x-x_i)³
        std::vector<T> x;           // knot points

        /// Evaluates the interpolated value
        T eval(T t) const {
            size_t n = x.size() - 1;
            // Find the interval
            size_t i = 0;
            for (size_t j = 0; j < n; ++j) {
                if (t >= x[j] && t <= x[j + 1]) { i = j; break; }
                if (j == n - 1) i = n - 1;
            }
            if (t <= x[0]) i = 0;
            if (t >= x[n]) i = n - 1;

            T dx = t - x[i];
            return a[i] + b[i] * dx + c[i] * dx * dx + d[i] * dx * dx * dx;
        }
    };

    /// Natural cubic spline interpolation (S''=0 at the endpoints)
    template<typename T>
    CubicSplineResult<T> cubicSpline(const std::vector<T>& xs, const std::vector<T>& ys) {
        size_t n = xs.size() - 1;
        assert(n >= 1 && xs.size() == ys.size());

        std::vector<T> h(n);
        for (size_t i = 0; i < n; ++i) h[i] = xs[i + 1] - xs[i];

        // Tridiagonal system Mc = rhs
        std::vector<T> alpha(n + 1, T{0});
        for (size_t i = 1; i < n; ++i)
            alpha[i] = T{3} / h[i] * (ys[i + 1] - ys[i])
                      - T{3} / h[i - 1] * (ys[i] - ys[i - 1]);

        // Forward elimination
        std::vector<T> l(n + 1, T{1}), mu(n + 1, T{0}), z(n + 1, T{0});
        for (size_t i = 1; i < n; ++i) {
            l[i] = T{2} * (xs[i + 1] - xs[i - 1]) - h[i - 1] * mu[i - 1];
            mu[i] = h[i] / l[i];
            z[i] = (alpha[i] - h[i - 1] * z[i - 1]) / l[i];
        }

        // Back substitution
        std::vector<T> cc(n + 1, T{0}), bb(n), dd(n);
        for (int j = static_cast<int>(n) - 1; j >= 0; --j) {
            cc[j] = z[j] - mu[j] * cc[j + 1];
            bb[j] = (ys[j + 1] - ys[j]) / h[j] - h[j] * (cc[j + 1] + T{2} * cc[j]) / T{3};
            dd[j] = (cc[j + 1] - cc[j]) / (T{3} * h[j]);
        }

        CubicSplineResult<T> result;
        result.x = xs;
        result.a.resize(n);
        result.b = std::move(bb);
        result.c.resize(n);
        result.d = std::move(dd);
        for (size_t i = 0; i < n; ++i) {
            result.a[i] = ys[i];
            result.c[i] = cc[i];
        }
        return result;
    }

    // ================================================================
    // B-spline regression (least-squares fitting)
    // ================================================================

    /// Result of B-spline regression
    template<typename T>
    struct BSplineRegressionResult {
        std::vector<T> coefficients;  // B-spline coefficients
        std::vector<T> knots;         // knot vector
        int degree;                    // degree

        /// Evaluates the predicted value
        T eval(T t) const {
            auto basis = bsplineBasisAll(degree, t, knots);
            T val = T{0};
            for (size_t i = 0; i < coefficients.size(); ++i)
                val += coefficients[i] * basis[i];
            return val;
        }
    };

    /// B-spline regression
    /// @param x       input data (N samples)
    /// @param y       output data (N samples)
    /// @param nBasis  number of B-spline bases
    /// @param degree  degree of the B-spline (typically 3)
    template<typename T>
    BSplineRegressionResult<T> bsplineRegression(
        const std::vector<T>& x, const std::vector<T>& y,
        int nBasis = 10, int degree = 3)
    {
        size_t N = x.size();
        assert(N == y.size());

        T xMin = *std::min_element(x.begin(), x.end());
        T xMax = *std::max_element(x.begin(), x.end());
        auto knots = uniformKnots<T>(nBasis, degree, xMin, xMax);

        // Design matrix B (N × nBasis)
        Matrix<T> B(N, nBasis, T{0});
        for (size_t i = 0; i < N; ++i) {
            auto basis = bsplineBasisAll(degree, x[i], knots);
            for (int j = 0; j < nBasis; ++j)
                B(i, j) = basis[j];
        }

        // Normal equations: B^T B c = B^T y
        Matrix<T> BtB(nBasis, nBasis, T{0});
        std::vector<T> Bty(nBasis, T{0});
        for (int j = 0; j < nBasis; ++j) {
            for (size_t i = 0; i < N; ++i) {
                Bty[j] += B(i, j) * y[i];
                for (int k = 0; k < nBasis; ++k)
                    BtB(j, k) += B(i, j) * B(i, k);
            }
        }

        // Small regularization
        for (int j = 0; j < nBasis; ++j)
            BtB(j, j) += T{1e-10};

        // Gaussian elimination
        Matrix<T> Aug(nBasis, nBasis + 1);
        for (int i = 0; i < nBasis; ++i) {
            for (int j = 0; j < nBasis; ++j) Aug(i, j) = BtB(i, j);
            Aug(i, nBasis) = Bty[i];
        }
        for (int k = 0; k < nBasis; ++k) {
            T pivot = Aug(k, k);
            if (std::abs(pivot) < T{1e-15}) continue;
            for (int j = k; j <= nBasis; ++j) Aug(k, j) /= pivot;
            for (int i = 0; i < nBasis; ++i) {
                if (i == k) continue;
                T f = Aug(i, k);
                for (int j = k; j <= nBasis; ++j) Aug(i, j) -= f * Aug(k, j);
            }
        }

        std::vector<T> coeffs(nBasis);
        for (int i = 0; i < nBasis; ++i) coeffs[i] = Aug(i, nBasis);

        return { std::move(coeffs), std::move(knots), degree };
    }

} // namespace sangi

#endif // SANGI_MATH_INTERPOLATION_BSPLINE_HPP