UnitVector.h

#pragma once
 
#include <span>
#include <tuple>
 
#include <Eigen/Core>
 
#include "valkyrie/util/eigenunit/Common.h"
 
namespace valor {
namespace EigenUnit {

template <typename... UnitTags>
class UnitVector {
   public:
    static constexpr std::size_t Size = sizeof...(UnitTags);
    using StorageType = Eigen::Matrix<double, Size, 1>;
    using UnitTuple = std::tuple<UnitTags...>;

    constexpr UnitVector() { data.setZero(); }

    constexpr explicit UnitVector(const StorageType& d) : data(d) {}

    constexpr explicit UnitVector(StorageType&& d) : data(std::move(d)) {}

    constexpr UnitVector(const UnitVector&) = default;
    constexpr UnitVector(UnitVector&&) = default;
    constexpr UnitVector& operator=(const UnitVector&) = default;
    constexpr UnitVector& operator=(UnitVector&&) = default;

    constexpr UnitVector(units::unit_t<UnitTags>... args) {
        std::size_t i = 0;
        ((data[i++] = args.value()), ...);
    }

    static constexpr UnitVector Zero() {
        StorageType d;
        d.setZero();
        return UnitVector(d);
    }

    static UnitVector FromArray(const std::array<double, Size>& arr) {
        StorageType d = Eigen::Map<const StorageType>(arr.data());
        return UnitVector(d);
    }

    static UnitVector FromSpan(const std::span<double, Size>& s) {
        StorageType d = Eigen::Map<const StorageType>(s.data());
        return UnitVector(d);
    }

    constexpr const StorageType& Raw() const { return data; }

    constexpr StorageType& Raw() { return data; }

    template <std::size_t N>
        requires ValidIndex<N, UnitTags...>
    constexpr auto Get() const {
        using T = std::tuple_element_t<N, UnitTuple>;
        return units::unit_t<T>(data[N]);
    }

    template <std::size_t N>
        requires ValidIndex<N, UnitTags...>
    void Set(units::unit_t<std::tuple_element_t<N, UnitTuple>> val) {
        data[N] = val.value();
    }

    template <typename T = std::tuple_element_t<0, UnitTuple>>
        requires AllSameUnits<UnitTags...>
    [[nodiscard]]
    auto Norm() const {
        return units::unit_t<T>(data.norm());
    }

    auto Normalized() const
        requires AllSameUnits<UnitTags...>
    {
        return data.normalized();
    }

    template <typename... OtherTags>
        requires(sizeof...(UnitTags) == sizeof...(OtherTags)) && AllSameUnits<MultiplyResult<UnitTags, OtherTags>...>
    [[nodiscard]]
    auto Dot(const UnitVector<OtherTags...>& other) const {
        using T1 = std::tuple_element_t<0, UnitTuple>;
        using T2 = std::tuple_element_t<0, std::tuple<OtherTags...>>;
        using ResultUnit = MultiplyResult<T1, T2>;
        return units::unit_t<ResultUnit>(data.dot(other.Raw()));
    }

    template <std::size_t N>
        requires ValidSize<N, UnitTags...>
    constexpr auto Head() const {
        return slice_impl<0>(std::make_index_sequence<N>{});
    }

    template <std::size_t N>
        requires ValidSize<N, UnitTags...>
    constexpr auto Tail() const {
        return slice_impl<Size - N>(std::make_index_sequence<N>{});
    }

    template <std::size_t Start, std::size_t N>
        requires ValidSize<Start + N, UnitTags...>
    constexpr auto Slice() const {
        return slice_impl<Start>(std::make_index_sequence<N>{});
    }

    template <typename... OtherTags>
    [[nodiscard]]
    constexpr UnitVector operator+(const UnitVector<OtherTags...>& other) const
        requires UnitEquality<UnitTuple, std::tuple<OtherTags...>>
    {
        return UnitVector<UnitTags...>(data + other.Raw());
    }

    template <typename... OtherTags>
    [[nodiscard]]
    constexpr UnitVector operator-(const UnitVector<OtherTags...>& other) const
        requires UnitEquality<UnitTuple, std::tuple<OtherTags...>>
    {
        return UnitVector<UnitTags...>(data - other.Raw());
    }

    constexpr UnitVector operator-() const { return UnitVector<UnitTags...>(-data); }

    template <typename... OtherTags>
    constexpr UnitVector& operator+=(const UnitVector<OtherTags...>& other)
        requires UnitEquality<UnitTuple, std::tuple<OtherTags...>>
    {
        data += other.Raw();
        return *this;
    }

    template <typename... OtherTags>
    constexpr UnitVector& operator-=(const UnitVector<OtherTags...>& other)
        requires UnitEquality<UnitTuple, std::tuple<OtherTags...>>
    {
        data -= other.Raw();
        return *this;
    }

    template <typename ScalarTag>
    [[nodiscard]]
    constexpr auto operator*(const units::unit_t<ScalarTag>& scalar) const {
        using ResultType = UnitVector<MultiplyResult<UnitTags, ScalarTag>...>;
        return ResultType(data * scalar.value());
    }

    template <typename ScalarTag>
    [[nodiscard]]
    constexpr auto operator/(const units::unit_t<ScalarTag>& scalar) const {
        using ResultType = UnitVector<DivideResult<UnitTags, ScalarTag>...>;
        return ResultType(data / scalar.value());
    }

    [[nodiscard]]
    constexpr UnitVector operator*(double scalar) const {
        return UnitVector(data * scalar);
    }

    [[nodiscard]]
    constexpr UnitVector operator/(double scalar) const {
        return UnitVector(data / scalar);
    }

    friend std::ostream& operator<<(std::ostream& os, const UnitVector& vec) {
        os << "[";
        vec.print_impl(os, std::make_index_sequence<Size>{});
        os << "]";
        return os;
    }

    template <typename... OtherTags>
    bool operator==(const UnitVector<OtherTags...>& other) const
        requires UnitEquality<UnitTuple, std::tuple<OtherTags...>>
    {
        return data == other.Raw();
    }

    template <typename... OtherTags>
    [[nodiscard]]
    constexpr auto CWiseProduct(const UnitVector<OtherTags...>& other) const {
        return UnitVector<MultiplyResult<UnitTags, OtherTags>...>(data.CWiseProduct(other.Raw()));
    }

    template <typename... OtherTags>
    [[nodiscard]]
    constexpr auto CWiseQuotient(const UnitVector<OtherTags...>& other) const {
        return UnitVector<DivideResult<UnitTags, OtherTags>...>(data.CWiseQuotient(other.Raw()));
    }

    [[nodiscard]]
    constexpr UnitVector CWiseAbs() const {
        return UnitVector(data.CWiseAbs());
    }

    [[nodiscard]]
    constexpr UnitVector CWiseMin(const UnitVector& other) const {
        return UnitVector(data.CWiseMin(other.Raw()));
    }

    [[nodiscard]]
    constexpr UnitVector CWiseMax(const UnitVector& other) const {
        return UnitVector(data.CWiseMax(other.Raw()));
    }

    [[nodiscard]]
    constexpr UnitVector Clamp(const UnitVector& lo, const UnitVector& hi) const
        requires AllSameUnits<UnitTags...>
    {
        return UnitVector(data.CWiseMax(lo.Raw()).CWiseMin(hi.Raw()));
    }

    [[nodiscard]]
    constexpr UnitVector Clamp(units::unit_t<std::tuple_element_t<0, UnitTuple>> lo,
                               units::unit_t<std::tuple_element_t<0, UnitTuple>> hi) const
        requires AllSameUnits<UnitTags...>
    {
        return UnitVector(data.cwiseMax(lo.value()).cwiseMin(hi.value()));
    }

    bool IsApprox(const UnitVector& other,
                  units::unit_t<std::tuple_element_t<0, UnitTuple>> tol = units::unit_t<std::tuple_element_t<0, UnitTuple>>{1e-9}) const
        requires AllSameUnits<UnitTags...>
    {
        return data.isApprox(other.Raw(), tol.value());
    }

    static UnitVector Constant(units::unit_t<typename std::tuple_element_t<0, UnitTuple>> val)
        requires AllSameUnits<UnitTags...>
    {
        return UnitVector(StorageType::Constant(val.value()));
    }

    template <typename T = std::tuple_element_t<0, UnitTuple>>
        requires AllSameUnits<UnitTags...>
    [[nodiscard]]
    auto Sum() const {
        return units::unit_t<T>(data.sum());
    }

    template <typename T = std::tuple_element_t<0, UnitTuple>>
        requires AllSameUnits<UnitTags...>
    [[nodiscard]]
    auto Mean() const {
        return units::unit_t<T>(data.mean());
    }

    constexpr auto X() const
        requires(Size >= 1)
    {
        return this->Get<0>();
    }

    void X(units::unit_t<std::tuple_element_t<Size >= 1 ? 0 : 0, UnitTuple>> val)
        requires(Size >= 1)
    {
        this->Set<0>(val);
    }

    constexpr auto Y() const
        requires(Size >= 2)
    {
        return this->Get<1>();
    }

    void Y(units::unit_t<std::tuple_element_t<Size >= 2 ? 1 : 0, UnitTuple>> val)
        requires(Size >= 2)
    {
        this->Set<1>(val);
    }

    constexpr auto Z() const
        requires(Size >= 3)
    {
        return this->Get<2>();
    }

    void Z(units::unit_t<std::tuple_element_t<Size >= 3 ? 2 : 0, UnitTuple>> val)
        requires(Size >= 3)
    {
        this->Set<2>(val);
    }

    constexpr auto W() const
        requires(Size >= 4)
    {
        return this->Get<3>();
    }

    void W(units::unit_t<std::tuple_element_t<Size >= 4 ? 3 : 0, UnitTuple>> val)
        requires(Size >= 4)
    {
        this->Set<3>(val);
    }

    auto AsDiagonal() const;

    static auto Skew(const UnitVector<UnitTags...>& vec)
        requires(Size == 3) && AllSameUnits<UnitTags...>
    {
        using Mat = UnitMatrix<UnitTuple, UnitTuple>;
        typename Mat::StorageType d;
        d << 0, -vec.data[2], vec.data[1], vec.data[2], 0, -vec.data[0], -vec.data[1], vec.data[0], 0;
        return Mat(std::move(d));
    }

    template <typename ScalarTag>
    [[nodiscard]]
    constexpr UnitVector operator+(const units::unit_t<ScalarTag>& scalar) const
        requires AllSameUnits<UnitTags...> && UnitEquality<ScalarTag, std::tuple_element_t<0, UnitTuple>>
    {
        return UnitVector<UnitTags...>(data.array() + scalar.value());
    }

    template <typename ScalarTag>
    [[nodiscard]]
    constexpr UnitVector operator-(const units::unit_t<ScalarTag>& scalar) const
        requires AllSameUnits<UnitTags...> && UnitEquality<ScalarTag, std::tuple_element_t<0, UnitTuple>>
    {
        return UnitVector<UnitTags...>(data.array() - scalar.value());
    }

    template <typename ScalarTag>
    constexpr UnitVector& operator+=(const units::unit_t<ScalarTag>& scalar)
        requires AllSameUnits<UnitTags...> && UnitEquality<ScalarTag, std::tuple_element_t<0, UnitTuple>>
    {
        data.array() += scalar.value();
        return *this;
    }

    template <typename ScalarTag>
    constexpr UnitVector& operator-=(const units::unit_t<ScalarTag>& scalar)
        requires AllSameUnits<UnitTags...> && UnitEquality<ScalarTag, std::tuple_element_t<0, UnitTuple>>
    {
        data.array() -= scalar.value();
        return *this;
    }

    template <int P, typename T = std::tuple_element_t<0, UnitTuple>>
        requires AllSameUnits<UnitTags...>
    [[nodiscard]]
    auto LPNorm() const {
        return units::unit_t<T>(data.template lpNorm<P>());
    }

    template <typename T = std::tuple_element_t<0, UnitTuple>>
        requires AllSameUnits<UnitTags...>
    [[nodiscard]]
    auto StableNorm() const {
        return units::unit_t<T>(data.stableNorm());
    }
 
    constexpr auto Begin() { return data.begin(); }
 
    constexpr auto End() { return data.end(); }
 
    constexpr auto Begin() const { return data.begin(); }
 
    constexpr auto End() const { return data.end(); }
 
   private:
    template <std::size_t Offset, std::size_t... Is>
    auto slice_impl(std::index_sequence<Is...>) const {
        return UnitVector<std::tuple_element_t<Offset + Is, UnitTuple>...>(data.template segment<sizeof...(Is)>(Offset));
    }
 
    template <std::size_t... Is>
    void print_impl(std::ostream& os, std::index_sequence<Is...>) const {
        ((os << units::unit_t<std::tuple_element_t<Is, UnitTuple>>(data[Is]).value() << (Is == Size - 1 ? "" : ", ")), ...);
    }
 
    StorageType data;
};

template <typename T, std::size_t N>
using VectorN = typename detail::VectorN_Builder<T, std::make_index_sequence<N>>::type;

using UnitQuaternion = VectorN<units::dimensionless::scalar, 4>;

template <typename... T>
    requires(requires { typename T::unit_type; } && ...)
UnitVector(T...) -> UnitVector<typename T::unit_type...>;

template <std::size_t I, typename... UnitTags>
constexpr auto Get(const UnitVector<UnitTags...>& vec) {
    return vec.template Get<I>();
}

template <typename... TagsA, typename... TagsB>
[[nodiscard]]
constexpr auto Cross(const UnitVector<TagsA...>& a, const UnitVector<TagsB...>& b)
    requires ThreeDimensionalVector<TagsA...> && ThreeDimensionalVector<TagsB...>
{
    using T1 = std::tuple_element_t<0, std::tuple<TagsA...>>;
    using T2 = std::tuple_element_t<0, std::tuple<TagsB...>>;
    using ResultUnit = MultiplyResult<T1, T2>;
 
    return UnitVector<ResultUnit, ResultUnit, ResultUnit>(a.Raw().cross(b.Raw()));
}

template <typename ScalarTag, typename... UnitTags>
[[nodiscard]]
constexpr auto operator*(const units::unit_t<ScalarTag>& scalar, const UnitVector<UnitTags...>& vec) {
    return vec * scalar;
}

template <typename... UnitTags>
[[nodiscard]]
constexpr auto operator*(double scalar, const UnitVector<UnitTags...>& vec) {
    return vec * scalar;
}

template <typename ScalarTag, typename... UnitTags>
[[nodiscard]]
constexpr auto operator+(const units::unit_t<ScalarTag>& scalar, const UnitVector<UnitTags...>& vec)
    requires AllSameUnits<UnitTags...> && UnitEquality<ScalarTag, std::tuple_element_t<0, std::tuple<UnitTags...>>>
{
    return vec + scalar;
}

template <typename ScalarTag, typename... UnitTags>
[[nodiscard]]
constexpr auto operator-(const units::unit_t<ScalarTag>& scalar, const UnitVector<UnitTags...>& vec)
    requires AllSameUnits<UnitTags...> && UnitEquality<ScalarTag, std::tuple_element_t<0, std::tuple<UnitTags...>>>
{
    return UnitVector<UnitTags...>(scalar.value() - vec.Raw().array());
}
 
template <typename Derived, typename Base>
struct NamedPhysicalVector : Base {
    using Base::Base;
 
    NamedPhysicalVector(const Base& b) : Base(b.Raw()) {}
};
 
}  // namespace EigenUnit
}  // namespace valor

template <typename... UnitTags>
struct std::tuple_size<valor::EigenUnit::UnitVector<UnitTags...>> : std::integral_constant<std::size_t, sizeof...(UnitTags)> {};

template <std::size_t N, typename... UnitTags>
struct std::tuple_element<N, valor::EigenUnit::UnitVector<UnitTags...>> {
    using type = decltype(std::declval<valor::EigenUnit::UnitVector<UnitTags...>>().template get<N>());
};