path_builder.hpp

// Copyright (c) TrajoptLib contributors
 
#pragma once
 
#include <stdint.h>
 
#include <cassert>
#include <cstddef>
#include <functional>
#include <span>
#include <utility>
#include <vector>
 
#include "trajopt/constraint/angular_velocity_max_magnitude_constraint.hpp"
#include "trajopt/constraint/lane_constraint.hpp"
#include "trajopt/constraint/linear_acceleration_max_magnitude_constraint.hpp"
#include "trajopt/constraint/linear_velocity_direction_constraint.hpp"
#include "trajopt/constraint/linear_velocity_max_magnitude_constraint.hpp"
#include "trajopt/constraint/point_at_constraint.hpp"
#include "trajopt/constraint/point_point_max_constraint.hpp"
#include "trajopt/constraint/point_point_min_constraint.hpp"
#include "trajopt/constraint/line_point_constraint.hpp"
#include "trajopt/constraint/point_line_region_constraint.hpp"
#include "trajopt/constraint/constraint.hpp"
#include "trajopt/constraint/pose_equality_constraint.hpp"
#include "trajopt/constraint/translation_equality_constraint.hpp"
#include "trajopt/geometry/pose2.hpp"
#include "trajopt/geometry/translation2.hpp"
#include "trajopt/path/path.hpp"
#include "trajopt/util/generate_linear_initial_guess.hpp"
#include "trajopt/util/generate_spline_initial_guess.hpp"
#include "trajopt/util/symbol_exports.hpp"
 
namespace trajopt {

template <typename Drivetrain, typename Solution>
class TRAJOPT_DLLEXPORT PathBuilder {
 public:
  void set_drivetrain(Drivetrain drivetrain) {
    path.drivetrain = std::move(drivetrain);
  }

  void set_bumpers(double front, double left, double right, double back) {
    bumpers.emplace_back(KeepOutRegion{.safety_distance = 0.01,
                                       .points = {{+front, +left},
                                                  {-back, +left},
                                                  {-back, -right},
                                                  {+front, -right}}});
  }

  void set_control_interval_counts(std::vector<size_t>&& counts) {
    control_interval_counts = std::move(counts);
  }

  const std::vector<size_t>& get_control_interval_counts() const {
    return control_interval_counts;
  }

  void sgmt_initial_guess_points(size_t from_index,
                                 std::span<const Pose2d> guess_points) {
    new_wpts(from_index + 1);
    std::vector<Pose2d>& to_initial_guess_points =
        initial_guess_points.at(from_index + 1);
    to_initial_guess_points.insert(to_initial_guess_points.begin(),
                                   guess_points.begin(), guess_points.end());
  }

  void pose_wpt(size_t index, const Pose2d& pose) {
    wpt_constraint(index, PoseEqualityConstraint{pose});
    empty_wpt(index, pose);
  }

  void translation_wpt(size_t index, const Translation2d& translation,
                       double heading_guess = 0.0) {
    wpt_constraint(index, TranslationEqualityConstraint{translation});
    empty_wpt(index, {translation, {heading_guess}});
  }

  void empty_wpt(size_t index, const Pose2d& pose_guess) {
    new_wpts(index);
    initial_guess_points.at(index).back() = pose_guess;
  }

  void wpt_linear_velocity_direction(size_t index, double angle) {
    wpt_constraint(index, LinearVelocityDirectionConstraint{angle});
  }

  void wpt_linear_velocity_max_magnitude(size_t index, double magnitude) {
    wpt_constraint(index, LinearVelocityMaxMagnitudeConstraint{magnitude});
  }

  void wpt_angular_velocity_max_magnitude(size_t index, double magnitude) {
    wpt_constraint(index, AngularVelocityMaxMagnitudeConstraint{magnitude});
  }

  void wpt_linear_acceleration_max_magnitude(size_t index, double magnitude) {
    wpt_constraint(index, LinearAccelerationMaxMagnitudeConstraint{magnitude});
  }

  void wpt_point_at(size_t index, const Translation2d& field_point,
                    double heading_tolerance, bool flip) {
    wpt_constraint(index,
                   PointAtConstraint{field_point, heading_tolerance, flip});
  }

  void wpt_keep_in_circle(size_t index, const Translation2d& center,
                          double radius) {
    for (const auto& bumper : bumpers) {
      for (const auto& point : bumper.points) {
        wpt_constraint(index, PointPointMaxConstraint{point, center, radius});
      }
    }
    wpt_constraint(index, PointPointMaxConstraint{{0.0, 0.0}, center, radius});
  }

  void wpt_keep_in_polygon(size_t index,
                           std::span<Translation2d> field_points) {
    for (size_t i = 0; i < field_points.size(); ++i) {
      size_t j = (i + 1) % field_points.size();
      wpt_constraint(
          index,
          PointLineRegionConstraint{
              {0.0, 0.0}, field_points[i], field_points[j], Side::ABOVE});
      for (const auto& bumper : bumpers) {
        for (const auto& corner : bumper.points) {
          wpt_constraint(
              index, PointLineRegionConstraint{corner, field_points[i],
                                               field_points[j], Side::ABOVE});
        }
      }
    }
  }

  void wpt_keep_in_lane(size_t index, const Translation2d& center_line_start,
                        const Translation2d& center_line_end,
                        double tolerance) {
    wpt_constraint(
        index, LaneConstraint{center_line_start, center_line_end, tolerance});
  }

  void wpt_keep_out_circle(size_t index, const Translation2d& center,
                           double radius) {
    for (const auto& bumper : bumpers) {
      for (size_t i = 0; i < bumper.points.size(); ++i) {
        size_t j = (i + 1) % bumper.points.size();
        wpt_constraint(index, PointPointMinConstraint{bumper.points.at(i),
                                                      center, radius});
        wpt_constraint(
            index, LinePointConstraint{bumper.points.at(i), bumper.points.at(j),
                                       center, radius});
      }
    }
  }

  void sgmt_linear_velocity_direction(size_t from_index, size_t to_index,
                                      double angle) {
    sgmt_constraint(from_index, to_index,
                    LinearVelocityDirectionConstraint{angle});
  }

  void sgmt_linear_velocity_max_magnitude(size_t from_index, size_t to_index,
                                          double magnitude) {
    sgmt_constraint(from_index, to_index,
                    LinearVelocityMaxMagnitudeConstraint{magnitude});
  }

  void sgmt_angular_velocity_max_magnitude(size_t from_index, size_t to_index,
                                           double magnitude) {
    sgmt_constraint(from_index, to_index,
                    AngularVelocityMaxMagnitudeConstraint{magnitude});
  }

  void sgmt_linear_acceleration_max_magnitude(size_t from_index,
                                              size_t to_index,
                                              double magnitude) {
    sgmt_constraint(from_index, to_index,
                    LinearAccelerationMaxMagnitudeConstraint{magnitude});
  }

  void sgmt_point_at(size_t from_index, size_t to_index,
                     Translation2d field_point, double heading_tolerance,
                     bool flip) {
    sgmt_constraint(from_index, to_index,
                    PointAtConstraint{field_point, heading_tolerance, flip});
  }

  void sgmt_keep_in_circle(size_t from_index, size_t to_index,
                           const Translation2d& center, double radius) {
    for (const auto& bumper : bumpers) {
      for (size_t i = 0; i < bumper.points.size(); ++i) {
        sgmt_constraint(
            from_index, to_index,
            PointPointMaxConstraint{bumper.points.at(i), center, radius});
      }
    }
  }

  void sgmt_keep_in_polygon(size_t from_index, size_t to_index,
                            std::span<Translation2d> field_points) {
    for (size_t i = 0; i < field_points.size(); ++i) {
      size_t j = (i + 1) % field_points.size();
      sgmt_constraint(
          from_index, to_index,
          PointLineRegionConstraint{
              {0.0, 0.0}, field_points[i], field_points[j], Side::ABOVE});
      for (const auto& bumper : bumpers) {
        for (const auto& corner : bumper.points) {
          sgmt_constraint(
              from_index, to_index,
              PointLineRegionConstraint{corner, field_points[i],
                                        field_points[j], Side::ABOVE});
        }
      }
    }
  }

  void sgmt_keep_in_lane(size_t from_index, size_t to_index,
                         const Translation2d& center_line_start,
                         const Translation2d& center_line_end,
                         double tolerance) {
    sgmt_constraint(
        from_index, to_index,
        LaneConstraint{center_line_start, center_line_end, tolerance});
  }

  void sgmt_keep_out_circle(size_t from_index, size_t to_index,
                            Translation2d center, double radius) {
    for (const auto& bumper : bumpers) {
      for (size_t i = 0; i < bumper.points.size(); ++i) {
        size_t j = (i + 1) % bumper.points.size();
        sgmt_constraint(
            from_index, to_index,
            PointPointMinConstraint{bumper.points.at(i), center, radius});
        sgmt_constraint(
            from_index, to_index,
            LinePointConstraint{bumper.points.at(i), bumper.points.at(j),
                                center, radius});
      }
    }
  }

  void add_callback(
      const std::function<void(const Solution& solution, int64_t handle)>
          callback) {
    path.callbacks.push_back(callback);
  }

  Path<Drivetrain, Solution>& get_path() { return path; }

  Solution calculate_linear_initial_guess() const {
    return generate_linear_initial_guess<Solution>(initial_guess_points,
                                                   control_interval_counts);
  }

  Solution calculate_spline_initial_guess() const {
    return generate_spline_initial_guess<Solution>(initial_guess_points,
                                                   control_interval_counts);
  }
 
 private:
  struct TRAJOPT_DLLEXPORT KeepOutRegion {
    double safety_distance;

    std::vector<Translation2d> points;
  };

  Path<Drivetrain, Solution> path;

  std::vector<KeepOutRegion> bumpers;

  std::vector<std::vector<Pose2d>> initial_guess_points;

  std::vector<size_t> control_interval_counts;

  void wpt_constraint(size_t index, const Constraint& constraint) {
    new_wpts(index);
    path.waypoints.at(index).waypoint_constraints.push_back(constraint);
  }

  void sgmt_constraint(size_t from_index, size_t to_index,
                       const Constraint& constraint) {
    assert(from_index < to_index);
 
    new_wpts(to_index);
    path.waypoints.at(from_index).waypoint_constraints.push_back(constraint);
    for (size_t index = from_index + 1; index <= to_index; ++index) {
      path.waypoints.at(index).waypoint_constraints.push_back(constraint);
      path.waypoints.at(index).segment_constraints.push_back(constraint);
    }
  }

  void new_wpts(size_t final_index) {
    if (final_index < path.waypoints.size()) {
      return;
    }
 
    for (size_t i = path.waypoints.size(); i <= final_index; ++i) {
      path.waypoints.emplace_back();
      initial_guess_points.emplace_back(std::vector<Pose2d>{{0.0, 0.0, {0.0}}});
      if (i != 0) {
        control_interval_counts.push_back(40);
      }
    }
  }
};
 
}  // namespace trajopt