/*
 * Copyright (c) 1999-2021 Hartmut Seichter
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 *
 */
#ifndef PW_CORE_MATRIX_HPP
#define PW_CORE_MATRIX_HPP

#include <pw/core/globals.hpp>
#include <pw/core/math.hpp>
#include <pw/core/vector.hpp>

#include <type_traits>
#include <utility>

namespace pw {

template <typename Scalar, std::size_t Rows, std::size_t Cols>
struct matrix final {

    static_assert(Rows > 0 && Cols > 0, "Undefined matrix type.");

    static constexpr unsigned int diag_size{std::min(Rows, Cols)};
    static constexpr unsigned int rows{Rows};
    static constexpr unsigned int cols{Cols};
    static constexpr bool is_square{Rows == Cols};

    using size_type      = std::common_type_t<decltype(Rows), decltype(Cols)>;
    using row_type       = vector<Scalar, Cols>;
    using column_type    = vector<Scalar, Rows>;
    using diag_type      = vector<Scalar, diag_size>;
    using transpose_type = matrix<Scalar, Cols, Rows>;
    using minor_type     = matrix<Scalar, Cols - 1, Rows - 1>;

    vector<vector<Scalar, Cols>, Rows> m_{};

    auto&& data(this auto&& self) {
        return std::forward<decltype(self)>(self).m_[0].data();
    }

    auto&& operator[](this auto&& self, int r) {
        return std::forward<decltype(self)>(self).m_[r];
    }

    //
    // part of tuple-protocol
    //
    template <size_type idx> auto get(this auto&& self) -> decltype(auto) {
        static_assert(idx < Rows, "Out of bounds access to a member.");
        return std::forward_like<decltype(self)>(self.m_[idx]);
    }

    constexpr auto diagonal() const noexcept -> diag_type {
        return [this]<std::size_t... Is>(std::index_sequence<Is...>) {
            return diag_type{(*this)[Is][Is]...};
        }(std::make_index_sequence<diag_size>{});
    }

    constexpr auto trace() const noexcept -> Scalar {
        static_assert(Rows == Cols,
                      "trace is only defined for square matrices");
        return [this, sum = Scalar{0}]<std::size_t... Is>(
                   std::index_sequence<Is...>) mutable {
            return ((*this)[Is][Is] + ... + sum);
        }(std::make_index_sequence<diag_size>{});
    }

    constexpr auto column(const unsigned int& c) const noexcept -> column_type {
        return [this, &c]<std::size_t... Rs>(std::index_sequence<Rs...>) {
            return column_type{(*this)[Rs][c]...};
        }(std::make_index_sequence<Cols>{});
    }

    constexpr auto transposed() const noexcept -> transpose_type {
        return [this]<std::size_t... Cs>(std::index_sequence<Cs...>) {
            return transpose_type{(*this).column(Cs)...};
        }(std::make_index_sequence<Cols>{});
    }

    static constexpr auto all(const Scalar& value) {
        return [&value]<std::size_t... Rs>(std::index_sequence<Rs...>) {
            return matrix{vector<Scalar, Cols>::all((value) + Rs - Rs)...};
        }(std::make_index_sequence<Rows>{});
    }

    static constexpr auto identity() noexcept -> matrix {
        static_assert(Rows == Cols,
                      "identity only defined for square matrices");
        return []<std::size_t... Rs>(std::index_sequence<Rs...>) {
            return matrix{vector<Scalar, Cols>::basis(Rs)...};
        }(std::make_index_sequence<Rows>{});
    }

    static constexpr auto from_diagonal(const diag_type& d) noexcept -> matrix {
        static_assert(
            Rows == Cols,
            "creating from diagonal vector only defined for square matrices");
        return [&d]<std::size_t... Rs>(std::index_sequence<Rs...>) {
            return matrix{vector<Scalar, Cols>::basis(Rs) * d[Rs]...};
        }(std::make_index_sequence<Rows>{});
    }

    constexpr auto minor(std::unsigned_integral auto r0,
                         std::unsigned_integral auto c0) const noexcept {
        static_assert(Rows > 1 && Cols > 1, "cannot create minor matrix");
        return [this, &r0, &c0]<std::size_t... Rs>(std::index_sequence<Rs...>) {
            return matrix<Scalar, Rows - 1, Cols - 1>{
                (*this)[(Rs < r0) ? Rs : Rs + 1].minor(c0)...};
        }(std::make_index_sequence<Rows - 1>{});
    }

    constexpr auto determinant() const noexcept -> Scalar {
        static_assert(Rows == Cols,
                      "determinant only defined for square matrices");
        return [this]<std::size_t... Cs>(std::index_sequence<Cs...>) {
            return ((((Cs % 2 == 0) ? (*this)[0][Cs] : -(*this)[0][Cs]) *
                     (*this).minor(0u, Cs).determinant()) +
                    ...);
        }(std::make_index_sequence<Cols>{});
    }

    constexpr auto inverse() const noexcept -> matrix {
        const auto inv_det{Scalar(1) / this->determinant()};
        matrix<Scalar, Rows, Cols> inv; // no initialization needed!
        for (auto c = 0u; c < Cols; c++) {
            for (auto r = 0u; r < Rows; r++) {
                const auto min_det = this->minor(r, c).determinant();
                const auto co_fact{((r + c) % 2 == 1) ? -min_det : min_det};
                inv[r][c] = inv_det * co_fact;
            }
        }
        return inv;
    }

    constexpr matrix operator*(const Scalar& v) const noexcept {
        return [this, &v]<std::size_t... Rs>(std::index_sequence<Rs...>) {
            return matrix{((*this)[Rs] * v)...};
        }(std::make_index_sequence<Rows>{});
    }
    constexpr matrix operator/(const Scalar& v) const noexcept {
        return operator*(Scalar{1} / v);
    }

    constexpr matrix operator+(const Scalar& v) const noexcept {
        return [this, &v]<std::size_t... Rs>(std::index_sequence<Rs...>) {
            return matrix{((*this)[Rs] + v)...};
        }(std::make_index_sequence<Rows>{});
    }
    constexpr matrix operator-(const Scalar& v) const noexcept {
        return operator+(-v);
    }

    constexpr auto operator*=(const Scalar& v) noexcept {
        [this, &v]<std::size_t... Rs>(std::index_sequence<Rs...>) {
            (((*this)[Rs] *= v), ...);
        }(std::make_index_sequence<Rows>{});
        return *this;
    }

    constexpr auto operator/=(const Scalar& v) noexcept {
        return operator*=(Scalar{1} / v);
    }

    constexpr auto operator+=(const Scalar& v) noexcept {
        [this, &v]<std::size_t... Rs>(std::index_sequence<Rs...>) {
            (((*this)[Rs] += v), ...);
        }(std::make_index_sequence<Rows>{});
        return *this;
    }

    constexpr auto operator-=(const Scalar& v) noexcept {
        return operator+=(-v);
    }
};

template <> constexpr float matrix<float, 1, 1>::determinant() const noexcept {
    return (*this)[0][0];
}

template <>
constexpr double matrix<double, 1, 1>::determinant() const noexcept {
    return (*this)[0][0];
}

template <typename Scalar, unsigned int Ra, unsigned int CaRb, unsigned int Cb>
constexpr auto operator*(const matrix<Scalar, Ra, CaRb>& A,
                         const matrix<Scalar, CaRb, Cb>& B) noexcept {
    matrix<Scalar, Ra, Cb> result{};
    [&]<std::size_t... Ia>(std::index_sequence<Ia...>) {
        (
            [&]<std::size_t... Ib>(auto&& Ai, std::index_sequence<Ib...>) {
            (
                [&]<std::size_t... iI>(auto&& Ai, auto&& Bi,
                                       std::index_sequence<iI...>) {
                ((result[Ai][Bi] += A[Ai][iI] * B[iI][Bi]),
                 ...); // do the actual multiplication
            }(Ai, Ib, std::make_index_sequence<CaRb>{}),
                ...); // forward row of A, colum of B and unpack inner loop
        }(Ia, std::make_index_sequence<Cb>{}),
            ...); // forward current row, unpack A
    }(std::make_index_sequence<Ra>{}); // rows of A as input
    return result;
}

template <typename ScalarA, typename ScalarB, unsigned int Ra, unsigned int Ca>
constexpr auto operator*(const matrix<ScalarA, Ra, Ca>& A,
                         const vector<ScalarB, Ca>& B) noexcept {
    // first step - should move to concepts to allow for std::array and
    // std::span to be allowed
    vector<std::common_type_t<ScalarA, ScalarB>, Ca> result{};
    [&]<std::size_t... Ia>(std::index_sequence<Ia...>) {
        (
            [&]<std::size_t... Ib>(auto&& Ai, std::index_sequence<Ib...>) {
            ((result[Ib] += A[Ai][Ib] * B[Ib]), ...);
        }(Ia, std::make_index_sequence<Ca>{}),
            ...); // forward current row, unpack A
    }(std::make_index_sequence<Ra>{}); // rows of A as input
    return result;
}

} // namespace pw

//
// tuple - protocol
//
template <class Scalar, std::size_t R, std::size_t C>
struct std::tuple_size<pw::matrix<Scalar, R, C>>
    : std::integral_constant<std::size_t, R> {};

template <std::size_t I, class Scalar, std::size_t R, std::size_t C>
struct std::tuple_element<I, pw::matrix<Scalar, R, C>> {
    using type = pw::matrix<Scalar, R, C>::row_type;
};

#endif
