PAZ_Math

#ifndef PAZ_MATH
#define PAZ_MATH

#include <cmath>
#include <vector>
#include <ostream>
#include <numeric>
#include <algorithm>

namespace paz
{
    template<typename T>
    using require_iterable = std::enable_if_t<std::is_same<decltype(std::begin(
        std::declval<const T&>())), decltype(std::end(std::declval<const
        T&>()))>::value>;

    template<typename T>
    using require_integral = std::enable_if_t<std::is_integral<T>::value>;

    template<typename T>
    using require_not_integral = std::enable_if_t<!std::is_integral<T>::value>;

    constexpr double NatBase   = 2.71828182845904523536028747135; // M_E
    constexpr double SqrtTwo   = 1.41421356237309504880168872421; // M_SQRT2
    constexpr double SqrtHalf  = 0.70710678118654752440084436210; // M_SQRT1_2
    constexpr double SqrtThree = 1.73205080756887729352744634151; // M_SQRT3
    constexpr double Pi        = 3.14159265358979323846264338328; // M_PI
    constexpr double TwoPi     = 6.28318530717958647692528676656; // 2.*M_PI

    constexpr std::size_t None = std::numeric_limits<std::size_t>::max();

    std::size_t random_seed();
    double randn();
    std::size_t randi(std::size_t n);
    double uniform();
    double uniform(double a, double b);
    std::size_t pmf_rand(const std::vector<double>& probs);
    std::vector<std::size_t> rand_seq(std::size_t length);
    std::size_t poissrnd(double lambda);

    // log(w) -> w
    void normalize_log_weights(std::vector<double>& logWeights);
    // w -> w
    void normalize_weights(std::vector<double>& weights);

    template<typename T0, typename T1, typename T2>
    constexpr auto mix(T0 a, T1 b, T2 k)
    {
        return a + k*(b - a);
    }

    template<typename T>
    constexpr T clamp(T x, T a, T b)
    {
        return std::max(a, std::min(b, x));
    }

    template<typename T, require_integral<T>* = nullptr>
    constexpr T fract(T /* x */)
    {
        return 0;
    }
    template<typename T, require_not_integral<T>* = nullptr>
    T fract(T x)
    {
        return x - std::floor(x);
    }

    template<typename T>
    constexpr T sign(T x)
    {
        return (T{0} < x) - (x < T{0});
    }

    template<typename T, require_not_integral<T>* = nullptr>
    T normalize_angle(T x)
    {
        return fract(x/TwoPi)*TwoPi;
    }

    template<typename T = double, require_not_integral<T>* = nullptr>
    constexpr T inf()
    {
        return std::numeric_limits<T>::infinity();
    }

    template<typename T = double, require_not_integral<T>* = nullptr>
    constexpr T eps()
    {
        return std::numeric_limits<T>::epsilon();
    }

    template<typename T = double, require_not_integral<T>* = nullptr>
    T eps(T x)
    {
        const T x0 = std::abs(x);
        return std::nextafter(x0, inf<T>()) - x0;
    }

    template<typename T = double, require_not_integral<T>* = nullptr>
    constexpr T nan()
    {
        return std::numeric_limits<T>::quiet_NaN();
    }

    class Vec;
    class MatRef;
    class Mat
    {
        friend class MatRef;

        std::vector<double> _vals;
        std::size_t _rows = 0;
        std::size_t _cols = 0;

    public:
        static Mat Constant(std::size_t rows, std::size_t cols, double c);
        static Mat Constant(std::size_t side, double c);
        static Mat Zero(std::size_t rows, std::size_t cols);
        static Mat Zero(std::size_t side);
        static Mat Ones(std::size_t rows, std::size_t cols);
        static Mat Ones(std::size_t side);
        static Mat Identity(std::size_t side);
        static Mat Diag(const MatRef& vals);
        static Mat BlockDiag(const MatRef& a, const MatRef& b);
        static Mat Cross(const MatRef& vals);
        static Mat Hcat(const MatRef& a, const MatRef& b);
        static Mat Vcat(const MatRef& a, const MatRef& b);
        static Mat Randn(std::size_t rows, std::size_t cols);
        static Mat Randn(std::size_t side);

        Mat() = default;
        Mat(std::size_t rows, std::size_t cols);
        Mat(std::size_t side);
        Mat(const Vec& v);
        Mat(const MatRef& m);
        Mat(const std::initializer_list<std::initializer_list<double>>& list);

        double det() const;
        Mat inv() const;
        Mat solve(const Mat& b) const; // x : (*this)*x = b (least squares)
        Mat chol() const;
        Mat cholUpdate(const Mat& m, double a) const;
        Vec eig() const;
        Vec eig(Mat& vecs) const;
        void qr(Mat& q, Mat& r) const;
        void qr(Mat& q, Mat& r, std::vector<std::size_t>& p) const;
        Mat trans() const;
        Vec diag() const;
        Mat rep(std::size_t m, std::size_t n) const;
        double& operator()(std::size_t i, std::size_t j)
        {
            if(i >= rows())
            {
                throw std::runtime_error("Row index out of range.");
            }
            if(j >= cols())
            {
                throw std::runtime_error("Column index out of range.");
            }
            return _vals[i + _rows*j];
        }
        double operator()(std::size_t i, std::size_t j) const
        {
            if(i >= rows())
            {
                throw std::runtime_error("Row index out of range.");
            }
            if(j >= cols())
            {
                throw std::runtime_error("Column index out of range.");
            }
            return _vals[i + _rows*j];
        }
        double& operator()(std::size_t i)
        {
            if(i >= size())
            {
                throw std::runtime_error("Index is out of range.");
            }
            return _vals[i];
        }
        double operator()(std::size_t i) const
        {
            if(i >= size())
            {
                throw std::runtime_error("Index is out of range.");
            }
            return _vals[i];
        }
        std::size_t size() const
        {
            return _vals.size();
        }
        std::size_t rows() const
        {
            return _rows;
        }
        std::size_t cols() const
        {
            return _cols;
        }
        double* data()
        {
            return _vals.data();
        }
        const double* data() const
        {
            return _vals.data();
        }
        bool empty() const
        {
            return _vals.empty();
        }
        auto begin()
        {
            return _vals.begin();
        }
        auto begin() const
        {
            return _vals.begin();
        }
        auto end()
        {
            return _vals.end();
        }
        auto end() const
        {
            return _vals.end();
        }
        double normSq() const;
        double norm() const;
        double sum() const;
        Vec rowSum() const;
        Mat colSum() const;
        double min() const;
        double max() const;
        Mat normalized() const;
        Mat prod(const MatRef& rhs) const; // elementwise
        Mat quot(const MatRef& rhs) const; // elementwise
        Mat& operator*=(const MatRef& rhs);
        Mat operator*(const MatRef& rhs) const;
        Mat& operator+=(const MatRef& rhs);
        Mat operator+(const MatRef& rhs) const;
        Mat& operator-=(const MatRef& rhs);
        Mat operator-(const MatRef& rhs) const;
        Mat& operator*=(double rhs);
        Mat operator*(double rhs) const;
        Mat& operator/=(double rhs);
        Mat operator/(double rhs) const;
        Mat operator-() const;
        double dot(const MatRef& rhs) const;
        Vec cross(const MatRef& rhs) const;
        MatRef block(std::size_t startRow, std::size_t startCol, std::size_t
            numRows, std::size_t numCols) const;
        void setBlock(std::size_t startRow, std::size_t startCol, std::size_t
            numRows, std::size_t numCols, const MatRef& rhs);
        MatRef row(std::size_t m) const;
        void setRow(std::size_t m, const MatRef& rhs);
        MatRef col(std::size_t n) const;
        void setCol(std::size_t n, const MatRef& rhs);
        void resize(std::size_t newRows, std::size_t newCols);
        void resizeRows(std::size_t newRows);
        void resizeCols(std::size_t newCols);
        bool hasNan() const;
        void shuffleCols();
    };

    std::ostream& operator<<(std::ostream& out, const MatRef& x);

    class Vec : public Mat
    {
    public:
        static Vec Constant(std::size_t rows, double c);
        static Vec Zero(std::size_t rows);
        static Vec Ones(std::size_t rows);
        static Vec IdQuat();
        static Vec Randn(std::size_t rows);
        static Vec Cat(const MatRef& a, const MatRef& b);

        Vec() = default;
        Vec(std::size_t rows);
        Vec(const Mat& m);
        Vec(const MatRef& m);
        Vec(const std::initializer_list<std::initializer_list<double>>& list);

        MatRef segment(std::size_t start, std::size_t n) const;
        void setSegment(std::size_t start, std::size_t n, const MatRef& rhs);
        MatRef head(std::size_t n) const;
        void setHead(std::size_t n, const Vec& rhs);
        MatRef tail(std::size_t n) const;
        void setTail(std::size_t n, const Vec& rhs);
        void resize(std::size_t newRows);
    };

    class MatRef
    {
        friend class Mat;

    public:
        struct iterator
        {
            using value_type = double;
            using difference_type = std::ptrdiff_t;
            using reference = const double&;
            using pointer = const double*;
            using iterator_category = std::random_access_iterator_tag;

            pointer ptr;
            difference_type row;
            difference_type origRows;
            difference_type blockRows;

            iterator& operator--()
            {
                --row;
                if(row < 0)
                {
                    ptr += blockRows;
                    ptr -= origRows + 1;
                    row = blockRows - 1;
                }
                else
                {
                    --ptr;
                }
                return *this;
            }
            iterator operator--(int)
            {
                auto temp = *this;
                --(*this);
                return temp;
            }
            iterator& operator++()
            {
                ++row;
                if(row == static_cast<difference_type>(blockRows))
                {
                    ptr += origRows + 1;
                    ptr -= blockRows;
                    row = 0;
                }
                else
                {
                    ++ptr;
                }
                return *this;
            }
            iterator operator++(int)
            {
                auto temp = *this;
                ++(*this);
                return temp;
            }
            iterator& operator-=(difference_type n)
            {
                return *this += -n;
            }
            iterator& operator+=(difference_type n)
            {
                const difference_type temp = row + n;
                difference_type deltaCol = temp/blockRows;
                difference_type deltaRow = n - blockRows*deltaCol;
                if(temp < 0 && temp%blockRows)
                {
                    --deltaCol;
                    deltaRow += blockRows;
                }
                ptr += deltaRow + origRows*deltaCol;
                row += deltaRow;
                return *this;
            }
            iterator operator-(difference_type n) const
            {
                auto temp = *this;
                temp -= n;
                return temp;
            }
            iterator operator+(difference_type n) const
            {
                auto temp = *this;
                temp += n;
                return temp;
            }
            difference_type operator-(const iterator& it) const
            {
                difference_type deltaPtr;
                if(ptr < it.ptr)
                {
                    deltaPtr = -static_cast<difference_type>(it.ptr - ptr);
                }
                else
                {
                    deltaPtr = ptr - it.ptr;
                }
                difference_type deltaCol = deltaPtr/origRows;
                if(deltaPtr < 0 && deltaPtr%origRows)
                {
                    --deltaCol;
                }
                const difference_type deltaRow = row - it.row;
                return deltaRow + blockRows*deltaCol;
            }
            reference operator*() const
            {
                return *ptr;
            }
            reference operator[](difference_type n) const
            {
                return *(*this + n);
            }
            bool operator==(const iterator& it) const
            {
                return ptr == it.ptr;
            }
            bool operator!=(const iterator& it) const
            {
                return ptr != it.ptr;
            }
            bool operator<=(const iterator& it) const
            {
                return ptr <= it.ptr;
            }
            bool operator>=(const iterator& it) const
            {
                return ptr >= it.ptr;
            }
            bool operator<(const iterator& it) const
            {
                return ptr < it.ptr;
            }
            bool operator>(const iterator& it) const
            {
                return ptr > it.ptr;
            }
        };

    private:
        const iterator _begin;
        const std::size_t _origCols;
        const std::size_t _blockCols;

        MatRef(const double* ptr, std::size_t origRows, std::size_t origCols,
            std::size_t blockRows, std::size_t blockCols);

    public:
        MatRef(const Mat& m);

        // All `const` methods of `paz::Mat` except `paz::Mat::data`.
        double det() const;
        Mat inv() const;
        Mat solve(const Mat& b) const; // x : (*this)*x = b (least squares)
        Mat chol() const;
        Mat cholUpdate(const Mat& m, double a) const;
        Vec eig() const;
        Vec eig(Mat& vecs) const;
        void qr(Mat& q, Mat& r) const;
        void qr(Mat& q, Mat& r, std::vector<std::size_t>& p) const;
        Mat trans() const;
        Vec diag() const;
        Mat rep(std::size_t m, std::size_t n) const;
        double operator()(std::size_t i, std::size_t j) const
        {
            if(i >= rows())
            {
                throw std::runtime_error("Row index out of range.");
            }
            if(j >= cols())
            {
                throw std::runtime_error("Column index out of range.");
            }
            return _begin[i + rows()*j];
        }
        double operator()(std::size_t i) const
        {
            if(i >= size())
            {
                throw std::runtime_error("Index is out of range.");
            }
            return _begin[i];
        }
        std::size_t size() const
        {
            return rows()*cols();
        }
        std::size_t rows() const
        {
            return _begin.blockRows;
        }
        std::size_t cols() const
        {
            return _blockCols;
        }
        bool empty() const
        {
            return !rows() || !cols();
        }
        const iterator& begin() const
        {
            return _begin;
        }
        iterator end() const
        {
            return _begin + size();
        }
        double normSq() const;
        double norm() const;
        double sum() const;
        Vec rowSum() const;
        Mat colSum() const;
        double min() const;
        double max() const;
        Mat normalized() const;
        Mat prod(const MatRef& rhs) const; // elementwise
        Mat quot(const MatRef& rhs) const; // elementwise
        Mat operator*(const MatRef& rhs) const;
        Mat operator+(const MatRef& rhs) const;
        Mat operator-(const MatRef& rhs) const;
        Mat operator*(double rhs) const;
        Mat operator/(double rhs) const;
        Mat operator-() const;
        double dot(const MatRef& rhs) const;
        Vec cross(const MatRef& rhs) const;
        MatRef block(std::size_t startRow, std::size_t startCol, std::size_t
            numRows, std::size_t numCols) const;
        MatRef row(std::size_t m) const;
        MatRef col(std::size_t n) const;
        bool hasNan() const;
    };

    inline MatRef::iterator operator+(MatRef::iterator::difference_type n, const
        MatRef::iterator& it)
    {
        return it + n;
    }

    inline Mat operator*(double lhs, const MatRef& rhs)
    {
        return rhs*lhs;
    }

    inline bool operator==(const MatRef& a, const MatRef& b)
    {
        if(a.rows() != b.rows() || a.cols() != b.cols())
        {
            return false;
        }
        return std::equal(a.begin(), a.end(), b.begin());
    }
    inline bool operator!=(const MatRef& a, const MatRef& b)
    {
        return !(a == b);
    }

    Mat to_mat(const MatRef& q);
    Vec to_quat(const MatRef& m);
    Vec qinv(const MatRef& q);
    Mat xi(const MatRef& q);
    Vec qmult(const MatRef& p, const MatRef& q);
    Vec axis_angle(const MatRef& axis, double angle);
    Vec nlerp(const MatRef& p, const MatRef& q, double k);

    template<typename T, require_iterable<T>* = nullptr>
    bool any_true(const T& v)
    {
        return std::any_of(std::begin(v), std::end(v), [](bool x){ return x; });
    }

    template<typename T, require_iterable<T>* = nullptr>
    bool all_true(const T& v)
    {
        return std::all_of(std::begin(v), std::end(v), [](bool x){ return x; });
    }

    template<typename T, require_iterable<T>* = nullptr>
    bool none_true(const T& v)
    {
        return std::none_of(std::begin(v), std::end(v), [](bool x){ return x;
            });
    }

    template<typename T, require_iterable<T>* = nullptr>
    std::size_t num_true(const T& v)
    {
        return std::accumulate(std::begin(v), std::end(v), std::size_t{0});
    }

    double gaussian_pdf(const MatRef& mean, const MatRef& cholCov, const MatRef&
        x);
    double cs_divergence(const MatRef& meanA, const MatRef& cholCovA, const
        MatRef& meanB, const MatRef& cholCovB);
    Vec gmm_rand(const std::vector<double>& weights, const std::vector<Vec>&
        means, const std::vector<Mat>& cholCovs);

    double jv(const MatRef& costMat, std::vector<std::size_t>& rowSols);
    void murty(const MatRef& costMat, std::size_t numBest, std::vector<std::
        vector<std::size_t>>& rowSols, std::vector<double>& costs);

    Mat rot1(double angle);
    Mat rot2(double angle);
    Mat rot3(double angle);
}

#endif