trajectory.h

#pragma once
 
#include <algorithm>
#include <string>
 
#include "core/math/geometry/pose2d.h"
#include "core/math/kinematics/chassis_accelerations.h"
#include "core/math/kinematics/chassis_velocities.h"
#include "core/math/kinematics/tank_kinematics.h"
#include "core/math/kinematics/tank_wheel_velocities.h"
#include "core/math/units.h"
 
using namespace units::literals;

struct PoseWithCurvature {
  Pose2d pose;
  units::Curvature curvature;
};

struct TrajectoryEvent {
  units::Time time;
  std::string name;
};

class Trajectory {
 public:
  struct State {
    units::Time t = 0_s;                         
    units::Length s = 0_in;                      
    units::Velocity velocity = 0_inps;           
    units::Acceleration acceleration = 0_inps2;  
    Pose2d pose{};                               
    units::Curvature curvature = 0_radpm;        

    State() = default;

    State(units::Time t, units::Velocity velocity, units::Acceleration acceleration, Pose2d pose,
          units::Curvature curvature, units::Length s = 0_in)
        : t(t),
          s(s),
          velocity(velocity),
          acceleration(acceleration),
          pose(pose),
          curvature(curvature) {}

    bool operator==(const State& other) const {
      return t == other.t && velocity == other.velocity && acceleration == other.acceleration &&
             pose == other.pose && curvature == other.curvature;
    }

    State interpolate(State end_value, double i) const {
      const units::Time new_t = t + (end_value.t - t) * i;
      const units::Time delta_t = new_t - t;
 
      if (delta_t < 0_s) {
        return end_value.interpolate(*this, 1.0 - i);
      }
 
      const bool reversing =
          velocity < 0_inps || (units::abs(velocity) < 1E-9_inps && acceleration < 0_inps2);
      const units::Velocity new_v = velocity + (acceleration * delta_t);
      const units::Length new_s =
          (velocity * delta_t + 0.5 * acceleration * delta_t * delta_t) * (reversing ? -1.0 : 1.0);
 
      const units::Length distance = end_value.pose.translation().distance(pose.translation());
      const double interpolation_frac = distance > 1E-9_in ? double(new_s / distance) : i;
 
      Pose2d new_pose = pose.interpolate(end_value.pose, interpolation_frac);
 
      const units::Curvature new_curvature = curvature + (end_value.curvature - curvature) * interpolation_frac;
 
      return State(new_t, new_v, acceleration, new_pose, new_curvature, s + new_s);
    }
 
    ChassisVelocities chassis_velocities() const {
      return ChassisVelocities(velocity, 0_inps, velocity * curvature);
    }
 
    ChassisAccelerations chassis_accelerations() const {
      return ChassisAccelerations(acceleration, 0_inps2, acceleration * curvature);
    }
 
    TankWheelVelocities to_wheel_velocities(const TankKinematics& kin) const {
      return kin.to_wheel_velocities(chassis_velocities());
    }
 
    TankWheelAccelerations to_wheel_accelerations(const TankKinematics& kin) const {
      return kin.to_wheel_accelerations(chassis_accelerations());
    }
  };

  Trajectory() = default;

  explicit Trajectory(std::vector<State> states) : states_(std::move(states)) {
    if (!states_.empty()) {
      total_time_ = states_.back().t;
    }
  }

  bool empty() const { return states_.empty(); }

  units::Time total_time() const { return total_time_; }

  const std::vector<State>& states() const { return states_; }

  const std::vector<TrajectoryEvent>& events() const { return events_; }

  void set_events(std::vector<TrajectoryEvent> events) { events_ = std::move(events); }

  units::Time time_from_distance(units::Length s) const {
    if (states_.empty() || s <= 0_in) return 0_s;
    if (s >= states_.back().s) return total_time_;
 
    std::vector<State>::const_iterator it =
        std::lower_bound(states_.begin() + 1, states_.end(), s,
                         [](const State& st, units::Length val) { return st.s < val; });
 
    const State& s0 = *(it - 1);
    const State& s1 = *it;
 
    if (s1.s == s0.s) return s1.t;
    double alpha = (s - s0.s) / (s1.s - s0.s);
    return s0.t + (s1.t - s0.t) * alpha;
  }

  State sample(units::Time t) const {
    if (states_.empty()) {
      return State{};
    }
 
    if (t <= states_.front().t) {
      return states_.front();
    }
    if (t >= total_time_) {
      return states_.back();
    }
 
    std::vector<State>::const_iterator sample =
        std::lower_bound(states_.cbegin() + 1, states_.cend(), t,
                         [](const State& a, const units::Time& b) { return a.t < b; });
 
    std::vector<State>::const_iterator prev_sample = sample - 1;
 
    if (units::abs(sample->t - prev_sample->t) < 1E-9_s) {
      return *sample;
    }
 
    return prev_sample->interpolate(
        *sample, ((t - prev_sample->t) / (sample->t - prev_sample->t)));
  }

  Trajectory relative_to(const Pose2d& pose) const {
    std::vector<State> new_states = states_;
    for (State& state : new_states) {
      state.pose = state.pose.relative_to(pose);
    }
    return Trajectory(new_states);
  }

  Trajectory mirror(const Pose2d& line_pose) const {
    std::vector<State> new_states;
    new_states.reserve(states_.size());
    for (const State& st : states_) {
      State new_st = st;
      const Pose2d local_pose = st.pose.relative_to(line_pose);
      const Pose2d reflected_local_pose(local_pose.x(), -local_pose.y(), -local_pose.rotation());
      new_st.pose = line_pose.transform_by(reflected_local_pose - Pose2d{});
      new_st.curvature = -st.curvature;
      new_states.push_back(new_st);
    }
    return Trajectory(new_states);
  }

  Trajectory rotate_center(Translation2d center = Translation2d(3_tiles, 3_tiles)) const {
    std::vector<State> new_states;
    new_states.reserve(states_.size());
    for (const State& st : states_) {
      State new_st = st;
      const Translation2d new_translation = 2.0 * center - st.pose.translation();
      new_st.pose = Pose2d(new_translation, st.pose.rotation().opposite());
      // Curvature remains unchanged for 180 degree rotation
      new_states.push_back(new_st);
    }
    return Trajectory(new_states);
  }

  Trajectory reverse() const {
    std::vector<State> new_states;
    new_states.reserve(states_.size());
    for (std::vector<State>::const_reverse_iterator it = states_.rbegin(); it != states_.rend();
         ++it) {
      State st = *it;
      st.t = total_time_ - st.t;
      st.velocity = -st.velocity;
      st.curvature = -st.curvature;
      st.s = states_.back().s - st.s;
      new_states.push_back(st);
    }
    return Trajectory(new_states);
  }

  Trajectory operator+(const Trajectory& other) const {
    if (states_.empty()) return other;
    if (other.empty()) return *this;
 
    std::vector<State> states = states_;
    std::vector<State> other_states = other.states();
    const units::Length total_distance = states_.back().s;
    for (State& other_state : other_states) {
      other_state.t += total_time_;
      other_state.s += total_distance;
    }
 
    states.insert(states.end(), other_states.begin() + 1, other_states.end());
    Trajectory combined(states);
 
    std::vector<TrajectoryEvent> combined_events = events_;
    for (const TrajectoryEvent& ev : other.events()) {
      combined_events.push_back({ev.time + total_time_, ev.name});
    }
    combined.set_events(std::move(combined_events));
 
    return combined;
  }

  Pose2d initial_pose() const { return sample(0_s).pose; }

  bool operator==(const Trajectory& other) const {
    return total_time_ == other.total_time_ && states_ == other.states_;
  }
 
 private:
  std::vector<State> states_;
  std::vector<TrajectoryEvent> events_;
  units::Time total_time_ = 0_s;
};

class TrajectorySampler {
 public:
  explicit TrajectorySampler(const Trajectory& trajectory)
      : trajectory_(&trajectory), cached_index_(0) {}

  void reset() { cached_index_ = 0; }

  Trajectory::State sample(units::Time t) {
    if (!trajectory_ || trajectory_->empty()) {
      return Trajectory::State{};
    }
    const std::vector<Trajectory::State>& states = trajectory_->states();
    if (t <= states.front().t) {
      cached_index_ = 0;
      return states.front();
    }
    if (t >= trajectory_->total_time()) {
      cached_index_ = states.size() - 1;
      return states.back();
    }
 
    while (cached_index_ + 1 < states.size() && states[cached_index_ + 1].t <= t) {
      cached_index_++;
    }
    while (cached_index_ > 0 && states[cached_index_].t > t) {
      cached_index_--;
    }
 
    if (cached_index_ + 1 >= states.size()) {
      return states.back();
    }
 
    const Trajectory::State& prev = states[cached_index_];
    const Trajectory::State& next = states[cached_index_ + 1];
    if (units::abs(next.t - prev.t) < 1E-9_s) {
      return next;
    }
    return prev.interpolate(next, ((t - prev.t) / (next.t - prev.t)));
  }
 
 private:
  const Trajectory* trajectory_ = nullptr;
  size_t cached_index_ = 0;
};