//
//  matrix.hpp
//  math
//
//  Created by Sam Jaffe on 5/28/16.
//

#pragma once

#include <expect/expect.hpp>
#include <math/vector/vector.hpp>

#include "math/matrix/forward.h"
#include "math/matrix/row_reference.hpp"
#include "math/matrix/traits.hpp"

#include "math/matrix/macro.h"

namespace math::matrix {

template <typename T, std::size_t R, std::size_t C> class matrix {
public:
  using value_type = T;

  template <typename M>
  using mul_t = decltype(std::declval<T>() * std::declval<M>());
  template <typename M>
  using div_t = decltype(std::declval<T>() / std::declval<M>());

public:
  matrix() = default;

  MATRIX_CTOR_N_ARGS(1) {}
  MATRIX_CTOR_N_ARGS(2) {}
  MATRIX_CTOR_N_ARGS(3) {}
  MATRIX_CTOR_N_ARGS(4) {}

  DEFER_RESOLUTION matrix(vector::vector<T, R> const & other,
                          DEFERRED_ENABLE_IF_T(C == 1, bool) = true) {
    VECTOR_FOR_EACH_RANGE(i, R) { _data[i][0] = other[i]; }
  }

  matrix(std::array<std::array<T, C>, R> const & init) {
    MATRIX_FOR_EACH(i, j) { _data[i][j] = init[i][j]; }
  }

  matrix(std::array<vector::vector<T, C>, R> const & init) {
    MATRIX_FOR_EACH(i, j) { _data[i][j] = init[i][j]; }
  }

  template <size_t R2, size_t C2> matrix(matrix<T, R2, C2> const & other) {
    MATRIX_FOR_EACH_RANGE(i, std::min(R, R2), j, std::min(C, C2)) {
      _data[i][j] = other(i, j);
    }
  }

  matrix<T, C, R> transpose() const {
    matrix<T, C, R> out;
    MATRIX_FOR_EACH(i, j) { out(j, i) = _data[i][j]; }
    return out;
  }

  template <size_t C2>
  matrix<T, R, C + C2> concat(matrix<T, R, C2> const & other,
                              concat_strategy::horizontal_concat_t) const {
    matrix<T, R, C + C2> accum{*this};
    MATRIX_FOR_EACH_RANGE(i, R, j, C2) { accum(i, j + C) = other(i, j); }
    return accum;
  }

  template <size_t R2>
  matrix<T, R + R2, C> concat(matrix<T, R2, C> const & other,
                              concat_strategy::vertical_concat_t) const {
    matrix<T, R + R2, C> accum{*this};
    MATRIX_FOR_EACH_RANGE(i, R2, j, C) { accum(i + R, j) = other(i, j); }
    return accum;
  }

  template <size_t R2, size_t C2>
  matrix<T, R + R2, C + C2> concat(matrix<T, R2, C2> const & other,
                                   concat_strategy::diagonal_concat_t) const {
    matrix<T, R + R2, C + C2> accum{*this};
    MATRIX_FOR_EACH_RANGE(i, R2, j, C2) { accum(i + R, j + C) = other(i, j); }
    return accum;
  }

  // In C++23 - we can use operator[](row, col)
  T const & operator()(std::size_t row, std::size_t col) const {
    return _data[row][col];
  }
  T & operator()(std::size_t row, std::size_t col) { return _data[row][col]; }
  row_reference<const T, C> operator[](std::size_t row) const {
    return {_data[row]};
  }
  row_reference<T, C> operator[](std::size_t row) { return {_data[row]}; }
  row_reference<const T, C> at(std::size_t row) const {
    expects(row < R, std::out_of_range, "row index out of range");
    return operator[](row);
  }
  row_reference<T, C> at(std::size_t row) {
    expects(row < R, std::out_of_range, "row index out of range");
    return operator[](row);
  }
  value_type const & at(std::size_t row, std::size_t col) const {
    expects(row < R && col < C, std::out_of_range, "coordinates out of range");
    return _data[row][col];
  }
  value_type & at(std::size_t row, std::size_t col) {
    expects(row < R && col < C, std::out_of_range, "coordinates out of range");
    return _data[row][col];
  }

  matrix & operator+=(matrix const & other) {
    MATRIX_FOR_EACH(i, j) { _data[i][j] += other(i, j); }
    return *this;
  }
  matrix operator+(matrix const & other) const {
    return matrix{*this} += other;
  }
  matrix & operator-=(matrix const & other) {
    MATRIX_FOR_EACH(i, j) { _data[i][j] -= other(i, j); }
    return *this;
  }
  matrix operator-(matrix const & other) const {
    return matrix{*this} -= other;
  }

  matrix operator-() const {
    matrix tmp;
    MATRIX_FOR_EACH(i, j) { tmp(i, j) = -_data[i][j]; }
    return tmp;
  }

  vector::vector<T, C> operator*(vector::vector<T, C> const & vec) const {
    vector::vector<T, C> rval;
    MATRIX_FOR_EACH(i, j) { rval[i] += _data[i][j] * vec[j]; }
    return rval;
  }

  template <std::size_t C2>
  matrix<T, R, C2> operator*(matrix<T, C, C2> const & other) const {
    matrix<T, R, C2> rval;
    MATRIX_FOR_EACH(i, j) {
      for (size_t k = 0; k < C2; ++k) {
        rval(i, k) += _data[i][j] * other(j, k);
      }
    }
    return rval;
  }

  matrix<T, R, C> & operator*=(T c) {
    MATRIX_FOR_EACH(i, j) { _data[i][j] *= c; }
    return *this;
  }

  template <typename M>
  MATRIX_DISABLE_IF_MATRIX(M, mul_t<M>, R, C)
  operator*(M c) const {
    return matrix<mul_t<M>, R, C>{*this} *= c;
  }

  template <typename M>
  friend MATRIX_DISABLE_IF_MATRIX(M, mul_t<M>, R, C)
  operator*(M c, matrix const & matr) {
    return matrix<mul_t<M>, R, C>{matr} *= c;
  }

  template <typename M> matrix<div_t<M>, R, C> & operator/=(M c) {
    MATRIX_FOR_EACH(i, j) { _data[i][j] /= c; }
    return *this;
  }
  template <typename M> matrix<div_t<M>, R, C> operator/(M c) const {
    return matrix<mul_t<M>, R, C>{*this} /= c;
  }

  bool operator==(matrix const & other) const {
    MATRIX_FOR_EACH(i, j) {
      if (_data[i][j] != other(i, j)) { return false; }
    }
    return true;
  }
  bool operator!=(matrix const & other) const { return !operator==(other); }

private:
  value_type _data[R][C] = {value_type()};
};

MATRIX_CTOR_DEDUCTION(1);
MATRIX_CTOR_DEDUCTION(2);
MATRIX_CTOR_DEDUCTION(3);
MATRIX_CTOR_DEDUCTION(4);

}

#include "math/matrix/undef.h"