pose2d.h

#pragma once
 
#include <vector>
 
#include "cevalm.hpp"
#include "core/math/geometry/rotation2d.h"
#include "core/math/geometry/transform2d.h"
#include "core/math/geometry/translation2d.h"
#include "core/math/geometry/twist2d.h"
#include "core/math/units.h"

class Pose2d {
   private:
    Translation2d translation_;
    Rotation2d rotation_;
 
   public:
    constexpr Pose2d() : translation_{Translation2d()}, rotation_{Rotation2d()} {}

    constexpr Pose2d(Translation2d translation, Rotation2d rotation)
        : translation_{translation}, rotation_{rotation} {}

    constexpr Pose2d(units::Length x, units::Length y, Rotation2d rotation)
        : translation_{x, y}, rotation_{rotation} {}

    constexpr Pose2d(units::Length x, units::Length y, double radians)
        : translation_{x, y}, rotation_{radians} {}

    constexpr Pose2d(
            const Eigen::Vector3d& pose_vector, units::Length length_unit,
            units::Angle angle_unit = units::radians
    )
        : translation_{
                  units::Length(pose_vector[0], length_unit),
                  units::Length(pose_vector[1], length_unit)
          },
          rotation_{units::Angle(pose_vector[2], angle_unit)} {}

    constexpr units::Length x() const { return translation_.x(); }

    constexpr double x(units::Length unit) const { return translation_.x(unit); }

    constexpr void set_x(units::Length val) { translation_.set_x(val); }

    constexpr units::Length y() const { return translation_.y(); }

    constexpr double y(units::Length unit) const { return translation_.y(unit); }

    constexpr void set_y(units::Length val) { translation_.set_y(val); }

    constexpr Translation2d translation() const { return translation_; }

    constexpr void set_translation(Translation2d val) { translation_ = val; }

    constexpr Rotation2d rotation() const { return rotation_; }

    constexpr void set_rotation(Rotation2d val) { rotation_ = val; }

    constexpr units::Angle angle() const { return rotation_.angle(); }

    constexpr double angle(units::Angle unit) const { return rotation_.angle(unit); }

    constexpr void set_angle(units::Angle val) { rotation_ = Rotation2d(val); }

    EVec<3> as_vector(units::Length length_unit, units::Angle angle_unit = units::radians) const {
        return {translation_.x().to(length_unit), translation_.y().to(length_unit),
                rotation_.angle().to(angle_unit)};
    }

    constexpr units::Length distance(Translation2d point) const {
        return translation_.distance(point);
    }

    constexpr units::Length distance(const Pose2d& other) const {
        return distance(other.translation_);
    }

    constexpr Rotation2d bearing_to(Translation2d point) const {
        const auto delta = point - translation_;
        if (delta.x().internal() == 0 && delta.y().internal() == 0) {
            return rotation_;
        }
        return delta.theta();
    }

    constexpr units::Angle angle_to(Translation2d point) const {
        return (bearing_to(point) - rotation_).angle();
    }

    constexpr Translation2d local_to_world(Translation2d point) const {
        return translation_ + point.rotate_by(rotation_);
    }

    constexpr Translation2d world_to_local(Translation2d point) const {
        return (point - translation_).rotate_by(-rotation_);
    }

    constexpr Pose2d relative_to(const Pose2d& other) const {
        Transform2d transform = *this - other;
        return Pose2d{transform.translation(), transform.rotation()};
    }

    constexpr Pose2d transform_by(const Transform2d& transform) const {
        return Pose2d{
                translation_ + (transform.translation().rotate_by(rotation_)),
                rotation_ + transform.rotation()
        };
    }

    constexpr Pose2d interpolate(const Pose2d& end, double fraction) const {
        if (fraction <= 0) {
            return *this;
        }
        if (fraction >= 1) {
            return end;
        }
        return {translation_ + (end.translation_ - translation_) * fraction,
                rotation_ + (end.rotation_ - rotation_) * fraction};
    }

    constexpr Pose2d exp(const Twist2d& twist) const {
        const units::Length dx = twist.dx();
        const units::Length dy = twist.dy();
        const double dtheta = twist.dtheta().to(units::radians);
 
        const double sin_theta = cevalm::sin(dtheta);
        const double cos_theta = cevalm::cos(dtheta);
 
        double s, c;
        if (cevalm::abs(dtheta) < 1e-9) {
            s = 1.0 - 1.0 / 6.0 * dtheta * dtheta;
            c = 0.5 * dtheta;
        } else {
            s = sin_theta / dtheta;
            c = (1 - cos_theta) / dtheta;
        }
 
        const Transform2d transform{
                Translation2d{dx * s - dy * c, dx * c + dy * s}, Rotation2d{cos_theta, sin_theta}
        };
 
        return *this + transform;
    }

    constexpr Twist2d log(const Transform2d& transform) const {
        const double dtheta = transform.rotation().radians();
        const double half_dtheta = dtheta / 2.0;
 
        const double cos_minus_one = transform.rotation().f_cos() - 1;
 
        double half_theta_by_tan_of_half_dtheta;
 
        if (cevalm::abs(cos_minus_one) < 1e-9) {
            half_theta_by_tan_of_half_dtheta = 1.0 - 1.0 / 12.0 * dtheta * dtheta;
        } else {
            half_theta_by_tan_of_half_dtheta =
                    -(half_dtheta * transform.rotation().f_sin()) / cos_minus_one;
        }
 
        const Translation2d translation_part =
                transform.translation().rotate_by(
                        {half_theta_by_tan_of_half_dtheta, -half_dtheta}
                ) *
                cevalm::hypot(half_theta_by_tan_of_half_dtheta, half_dtheta);
 
        return Twist2d{
                translation_part.x(), translation_part.y(), units::Angle(dtheta, units::radians)
        };
    }

    static Pose2d wrapped_mean(const std::vector<Pose2d>& list) {
        if (list.size() == 0) {
            return Pose2d{};
        }
 
        units::Length sumx;
        units::Length sumy;
 
        double sum_sin = 0;
        double sum_cos = 0;
 
        for (Pose2d pose : list) {
            sumx += pose.x();
            sumy += pose.y();
 
            sum_sin += pose.rotation_.f_sin();
            sum_cos += pose.rotation_.f_cos();
        }
 
        return Pose2d{
                Translation2d{sumx / list.size(), sumy / list.size()},
                Rotation2d{sum_cos / list.size(), sum_sin / list.size()}
        };
    }

    constexpr Pose2d operator+(const Transform2d& transform) const {
        return Pose2d{
                translation_ + (transform.translation().rotate_by(rotation_)),
                transform.rotation() + rotation_
        };
    }

    constexpr Pose2d& operator+=(const Transform2d& transform) { return *this = *this + transform; }

    constexpr Transform2d operator-(const Pose2d& other) const {
        return Transform2d{
                (translation_ - other.translation_).rotate_by(-other.rotation_),
                rotation_ - other.rotation_
        };
    }

    constexpr bool operator==(const Pose2d& other) const {
        return (translation_ == other.translation_) && (rotation_ == other.rotation_);
    }

    constexpr bool is_near(
            const Pose2d& other, units::Length distance_tolerance = units::Length(1e-6),
            units::Angle angle_tolerance = units::Angle(1e-6)
    ) const {
        return translation_.is_near(other.translation_, distance_tolerance) &&
               rotation_.is_near(other.rotation_, angle_tolerance);
    }
};