#pragma once

#include <type_traits>
#include <vector>

namespace AMDiS
{
  /// Dense matrix with row-wise storage in flat data vector
  template <class T, class Allocator = std::allocator<T>>
  class FlatMatrix
      : public std::vector<T,Allocator>
  {
    using Super = std::vector<T,Allocator>;

  public:
    using value_type = T;
    using size_type = typename Super::size_type;
    using reference = typename Super::reference;
    using const_reference = typename Super::const_reference;

  private:
    /// Accessor to a row of the matrix.
    template <class Vector>
    struct AccessProxy
    {
      /// Access the column entry `col` of the row `row_` of this accessor.
      /// @{
      template <class V = Vector, std::enable_if_t<not std::is_const<V>::value,int> = 0>
      reference operator[](size_type col)
      {
        return (*vector_)[row_*cols_ + col];
      }

      const_reference operator[](size_type col) const
      {
        return (*vector_)[row_*cols_ + col];
      }
      // @}

      Vector* vector_;
      size_type row_;
      size_type cols_;
    };

  public:
    /// Default constructor, creates an empty 0x0 matrix.
    FlatMatrix() = default;

    /// Construct a new matrix with rows `r` and columns `c` and all values
    /// initialized by copies of value `v`.
    FlatMatrix(size_type r, size_type c, value_type v = {})
      : Super(r*c, v)
      , rows_(r)
      , cols_(c)
    {}

    /// Assign value `s` to all entries of the matrix
    FlatMatrix& operator=(value_type s)
    {
      Super::assign(Super::size(), s);
      return *this;
    }

    /// Resizes the container to contain r x c elements.
    /**
     * If the current `rows*cols` is greater than `r*c`, the container is reduced
     * to its first `r*c` elements. If the current `rows*cols` is less than `r*c`,
     * additional copies of value `v` are appended.
     **/
    void resize(size_type r, size_type c, value_type v = {})
    {
      Super::resize(r*c, v);
      rows_ = r;
      cols_ = c;
    }

    /// Return the number of rows of the matrix
    size_type rows() const noexcept
    {
      return rows_;
    }

    /// Return the number of columns of the matrix
    size_type cols() const noexcept
    {
      return cols_;
    }

    /// Return the number of entries in the matrix, i.e. `rows*cols`.
    using Super::size;

    /// Return accessor to the `row` of the matrix, see \ref AccessProxy
    /// @{
    AccessProxy<Super> operator[](size_type row)
    {
      return AccessProxy<Super>{static_cast<Super*>(this), row, cols_};
    }

    AccessProxy<Super const> operator[](size_type row) const
    {
      return AccessProxy<Super const>{static_cast<Super const*>(this), row, cols_};
    }
    /// @}

  private:
    size_type rows_ = 0;
    size_type cols_ = 0;
  };

} // end namespace AMDiS