#ifndef __MATRIX__
#define __MATRIX__

#include <iostream>
#include <vector>

/** This is a N-Dimension capable matrix class.
 *
 * Only to be used as N-dimensional lookup grid. No mathematical operations have
 * been implemented.
 */
template <typename T> class Matrix {

  private:
    std::vector<int> dim   = {0};
    size_t           numel = 0;
    T *              data  = NULL;

  public:
    // ====================================================================== //
    // Allocation and de-allocation

    /** Default cnstructor. */
    Matrix(){};
    /** Destructor. */
    ~Matrix() { delete[] this->data; };

    /** Construct matrix with a desired size. */
    Matrix(int I, int J = 1, int K = 1, int L = 1, int M = 1) {
        this->dim   = {I, J, K, L, M};
        this->numel = I * J * K * L * M;
        this->data  = new T[this->numel];
    };
    /** Construct a matrix with size designated with vector. */
    Matrix(std::vector<int> dim) {
        int N = 1;
        for (int i : dim) N *= i;

        this->dim   = dim;
        this->numel = N;
        this->data  = new T[this->numel];
    };

    /** Copy constructor. */
    Matrix(const Matrix &original) {

        // Clear old data
        delete[] this->data;

        // Assign regular fields
        this->dim   = original.dim;
        this->numel = original.numel;
        this->data  = new T[original.numel];
        std::copy(original.data, original.data + original.numel, this->data);
    };
    /** Move constructor. */
    Matrix(Matrix &&original) {

        // Clear old data
        delete[] this->data;

        // Copy over new values and pointers
        this->dim   = original.dim;
        this->numel = original.numel;
        this->data  = original.data;

        // Set old values to nothing
        original.dim   = {0};
        original.numel = 0;
        original.data  = NULL;
    };

    // ====================================================================== //
    // Operator overloads

    /** Copy assignment operator. */
    Matrix<T> &operator=(Matrix<T> &original) {

        // Handle pointer data
        delete[] this->data;

        // Assign regular fields
        this->dim   = original.dim;
        this->numel = original.numel;
        this->data  = new T[original.numel];
        std::copy(original.data, original.data + original.numel, this->data);

        return *this;
    };
    /** Move assignment operator. */
    Matrix<T> &operator=(Matrix<T> &&original) {

        // Clear old data
        delete[] this->data;

        // Copy over new values and pointers
        this->dim   = original.dim;
        this->numel = original.numel;
        this->data  = original.data;

        // Set old values to nothing
        original.dim   = std::vector<int>({0});
        original.numel = 0;
        original.data  = NULL;

        return *this;
    };

    // ====================================================================== //
    // Size, dimensions and Printing

    /** Return the number of elements in the matrix. */
    size_t size() { return this->numel; };

    /** Return the vector of dimension sizes. */
    std::vector<int> get_dim() { return this->dim; };

    /** Print the dimensions of the matrix. */
    void print_size() {
        std::cout << "Size:\n [ ";
        for (int i : this->dim) printf("%d ", i);
        std::cout << "] " << std::endl;
    };

    // ====================================================================== //
    // Access fields

    /** Linear access. */
    T get(int i) { return this->get({i}); };
    /** Multi dimensional access. */
    T get(int i, int j, int k = 0, int l = 0, int m = 0) {
        return this->get({i, j, k, l, m});
    };
    /** Vector based access in ND. */
    T get(std::vector<int> idx) {

        // Compute the array position
        int tmp = 1, i = 0;
        for (size_t j = 0; j < idx.size(); j++) {
            i += idx[j] * tmp;
            tmp *= this->dim[j];
        }

        // Access the array
        return this->data[i];
    };

    /** Global assignment. */
    void set(T v) {
        for (size_t i = 0; i < this->numel; i++) this->data[i] = v;
    };
    /** Linear access. */
    void set(T v, int i) { this->data[i] = v; };
    /** Multi dimensional access. */
    void set(T v, int i, int j, int k = 0, int l = 0, int m = 0) {
        this->set(v, {i, j, k, l, m});
    };
    /** Vector based acces in ND. */
    void set(T v, std::vector<int> idx) {

        // Compute the array position
        int tmp = 1, i = 0;
        for (size_t j = 0; j < idx.size(); j++) {
            i += idx[j] * tmp;
            tmp *= this->dim[j];
        }

        // Access data
        this->data[i] = v;
    };
};

#endif /* __MATRIX__ */