#pragma once

#include <stddef.h>
#include <ostream>
#include <utility>
#include <vector>
#include <cmath>
#include <assert.h>

inline size_t mod(int num, size_t module) {
    while (num < 0) { num += module; };

    return static_cast<size_t>(num % module);
}

enum Norm { Frob, Max };

template <typename T>
class Matrix {
public:
    // Default Constructor
    Matrix() = default;
    // Creation Constructor (0 Matrix)
    Matrix(size_t nrow, size_t ncol) :
        _nrow{nrow},
        _ncol{ncol},
        _data(nrow * ncol) { };
    // Creation Constructor with default Element (for example all elements 1)
    Matrix(size_t nrow, size_t ncol, const T& elem) :
        _nrow{nrow},
        _ncol{ncol},
        _data(nrow * ncol, elem) { };
    // Copy Constructor
    Matrix(const Matrix<T>& A) :
        _nrow{A._nrow},
        _ncol{A._ncol},
        _data(A._data) { };

    size_t nrow() const { return _nrow; };
    size_t ncol() const { return _ncol; };
    size_t nx() const { return _ncol; };
    size_t ny() const { return _nrow; };
    size_t size() const { return _nrow * _ncol; };

          T* data()       { return _data.data(); };
    const T* data() const { return _data.data(); };

          T& operator()(int i)       { return _data[i]; };
    const T& operator()(int i) const { return _data[i]; };

          T& operator()(int i, int j)       { return _data[i + j * _nrow]; }
    const T& operator()(int i, int j) const { return _data[i + j * _nrow]; }

          T& get(int i)       { return _data[index(i)]; };
    const T& get(int i) const { return _data[index(i)]; };

          T& get(int i, int j)       { return _data[index(i, j)]; }
    const T& get(int i, int j) const { return _data[index(i, j)]; }

    template <enum Norm N = Norm::Frob>
    double norm() const {
        T accum = static_cast<T>(0);    /*< result accumulator */

        // Enforce valid norm types (at compile time)
        static_assert(N == Norm::Frob || N == Norm::Max);

        if constexpr (N == Norm::Frob) {
            // Sum of Squares
            for (const T& elem : _data) {
                accum += elem * elem;
            }
            // The Frobenius norm is the square root of the sum of squares
            return std::sqrt(static_cast<double>(accum));
        } else { // Maximum Norm
            // Maximum absolute value
            for (const T& elem : _data) {
                if (accum < std::abs(elem)) {
                    accum = std::abs(elem);
                }
            }
            return static_cast<double>(accum);
        }
    }

    /**
     * Distance of this Matrix to given matrix A as || this - A ||_N
     *
     * Note: It there is a dimension mismatch, only the number of elements
     * of the smaller matrix are considered.
     *
     * @param A matrix to compute the "distance" to
     * @param mar_b bottom margins
     * @param mar_l left margins
     * @param mar_t top margins
     * @param mar_r right margins
     */
    template <enum Norm N = Norm::Frob>
    double dist(const Matrix<T>& A,
        size_t mar_b = 0, size_t mar_l = 0, size_t mar_t = 0, size_t mar_r = 0
    ) const {
        // Enforce valid norm types (at compile time)
        static_assert(N == Norm::Frob || N == Norm::Max);

        assert(this->_nrow == A._nrow);
        assert(this->_ncol == A._ncol);

        T accum = static_cast<T>(0);    /*< result accumulator */
        for (size_t j = mar_l; j + mar_r < _ncol; ++j) {
            for (size_t i = mar_t; i + mar_b < _nrow; ++i) {
                if constexpr (N == Norm::Frob) {
                    T diff = (*this)(i, j) - A(i, j);
                    accum += diff * diff;
                } else {
                    accum = std::max(std::abs((*this)(i, j) - A(i, j)), accum);
                }
            }
        }

        if constexpr (N == Norm::Frob) {
            return std::sqrt(static_cast<double>(accum));
        } else {
            return static_cast<double>(accum);
        }
    }

    /** overloaded standard swap for matrices (of same type)
     *
     * @example
     * Matrix<int> A(1, 2);
     * Matrix<int> B(3, 4);
     * std::swap(A, B);
     */
    friend void swap(Matrix<T>& A, Matrix<T>& B) {
        std::swap(A._nrow, B._nrow);
        std::swap(A._ncol, B._ncol);
        std::swap(A._data, B._data);
    }

private:
    size_t _nrow;           /*< Number of Rows */
    size_t _ncol;           /*< Number of Columns */
    std::vector<T> _data;   /*< Matrix elements / Data */

    size_t index(int i) const {
        const int nelem = static_cast<int>(_nrow * _ncol);

        while (i < 0)      { i += nelem; }
        while (i >= nelem) { i -= nelem; }

        return i;
    };

    size_t index(int i, int j) const {
        return static_cast<size_t>(mod(i, _nrow) + mod(j, _ncol) * _nrow);
    }
};

template <typename T>
std::ostream& operator<<(std::ostream& out, const Matrix<T>& mat) {
    for (size_t i = 0; i < mat.nrow(); ++i) {
        for (size_t j = 0; j < mat.ncol(); ++j) {
            out << mat(i, j) << ' ';
        }
        out << '\n';
    }
    return out;
}

template <typename T>
class MatrixView {
public:
    MatrixView(Matrix<T>& matrix, size_t index, size_t stride, size_t nelem) :
        _matrix(matrix), _index{index}, _stride{stride},
        _nelem{nelem} { };

    const size_t size() const { return _nelem; };

    const size_t stride() const { return _stride; };

          T* data()       { return _matrix.data() + _index; };
    const T* data() const { return _matrix.data() + _index; };

          T& operator()(int i)       { return _matrix(_index + i * _stride); };
    const T& operator()(int i) const { return _matrix(_index + i * _stride); };

          T& get(int i)       { return _matrix.get(_index + i * _stride); };
    const T& get(int i) const { return _matrix.get(_index + i * _stride); };

protected:
    Matrix<T>& _matrix;
    size_t _index;
    size_t _stride;
    size_t _nelem;
};

template <typename T>
std::ostream& operator<<(std::ostream& out, const MatrixView<T>& view) {
    for (size_t i = 0; i < view.size(); ++i) {
        out << view(i) << ' ';
    }

    return out;
}

template <typename T>
class Row : public MatrixView<T> {
public:
    Row(Matrix<T>& matrix, int index) :
        MatrixView<T>(matrix, mod(index, matrix.nrow()),
            matrix.nrow(), matrix.ncol()) { };
};

template <typename T>
class Col : public MatrixView<T> {
public:
    Col(Matrix<T>& matrix, int index) :
        MatrixView<T>(matrix, mod(index, matrix.ncol()) * matrix.nrow(),
            1, matrix.nrow()) { };
};

template <typename T>
class Diag : public MatrixView<T> {
public:
    Diag(Matrix<T>& matrix) :
        MatrixView<T>(matrix, 0, matrix.nrow() + 1,
            matrix.nrow() < matrix.ncol() ? matrix.nrow() : matrix.ncol()) { };
};