// Courtesy of RcppParallel: https://github.com/RcppCore/RcppParallel

#pragma once

#include <cstddef>
#include <iterator>

namespace RcppThread {


// Thread safe safe matrix wrapper, courtesy of RcppParallel:
// https://github.com/RcppCore/RcppParallel
template <typename T> class RMatrix {
public:
  class Row {

  public:
    template <typename V> class row_iterator {

    public:
      using iterator_category = std::random_access_iterator_tag;
      using value_type = V;
      using difference_type = std::size_t;
      using pointer = value_type *;
      using reference = value_type &;

      inline row_iterator(Row &row, difference_type i)
          : start_(row.start_), parentNrow_(row.parent_.nrow()), index_(i) {}

      inline row_iterator(pointer start, difference_type parentNrow,
                          difference_type index)
          : start_(start), parentNrow_(parentNrow), index_(index) {}

      inline row_iterator(const row_iterator &other)
          : start_(other.start_), parentNrow_(other.parentNrow_),
            index_(other.index_) {}

      inline row_iterator &operator++() {
        index_++;
        return *this;
      }

      inline row_iterator operator++(int) {
        row_iterator tmp(*this);
        operator++();
        return tmp;
      }

      inline row_iterator &operator--() {
        index_--;
        return *this;
      }

      inline row_iterator operator--(int) {
        row_iterator tmp(*this);
        index_--;
        return tmp;
      }

      row_iterator operator+(difference_type n) const {
        return row_iterator(start_, parentNrow_, index_ + n);
      }
      row_iterator operator-(difference_type n) const {
        return row_iterator(start_, parentNrow_, index_ - n);
      }

      difference_type operator+(const row_iterator &other) const {
        return index_ + other.index_;
      }

      difference_type operator-(const row_iterator &other) const {
        return index_ - other.index_;
      }

      row_iterator &operator+=(difference_type n) {
        index_ += n;
        return *this;
      }
      row_iterator &operator-=(difference_type n) {
        index_ -= n;
        return *this;
      }

      bool operator==(const row_iterator &other) const {
        return index_ == other.index_;
      }
      bool operator!=(const row_iterator &other) const {
        return index_ != other.index_;
      }
      bool operator<(const row_iterator &other) const {
        return index_ < other.index_;
      }
      bool operator>(const row_iterator &other) const {
        return index_ > other.index_;
      }
      bool operator<=(const row_iterator &other) const {
        return index_ <= other.index_;
      }
      bool operator>=(const row_iterator &other) const {
        return index_ >= other.index_;
      }

      inline reference operator*() { return start_[index_ * parentNrow_]; }

      inline pointer operator->() { return &(start_[index_ * parentNrow_]); }

      inline reference operator[](int i) {
        return start_[(index_ + i) * parentNrow_];
      }

    private:
      pointer start_;
      difference_type parentNrow_;
      difference_type index_;
    };

    typedef row_iterator<T> iterator;
    typedef row_iterator<const T> const_iterator;

    inline Row(RMatrix &parent, std::size_t i)
        : parent_(parent), start_(parent.begin() + i) {}

    inline Row(const Row &other)
        : parent_(other.parent_), start_(other.start_) {}

    inline iterator begin() { return iterator(*this, 0); }

    inline iterator end() { return iterator(*this, parent_.ncol()); }

    inline const_iterator begin() const { return const_iterator(*this, 0); }

    inline const_iterator end() const {
      return const_iterator(*this, parent_.ncol());
    }

    inline size_t length() const { return parent_.ncol(); }

    inline size_t size() const { return parent_.ncol(); }

    inline T &operator[](std::size_t i) { return start_[i * parent_.nrow()]; }

    inline const T &operator[](std::size_t i) const {
      return start_[i * parent_.nrow()];
    }

  private:
    RMatrix &parent_;
    T *start_;
  };

  class Column {

  public:
    typedef T *iterator;
    typedef const T *const_iterator;

    inline Column(RMatrix &parent, std::size_t i)
        : begin_(parent.begin() + (i * parent.nrow())),
          end_(begin_ + parent.nrow()) {}

    inline Column(const Column &other)
        : begin_(other.begin_), end_(other.end_) {}

    inline Column &operator=(const Column &rhs) {
      begin_ = rhs.begin_;
      end_ = rhs.end_;
      return *this;
    }

    inline iterator begin() { return begin_; }
    inline iterator end() { return end_; }

    inline const_iterator begin() const { return begin_; }
    inline const_iterator end() const { return end_; }

    inline size_t length() const { return end_ - begin_; }
    inline size_t size() const { return end_ - begin_; }

    inline T &operator[](std::size_t i) { return *(begin_ + i); }

    inline const T &operator[](std::size_t i) const { return *(begin_ + i); }

  private:
    T *begin_;
    T *end_;
  };

  typedef T *iterator;
  typedef const T *const_iterator;

  template <typename Source>
  inline explicit RMatrix(const Source &source)
      : data_(const_cast<Source &>(source).begin()), nrow_(source.nrow()),
        ncol_(source.ncol()) {}

  inline RMatrix(T *data, std::size_t nrow, std::size_t ncol)
      : data_(data), nrow_(nrow), ncol_(ncol) {}

  inline iterator begin() { return data_; }
  inline iterator end() { return data_ + length(); }

  inline const_iterator begin() const { return data_; }
  inline const_iterator end() const { return data_ + length(); }

  inline std::size_t length() const { return nrow_ * ncol_; }

  inline std::size_t nrow() const { return nrow_; }
  inline std::size_t ncol() const { return ncol_; }

  inline T &operator()(std::size_t i, std::size_t j) {
    return *(data_ + (i + j * nrow_));
  }

  inline const T &operator()(std::size_t i, std::size_t j) const {
    return *(data_ + (i + j * nrow_));
  }

  inline Row row(std::size_t i) { return Row(*this, i); }

  inline const Row row(std::size_t i) const {
    return Row(*const_cast<RMatrix *>(this), i);
  }

  inline Column column(std::size_t i) { return Column(*this, i); }

  inline const Column column(std::size_t i) const {
    return Column(*const_cast<RMatrix *>(this), i);
  }

  inline T &operator[](std::size_t i) { return *(data_ + i); }

  inline const T &operator[](std::size_t i) const { return *(data_ + i); }

private:
  T *data_;
  std::size_t nrow_;
  std::size_t ncol_;
};

} // namespace RcppThread
