CoordinateFrames.h

#pragma once
#include <concepts>
#include <type_traits>
#include <variant>
 
#include <frc/geometry/Rotation2d.h>
#include <frc/geometry/Rotation3d.h>
#include <frc/geometry/Transform3d.h>
#include <frc/geometry/Translation3d.h>
#include <frc/kinematics/ChassisSpeeds.h>
 
namespace valor {
 
namespace detail {
 
template <typename... Ts>
struct TypeList {};
 
template <typename T, typename... Ancestors>
struct GetAncestors {
    using type = TypeList<Ancestors..., T>;
};
 
template <typename T, typename... Ancestors>
    requires requires { typename T::Parent; }
struct GetAncestors<T, Ancestors...> {
    using Parent = typename T::Parent;
    using type = typename GetAncestors<Parent, Ancestors..., T>::type;
};
 
template <typename T, typename List>
struct Contains;
 
template <typename T, typename... Us>
struct Contains<T, TypeList<Us...>> : std::disjunction<std::is_same<T, Us>...> {};
 
template <typename ListA, typename ListB>
struct FindCommon;
 
template <typename ListB>
struct FindCommon<TypeList<>, ListB> {
    using type = void;
};
 
template <typename T, typename... Rest, typename ListB>
struct FindCommon<TypeList<T, Rest...>, ListB> {
    using type = std::conditional_t<Contains<T, ListB>::value, T, typename FindCommon<TypeList<Rest...>, ListB>::type>;
};
 
template <typename T>
static constexpr T GetInverse(const T& t) {
    return -t;
}
 
static constexpr frc::Transform2d GetInverse(const frc::Transform2d& t) {
    return frc::Transform2d{-t.Translation(), -t.Rotation()};
}
 
static constexpr frc::Pose2d GetInverse(const frc::Pose2d& t) {
    return frc::Pose2d{}.TransformBy(frc::Transform2d{t, {}});
}
 
static constexpr frc::Pose3d GetInverse(const frc::Pose3d& t) {
    return frc::Pose3d{}.TransformBy(frc::Transform3d{t, {}});
}
 
template <typename T, typename U>
static T ApplyTransform(const T& old, const U& trans) {
    return old + trans;
}
 
constexpr static frc::Translation3d ApplyTransform(const frc::Translation3d& old, const frc::Transform3d& trans) {
    return old.RotateBy(trans.Rotation()) + trans.Translation();
}
 
constexpr static frc::Rotation3d ApplyTransform(const frc::Rotation3d& old, const frc::Transform3d& trans) {
    return old + trans.Rotation();
}
 
constexpr static frc::Pose2d ApplyTransform(const frc::Pose2d& old, const frc::Transform3d& trans) {
    return frc::Pose2d{old.Translation() + trans.Translation().ToTranslation2d(), old.Rotation() + trans.Rotation().ToRotation2d()};
}
 
constexpr static frc::ChassisSpeeds ApplyTransform(const frc::ChassisSpeeds& old, const frc::Rotation2d& trans) {
    return frc::ChassisSpeeds::FromRobotRelativeSpeeds(old, trans);
}
 
constexpr static frc::ChassisSpeeds ApplyTransform(const frc::ChassisSpeeds& old, const frc::Transform3d& trans) {
    return ApplyTransform(old, trans.Rotation().ToRotation2d());
}
 
constexpr static frc::Translation3d ApplyTransform(const frc::Translation3d& old, const frc::Pose3d& trans) {
    return old + trans.Translation();
}
 
constexpr static frc::Translation2d ApplyTransform(const frc::Translation2d& old, const frc::Transform2d& trans) {
    return old + trans.Translation();
}
 
constexpr static frc::ChassisSpeeds ApplyTransform(const frc::ChassisSpeeds& old, const frc::Transform2d& trans) {
    return ApplyTransform(old, trans.Rotation());
}
 
template <class From, class Common, class Geometry, class Context>
Geometry UpTo(const Geometry& current, const Context& ctx) {
    if constexpr (std::is_same_v<From, Common>)
        return current;
    else
        return UpTo<typename From::Parent, Common>(ApplyTransform(current, ctx), ctx);
}
 
template <class From, class Common, class Geometry>
Geometry UpTo(const Geometry& current) {
    if constexpr (std::is_same_v<From, Common>)
        return current;
    else
        return UpTo<typename From::Parent, Common>(ApplyTransform(current, From::GetDefaultTransform()));
}
 
template <class To, class Common, class Geometry, class Context>
Geometry DownTo(const Geometry& current, const Context& ctx) {
    if constexpr (std::is_same_v<To, Common>)
        return current;
    else
        return ApplyTransform(DownTo<typename To::Parent, Common>(current, ctx), GetInverse(ctx));
}
 
template <class To, class Common, class Geometry>
Geometry DownTo(const Geometry& current) {
    if constexpr (std::is_same_v<To, Common>)
        return current;
    else
        return ApplyTransform(DownTo<typename To::Parent, Common>(current), GetInverse(To::GetDefaultTransform()));
}
 
}  // namespace detail
 
template <class Frame, class Geometry>
class FrameWrapper {
   public:
    template <typename... Args>
        requires std::constructible_from<Geometry, Args...>
    explicit FrameWrapper(Args&&... args) : geometry{std::forward<Args>(args)...} {}
 
    FrameWrapper(Geometry g) : geometry{g} {}
 
    template <class FromFrame, typename Context>
    FrameWrapper(FrameWrapper<FromFrame, Geometry> oldFrame, const Context& ctx) : geometry{oldFrame.template To<Frame>(ctx).Get()} {}
 
    template <class FromFrame>
    FrameWrapper(FrameWrapper<FromFrame, Geometry> oldFrame) : geometry{oldFrame.template To<Frame>().Get()} {}
 
    Geometry& Get() { return geometry; }
 
    Geometry* operator->() { return &geometry; }
 
    const Geometry& Get() const { return geometry; }
 
    const Geometry* operator->() const { return &geometry; }
 
    template <class ToFrame, typename Context>
    FrameWrapper<ToFrame, Geometry> To(const Context& ctx) const {
        using Common =
            typename detail::FindCommon<typename detail::GetAncestors<Frame>::type, typename detail::GetAncestors<ToFrame>::type>::type;
        static_assert(!std::is_void_v<Common>, "Frames do not share a common ancestor!");
        return detail::DownTo<ToFrame, Common>(detail::UpTo<Frame, Common>(geometry, ctx), ctx);
    }
 
    template <class ToFrame>
    FrameWrapper<ToFrame, Geometry> To() const {
        using Common =
            typename detail::FindCommon<typename detail::GetAncestors<Frame>::type, typename detail::GetAncestors<ToFrame>::type>::type;
        return detail::DownTo<ToFrame, Common>(detail::UpTo<Frame, Common>(geometry));
    }
 
   private:
    Geometry geometry;
};
 
struct FieldFrame {};
 
struct RobotFrame {
    using Parent = FieldFrame;
 
    static inline frc::Transform2d GetDefaultTransform() {
        frc::Pose2d pose = robotPoseGetter();
        return frc::Transform2d{{}, pose};
    }
 
    static std::function<frc::Pose2d()> robotPoseGetter;
};
 
}  // namespace valor