#pragma once

#include <cassert>
#include <type_traits>
#include <vector>
#include <ostream>

namespace amt {

  template <typename T>
    requires std::is_integral_v<T>
    struct Matrix {
      using base_type = std::vector<T>;
      using value_type = typename base_type::value_type;
      using pointer = typename base_type::pointer;
      using const_pointer = typename base_type::const_pointer;
      using reference = typename base_type::reference;
      using const_reference = typename base_type::const_reference;
      using iterator = typename base_type::iterator;
      using const_iterator = typename base_type::const_iterator;
      using size_type = std::size_t;

      template <bool IsConst>
        struct View {
          using base_type = std::conditional_t<IsConst, const_pointer, pointer>;
          base_type data;
          size_type size;

          constexpr auto operator[](size_type k) const noexcept {
            assert(k < size && "Out of bound access");
            return data[k];
          }

          constexpr auto operator[](size_type k) noexcept requires (!IsConst) {
            assert(k < size && "Out of bound access");
            return data[k];
          }
        };

      constexpr Matrix() noexcept = default;
      constexpr Matrix(Matrix const&) = default;
      constexpr Matrix& operator=(Matrix const&) = default;
      constexpr Matrix(Matrix &&) noexcept = default;
      constexpr Matrix& operator=(Matrix &&) noexcept = default;
      constexpr ~Matrix() = default;


      Matrix(size_type row, size_type col, value_type def = {})
        : m_data(row * col, def)
        , m_row(row)
        , m_col(col)
        {}

      constexpr bool empty() const noexcept { return m_data.empty(); };
      constexpr bool size() const noexcept { return m_data.size(); };
      constexpr size_type rows() const noexcept { return m_row; };
      constexpr size_type cols() const noexcept { return m_col; };

      constexpr iterator begin() noexcept { return m_data.begin(); }
      constexpr iterator end() noexcept { return m_data.end(); }
      constexpr const_iterator begin() const noexcept { return m_data.begin(); }
      constexpr const_iterator end() const noexcept { return m_data.end(); }

      constexpr auto operator()(size_type r, size_type c) noexcept -> reference {
        auto const index = r * m_col + c; // column-major;
        assert(index < size() && "Out of bound access");
        return m_data[index];
      }

      constexpr auto operator()(size_type r, size_type c) const noexcept -> const_reference {
        auto const index = r * m_col + c; // column-major;
        assert(index < size() && "Out of bound access");
        return m_data[index];
      }

      constexpr auto operator[](size_type r) noexcept -> View<false> {
        auto const base = r * m_col;
        assert(r < rows() && "Out of bound access");
        return { .data = m_data.data() + base, .size = m_col };
      }

      constexpr auto operator[](size_type r) const noexcept -> View<true> {
        auto const base = r * m_col;
        assert(r < rows() && "Out of bound access");
        return { .data = m_data.data() + base, .size = m_col };
      }

      friend std::ostream& operator<<(std::ostream& os, Matrix const& m) {
        os << "[\n";
        for (auto i = size_type{}; i < m.rows(); ++i) {
          for (auto j = size_type{}; j < m.cols(); ++j) {
            os << m[i][j] << ", ";
          }
          os << '\n';
        }
        return os << "]";
      }

      private:
      base_type m_data;
      size_type m_row{};
      size_type m_col{};
    };

} // namespace amt
