UnitMatrix.h

#pragma once
 
#include <tuple>
 
#include <units/base.h>
 
#include "valkyrie/util/eigenunit/Common.h"
#include "valkyrie/util/eigenunit/UnitVector.h"
 
namespace valor {
namespace EigenUnit {

template <typename OutputTagsTuple, typename InputTagsTuple>
class UnitMatrix;

template <typename... OutTags, typename... InTags>
class UnitMatrix<std::tuple<OutTags...>, std::tuple<InTags...>> {
   public:
    using StorageType = Eigen::Matrix<double, sizeof...(OutTags), sizeof...(InTags)>;
    using OutTuple = std::tuple<OutTags...>;
    using InTuple = std::tuple<InTags...>;
 
    static constexpr std::size_t RowCount = sizeof...(OutTags);
    static constexpr std::size_t ColCount = sizeof...(InTags);

    constexpr UnitMatrix() : UnitMatrix(StorageType::Zero()) {}

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

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

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

    static constexpr UnitMatrix Zero() { return UnitMatrix{}; }

    static constexpr UnitMatrix Identity()
        requires MatchingSquareMatrix<UnitMatrix>
    {
        return UnitMatrix(StorageType::Identity());
    }

    static UnitMatrix Constant(units::dimensionless::scalar_t val)
        requires MatchingSquareMatrix<UnitMatrix>
    {
        return UnitMatrix(StorageType::Constant(val.value()));
    }

    template <typename... ColVecs>
        requires(sizeof...(ColVecs) == ColCount) && ColVecsMatchOutput<OutTuple, ColVecs...>
    static constexpr UnitMatrix FromCols(const ColVecs&... cols) {
        StorageType d;
        std::size_t i = 0;
        ((d.col(i) = cols.Raw(), ++i), ...);
        return UnitMatrix(std::move(d));
    }

    template <typename... RowVecs>
        requires(sizeof...(RowVecs) == RowCount) && RowVecsMatchInput<InTuple, RowVecs...>
    static constexpr UnitMatrix FromRows(const RowVecs&... rows) {
        StorageType d;
        std::size_t i = 0;
        ((d.row(i) = rows.Raw().transpose(), ++i), ...);
        return UnitMatrix(std::move(d));
    }

    template <typename... Args>
        requires ValidElemArgs<OutTuple, InTuple, Args...>
    static constexpr UnitMatrix FromElems(Args&&... args) {
        StorageType d;
        std::size_t idx = 0;
        ((d(idx / ColCount, idx % ColCount) = args.value(), ++idx), ...);
        return UnitMatrix(std::move(d));
    }
 
    template <typename OutTuple, typename InTuple, typename... Args>
    auto MakeMatrix(Args&&... args) {
        return UnitMatrix<OutTuple, InTuple>::FromElems(std::forward<Args>(args)...);
    }

    template <typename TL, typename TR, typename BL, typename BR>
        requires ValidBlockLayout<TL, TR, BL, BR, RowCount, ColCount>
    static constexpr UnitMatrix FromBlocks(const TL& tl, const TR& tr, const BL& bl, const BR& br) {
        StorageType d;
        d.template block<TL::RowCount, TL::ColCount>(0, 0) = tl.Raw();
        d.template block<TR::RowCount, TR::ColCount>(0, TL::ColCount) = tr.Raw();
        d.template block<BL::RowCount, BL::ColCount>(TL::RowCount, 0) = bl.Raw();
        d.template block<BR::RowCount, BR::ColCount>(TL::RowCount, TL::ColCount) = br.Raw();
        return UnitMatrix(std::move(d));
    }

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

    constexpr StorageType& Raw() { return data; }

    template <std::size_t RowStart, std::size_t ColStart, std::size_t RowN, std::size_t ColN>
    constexpr auto Block() const {
        using RawSubOutTuple = typename SliceTuple<RowStart, RowN, OutTuple>::type;
        using RawSubInTuple = typename SliceTuple<ColStart, ColN, InTuple>::type;
 
        using Normalized = NormalizeBlockTuples<RawSubOutTuple, RawSubInTuple>;
 
        using SubOutTuple = typename Normalized::NormalizedOut;
        using SubInTuple = typename Normalized::NormalizedIn;
 
        return UnitMatrix<SubOutTuple, SubInTuple>(data.template block<RowN, ColN>(RowStart, ColStart));
    }

    template <std::size_t RowN, std::size_t ColN>
    constexpr auto TopLeft() const {
        return Block<0, 0, RowN, ColN>();
    }

    template <std::size_t RowN, std::size_t ColN>
    constexpr auto TopRight() const {
        return Block<0, ColCount - ColN, RowN, ColN>();
    }

    template <std::size_t RowN, std::size_t ColN>
    constexpr auto BottomLeft() const {
        return Block<RowCount - RowN, 0, RowN, ColN>();
    }

    template <std::size_t RowN, std::size_t ColN>
    constexpr auto BottomRight() const {
        return Block<RowCount - RowN, ColCount - ColN, RowN, ColN>();
    }

    auto Trace() const
        requires SquareMatrix<UnitMatrix>
    {
        using T = std::tuple_element_t<0, OutTuple>;
        return units::unit_t<T>(data.trace());
    }

    template <typename OtherOut, typename OtherIn>
    [[nodiscard]]
    constexpr auto operator+(const UnitMatrix<OtherOut, OtherIn>& other) const
        requires UnitEquality<OutTuple, OtherOut> && UnitEquality<InTuple, OtherIn>
    {
        return UnitMatrix<OutTuple, InTuple>(data + other.Raw());
    }

    template <typename OtherOut, typename OtherIn>
    [[nodiscard]]
    constexpr auto operator-(const UnitMatrix<OtherOut, OtherIn>& other) const
        requires UnitEquality<OutTuple, OtherOut> && UnitEquality<InTuple, OtherIn>
    {
        return UnitMatrix<OutTuple, InTuple>(data - other.Raw());
    }

    template <typename... VecTags>
    [[nodiscard]]
    constexpr auto operator*(const UnitVector<VecTags...>& vec) const
        requires UnitEquality<InTuple, std::tuple<VecTags...>>
    {
        return UnitVector<OutTags...>(data * vec.Raw());
    }

    template <typename ScalarTag>
    [[nodiscard]]
    constexpr auto operator*(const units::unit_t<ScalarTag>& scalar) const {
        using NewOutTuple = std::tuple<MultiplyResult<OutTags, ScalarTag>...>;
        return UnitMatrix<NewOutTuple, InTuple>(data * scalar.value());
    }

    template <typename ScalarTag>
    [[nodiscard]]
    constexpr auto operator/(const units::unit_t<ScalarTag>& scalar) const {
        using NewOutTuple = std::tuple<DivideResult<OutTags, ScalarTag>...>;
        return UnitMatrix<NewOutTuple, InTuple>(data / scalar.value());
    }

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

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

    template <typename OtherOut, typename OtherIn>
    constexpr UnitMatrix& operator+=(const UnitMatrix<OtherOut, OtherIn>& other)
        requires UnitEquality<OutTuple, OtherOut> && UnitEquality<InTuple, OtherIn>
    {
        data += other.Raw();
        return *this;
    }

    template <typename OtherOut, typename OtherIn>
    constexpr UnitMatrix& operator-=(const UnitMatrix<OtherOut, OtherIn>& other)
        requires UnitEquality<OutTuple, OtherOut> && UnitEquality<InTuple, OtherIn>
    {
        data -= other.Raw();
        return *this;
    }

    template <typename OtherOut, typename OtherIn>
    constexpr bool operator==(const UnitMatrix<OtherOut, OtherIn>& other) const
        requires UnitEquality<OutTuple, OtherOut> && UnitEquality<InTuple, OtherIn>
    {
        return data == other.Raw();
    }

    template <std::size_t R>
        requires ValidIndex<R, OutTags...>
    constexpr auto Row() const {
        return UnitVector<InTags...>(data.row(R).Transpose());
    }

    template <std::size_t C>
        requires ValidIndex<C, InTags...>
    constexpr auto Col() const {
        return UnitVector<OutTags...>(data.col(C));
    }

    template <std::size_t R, std::size_t C>
        requires ValidIndex<R, OutTags...> && ValidIndex<C, InTags...>
    auto Get() const {
        using RowT = std::tuple_element_t<R, OutTuple>;
        using ColT = std::tuple_element_t<C, InTuple>;
        using ElemUnit = DivideResult<RowT, ColT>;
        return units::unit_t<ElemUnit>(data(R, C));
    }

    template <std::size_t R, std::size_t C>
        requires ValidIndex<R, OutTags...> && ValidIndex<C, InTags...>
    void Set(units::unit_t<DivideResult<std::tuple_element_t<R, OutTuple>, std::tuple_element_t<C, InTuple>>> val) {
        data(R, C) = val.value();
    }

    [[nodiscard]]
    auto Inverse() const
        requires SquareMatrix<UnitMatrix>
    {
        return UnitMatrix<InTuple, OutTuple>(data.inverse());
    }

    [[nodiscard]]
    constexpr auto Transpose() const {
        return UnitMatrix<InTuple, OutTuple>(data.transpose());
    }

    auto Eigenvalues() const
        requires MatchingSquareMatrix<UnitMatrix>
    {
        return data.eigenvalues();
    }

    template <typename... VecTags>
    [[nodiscard]]
    auto Solve(const UnitVector<VecTags...>& b) const
        requires SquareMatrix<UnitMatrix> && UnitEquality<OutTuple, std::tuple<VecTags...>>
    {
        return UnitVector<InTags...>(data.partialPivLu().solve(b.Raw()));
    }

    template <typename... B_InTags>
    [[nodiscard]]
    auto Solve(const UnitMatrix<OutTuple, std::tuple<B_InTags...>>& B) const
        requires SquareMatrix<UnitMatrix>
    {
        return UnitMatrix<InTuple, std::tuple<B_InTags...>>(data.partialPivLu().solve(B.Raw()));
    }

    template <typename... VecTags>
    [[nodiscard]]
    auto SolveQR(const UnitVector<VecTags...>& b) const
        requires UnitEquality<OutTuple, std::tuple<VecTags...>>
    {
        return UnitVector<InTags...>(data.colPivHouseholderQr().solve(b.Raw()));
    }

    template <typename... B_InTags>
    [[nodiscard]]
    auto SolveQR(const UnitMatrix<OutTuple, std::tuple<B_InTags...>>& B) const
        requires SquareMatrix<UnitMatrix>
    {
        return UnitMatrix<InTuple, std::tuple<B_InTags...>>(data.colPivHouseholderQr().solve(B.Raw()));
    }

    template <typename... OtherOutTags, typename... OtherInTags>
        requires((sizeof...(OtherOutTags) == RowCount) && (sizeof...(OtherInTags) == ColCount))
    [[nodiscard]]
    constexpr auto CWiseProduct(const UnitMatrix<std::tuple<OtherOutTags...>, std::tuple<OtherInTags...>>& other) const {
        using ResultOut = std::tuple<MultiplyResult<OutTags, OtherOutTags>...>;
        using ResultIn = std::tuple<MultiplyResult<InTags, OtherInTags>...>;
        return UnitMatrix<ResultOut, ResultIn>(data.cwiseProduct(other.Raw()));
    }

    bool IsApprox(const UnitMatrix& other, double tol) const
        requires(!MatchingSquareMatrix<UnitMatrix>)
    {
        return data.isApprox(other.Raw(), tol);
    }

    bool IsApprox(const UnitMatrix& other, units::dimensionless::scalar_t tol = units::dimensionless::scalar_t{1e-9}) const
        requires MatchingSquareMatrix<UnitMatrix>
    {
        return data.IsApprox(other.Raw(), tol.value());
    }

    friend std::ostream& operator<<(std::ostream& os, const UnitMatrix& mat) {
        os << "UnitMatrix (\n";
        mat.print_all_rows(os, std::make_index_sequence<RowCount>{});
        os << ")";
        return os;
    }
 
   private:
    template <std::size_t... Rows>
    void print_all_rows(std::ostream& os, std::index_sequence<Rows...>) const {
        ((os << "  [", print_row<Rows>(os, std::make_index_sequence<ColCount>{}), os << "]\n"), ...);
    }
 
    template <std::size_t Row, std::size_t... Cols>
    void print_row(std::ostream& os, std::index_sequence<Cols...>) const {
        ((os << units::unit_t<DivideResult<std::tuple_element_t<Row, OutTuple>, std::tuple_element_t<Cols, InTuple>>>(data(Row, Cols))
             << (Cols == ColCount - 1 ? "" : ", ")),
         ...);
    }
 
    StorageType data;
};

template <typename... OutA, typename... InA, typename... OutB, typename... InB>
[[nodiscard]]
constexpr auto operator*(const UnitMatrix<std::tuple<OutA...>, std::tuple<InA...>>& A,
                         const UnitMatrix<std::tuple<OutB...>, std::tuple<InB...>>& B)
    requires UnitEquality<std::tuple<InA...>, std::tuple<OutB...>>
{
    return UnitMatrix<std::tuple<OutA...>, std::tuple<InB...>>(A.Raw() * B.Raw());
}

template <typename... OutTags, typename... InTags, typename... VecTags>
[[nodiscard]]
constexpr auto operator*(const UnitMatrix<std::tuple<OutTags...>, std::tuple<InTags...>>& mat, const UnitVector<VecTags...>& vec)
    requires std::is_same_v<std::tuple<OutTags...>, typename ScalarTuple_Builder<sizeof...(OutTags)>::type> &&
             std::is_same_v<std::tuple<InTags...>, typename ScalarTuple_Builder<sizeof...(InTags)>::type> &&
             (sizeof...(InTags) == sizeof...(VecTags))
{
    return UnitVector<VecTags...>(mat.Raw() * vec.Raw());
}
 
template <typename ScalarTag, typename OutTuple, typename InTuple>
[[nodiscard]]
constexpr auto operator*(const units::unit_t<ScalarTag>& scalar, const UnitMatrix<OutTuple, InTuple>& mat) {
    return mat * scalar;
}
 
template <typename OutTuple, typename InTuple>
[[nodiscard]]
constexpr auto operator*(double scalar, const UnitMatrix<OutTuple, InTuple>& mat) {
    return mat * scalar;
}

template <typename VecOut, typename VecIn>
using PhysicalMatrix = UnitMatrix<typename VecOut::UnitTuple, typename VecIn::UnitTuple>;
 
template <typename Derived, typename Base>
struct NamedPhysicalMatrix : Base {
    using Base::Base;
 
    NamedPhysicalMatrix(const Base& base) : Base(base.Raw()) {}
 
    Derived Inverse() const { return Derived(Base::Inverse()); }
 
    template <std::size_t RowStart, std::size_t ColStart, std::size_t RowN, std::size_t ColN>
    auto Block() const {
        return Base::template Block<RowStart, ColStart, RowN, ColN>();
    }
};
 
// Implementation of UnitVector::asDiagonal()
template <typename... UnitTags>
auto UnitVector<UnitTags...>::AsDiagonal() const {
    using Dimensionless = units::dimensionless::scalar_t;
    using DimensionlessTuple = typename detail::VectorN_Builder<Dimensionless, std::make_index_sequence<Size>>::type::UnitTuple;
    return UnitMatrix<UnitTuple, DimensionlessTuple>(data.asDiagonal());
}
 
}  // namespace EigenUnit
}  // namespace valor