spline_base.h

#pragma once
 
#include <vector>
 
#include "core/math/geometry/rotation2d.h"
#include "core/math/geometry/translation2d.h"
#include "core/math/units.h"
 
using namespace units::literals;

struct SplineSample {
  double u = 0.0;              
  double s = 0.0;              
  Translation2d position;      
  Translation2d velocity;      
  Translation2d acceleration;  
  Rotation2d heading;          
  units::Curvature curvature = 0_radpm;  
};

class SplineBase {
 public:
  virtual ~SplineBase() = default;

  virtual Translation2d position(double u) const = 0;

  virtual Translation2d velocity(double u) const = 0;

  virtual Translation2d acceleration(double u) const = 0;

  Rotation2d heading(double u) const {
    const Translation2d vel = velocity(clamp_u(u));
    return vel.norm().to(units::in) > 1e-9 ? vel.theta() : Rotation2d();
  }

  units::Curvature curvature(double u) const {
    const Translation2d vel = velocity(clamp_u(u));
    const Translation2d acc = acceleration(clamp_u(u));
    const units::Length speed = vel.norm();
    if (speed < 1e-4_in) {
      return 0_radpm;
    }
 
    const units::Area cross_product = (vel.x() * acc.y()) - (vel.y() * acc.x());
    return (cross_product * 1_rad) / units::cube(speed);
  }

  double length() const { return arc_lengths_.empty() ? 0.0 : arc_lengths_.back(); }

  double u_from_s(double s) const {
    if (arc_lengths_.empty() || us_.empty()) {
      return 0.0;
    }
    if (s <= 0.0) {
      return 0.0;
    }
    if (s >= length()) {
      return 1.0;
    }
 
    auto upper = std::lower_bound(arc_lengths_.begin(), arc_lengths_.end(), s);
    const size_t idx = static_cast<size_t>(upper - arc_lengths_.begin());
    if (idx == 0) {
      return us_.front();
    }
 
    const double s0 = arc_lengths_[idx - 1];
    const double s1 = arc_lengths_[idx];
    const double u0 = us_[idx - 1];
    const double u1 = us_[idx];
    const double alpha = (s - s0) / std::max(s1 - s0, 1e-9);
    return u0 + ((u1 - u0) * alpha);
  }

  SplineSample sample(double u) const {
    const double clamped_u = clamp_u(u);
    return SplineSample{
        .u = clamped_u,
        .s = s_from_u(clamped_u),
        .position = position(clamped_u),
        .velocity = velocity(clamped_u),
        .acceleration = acceleration(clamped_u),
        .heading = heading(clamped_u),
        .curvature = curvature(clamped_u),
    };
  }

  SplineSample sample_by_s(double s) const {
    SplineSample out = sample(u_from_s(s));
    out.s = clamp_value(s, 0.0, length());
    return out;
  }

  std::vector<Translation2d> sample_by_count(int count) const {
    std::vector<Translation2d> points;
    if (count <= 0) {
      return points;
    }
    if (count == 1) {
      points.push_back(position(0.0));
      return points;
    }
 
    points.reserve(static_cast<size_t>(count));
    for (int i = 0; i < count; ++i) {
      const double u = static_cast<double>(i) / static_cast<double>(count - 1);
      points.push_back(position(u));
    }
    return points;
  }

  std::vector<Translation2d> sample_by_spacing(double ds) const {
    std::vector<Translation2d> points;
    if (ds <= 0.0 || length() <= 0.0) {
      return points;
    }
 
    for (double s = 0.0; s < length(); s += ds) {
      points.push_back(sample_by_s(s).position);
    }
    points.push_back(position(1.0));
    return points;
  }
 
 protected:
  template <size_t N>
  static double eval_poly(const std::array<double, N>& coeffs, double u) {
    double out = 0.0;
    for (size_t i = N; i-- > 0;) {
      out = (out * u) + coeffs[i];
    }
    return out;
  }

  template <size_t N>
  static double eval_poly_derivative(const std::array<double, N>& coeffs, double u) {
    double out = 0.0;
    for (size_t i = N - 1; i-- > 0;) {
      out = (out * u) + (coeffs[i + 1] * static_cast<double>(i + 1));
    }
    return out;
  }

  template <size_t N>
  static double eval_poly_second_derivative(const std::array<double, N>& coeffs, double u) {
    double out = 0.0;
    for (size_t i = N - 2; i-- > 0;) {
      out = (out * u) + (coeffs[i + 2] * static_cast<double>((i + 1) * (i + 2)));
    }
    return out;
  }

  double integrate_segment_length(double u0, double u1) const {
    constexpr double x[5] = {0.0, -0.5384693101056831, 0.5384693101056831, -0.9061798459386640,
                             0.9061798459386640};
    constexpr double w[5] = {0.5688888888888889, 0.4786286704993665, 0.4786286704993665,
                             0.2369268850561891, 0.2369268850561891};
 
    const double half_len = (u1 - u0) * 0.5;
    const double mid = (u0 + u1) * 0.5;
    double sum = 0.0;
    for (int i = 0; i < 5; ++i) {
      const double u = mid + (half_len * x[i]);
      sum += w[i] * velocity(u).norm().to(units::in);
    }
    return half_len * sum;
  }

  void build_arc_table_recursive(double u0, double u1, double s0, double segment_length,
                                 double max_err) {
    double mid_u = (u0 + u1) * 0.5;
    double len1 = integrate_segment_length(u0, mid_u);
    double len2 = integrate_segment_length(mid_u, u1);
 
    double actual_mid_s = s0 + len1;
    double linear_mid_s = s0 + segment_length * 0.5;
 
    double integral_err = units::abs((len1 + len2) - segment_length);
    double interp_err = units::abs(actual_mid_s - linear_mid_s);
 
    if ((integral_err > max_err || interp_err > max_err) && (u1 - u0) > 1e-5) {
      build_arc_table_recursive(u0, mid_u, s0, len1, max_err);
      build_arc_table_recursive(mid_u, u1, actual_mid_s, len2, max_err);
    } else {
      us_.push_back(u1);
      arc_lengths_.push_back(s0 + len1 + len2);
    }
  }

  void build_arc_table(double max_err = 1e-4) {
    us_.clear();
    arc_lengths_.clear();
 
    us_.push_back(0.0);
    arc_lengths_.push_back(0.0);
 
    const int initial_segments = 10;
    double s0 = 0.0;
    for (int i = 0; i < initial_segments; ++i) {
      double u0 = static_cast<double>(i) / initial_segments;
      double u1 = static_cast<double>(i + 1) / initial_segments;
      double len = integrate_segment_length(u0, u1);
      build_arc_table_recursive(u0, u1, s0, len, max_err);
      s0 = arc_lengths_.back();
    }
  }
 
  static double clamp_u(double u) { return clamp_value(u, 0.0, 1.0); }
 
  static double clamp_value(double value, double lo, double hi) {
    if (value < lo) {
      return lo;
    }
    if (value > hi) {
      return hi;
    }
    return value;
  }
 
 private:
  double s_from_u(double u) const {
    if (us_.empty() || arc_lengths_.empty()) {
      return 0.0;
    }
    if (u <= 0.0) {
      return 0.0;
    }
    if (u >= 1.0) {
      return length();
    }
 
    auto upper = std::lower_bound(us_.begin(), us_.end(), u);
    const size_t idx = static_cast<size_t>(upper - us_.begin());
    if (idx == 0) {
      return arc_lengths_.front();
    }
 
    const double u0 = us_[idx - 1];
    const double u1 = us_[idx];
    const double s0 = arc_lengths_[idx - 1];
    const double s1 = arc_lengths_[idx];
    const double alpha = (u - u0) / std::max(u1 - u0, 1e-9);
    return s0 + ((s1 - s0) * alpha);
  }
 
  std::vector<double> us_;
  std::vector<double> arc_lengths_;
};

template <size_t Order>
class HermiteSpline : public SplineBase {
 public:
  Translation2d position(double u) const override {
    const double clamped_u = clamp_u(u);
    return Translation2d(eval_poly(x_, clamped_u) * 1_in, eval_poly(y_, clamped_u) * 1_in);
  }
 
  Translation2d velocity(double u) const override {
    const double clamped_u = clamp_u(u);
    return Translation2d(eval_poly_derivative(x_, clamped_u) * 1_in,
                         eval_poly_derivative(y_, clamped_u) * 1_in);
  }
 
  Translation2d acceleration(double u) const override {
    const double clamped_u = clamp_u(u);
    return Translation2d(eval_poly_second_derivative(x_, clamped_u) * 1_in,
                         eval_poly_second_derivative(y_, clamped_u) * 1_in);
  }
 
 protected:
  std::array<double, Order + 1> x_{};
  std::array<double, Order + 1> y_{};
};