#ifndef VECTOR_HPP
#define VECTOR_HPP

#include <utils/PugsMacros.hpp>
#include <utils/PugsUtils.hpp>

#include <utils/PugsAssert.hpp>

#include <utils/Array.hpp>

template <typename DataType>
class Vector   // LCOV_EXCL_LINE
{
 public:
  using data_type  = DataType;
  using index_type = size_t;

 private:
  Array<DataType> m_values;
  const bool m_deep_copies;

  static_assert(std::is_same_v<typename decltype(m_values)::index_type, index_type>);

  // Allows const version to access our data
  friend Vector<std::add_const_t<DataType>>;

 public:
  friend std::ostream&
  operator<<(std::ostream& os, const Vector& x)
  {
    for (size_t i = 0; i < x.size(); ++i) {
      os << ' ' << x[i];
    }
    return os;
  }

  friend Vector<std::remove_const_t<DataType>>
  copy(const Vector& x)
  {
    auto values = copy(x.m_values);

    Vector<std::remove_const_t<DataType>> x_copy{0};
    x_copy.m_values = values;
    return x_copy;
  }

  friend Vector
  operator*(const DataType& a, const Vector& x)
  {
    Vector<std::remove_const_t<DataType>> y = copy(x);
    return y *= a;
  }

  template <typename DataType2>
  PUGS_INLINE friend auto
  dot(const Vector& x, const Vector<DataType2>& y)
  {
    Assert(x.size() == y.size());
    // Quite ugly, TODO: fix me in C++20
    auto promoted = [] {
      DataType a{0};
      DataType2 b{0};
      return a * b;
    }();

    decltype(promoted) sum = 0;

    // Does not use parallel_for to preserve sum order
    for (index_type i = 0; i < x.size(); ++i) {
      sum += x[i] * y[i];
    }

    return sum;
  }

  template <typename DataType2>
  PUGS_INLINE Vector&
  operator/=(const DataType2& a)
  {
    const auto inv_a = 1. / a;
    return (*this) *= inv_a;
  }

  template <typename DataType2>
  PUGS_INLINE Vector&
  operator*=(const DataType2& a)
  {
    parallel_for(
      this->size(), PUGS_LAMBDA(index_type i) { m_values[i] *= a; });
    return *this;
  }

  template <typename DataType2>
  PUGS_INLINE Vector&
  operator-=(const Vector<DataType2>& y)
  {
    Assert(this->size() == y.size());

    parallel_for(
      this->size(), PUGS_LAMBDA(index_type i) { m_values[i] -= y[i]; });

    return *this;
  }

  template <typename DataType2>
  PUGS_INLINE Vector&
  operator+=(const Vector<DataType2>& y)
  {
    Assert(this->size() == y.size());

    parallel_for(
      this->size(), PUGS_LAMBDA(index_type i) { m_values[i] += y[i]; });

    return *this;
  }

  template <typename DataType2>
  PUGS_INLINE Vector<std::remove_const_t<DataType>>
  operator+(const Vector<DataType2>& y) const
  {
    Assert(this->size() == y.size());
    Vector<std::remove_const_t<DataType>> sum{y.size()};

    parallel_for(
      this->size(), PUGS_LAMBDA(index_type i) { sum.m_values[i] = m_values[i] + y[i]; });

    return sum;
  }

  template <typename DataType2>
  PUGS_INLINE Vector<std::remove_const_t<DataType>>
  operator-(const Vector<DataType2>& y) const
  {
    Assert(this->size() == y.size());
    Vector<std::remove_const_t<DataType>> sum{y.size()};

    parallel_for(
      this->size(), PUGS_LAMBDA(index_type i) { sum.m_values[i] = m_values[i] - y[i]; });

    return sum;
  }

  PUGS_INLINE
  DataType&
  operator[](index_type i) const noexcept(NO_ASSERT)
  {
    return m_values[i];
  }

  PUGS_INLINE
  size_t
  size() const noexcept
  {
    return m_values.size();
  }

  PUGS_INLINE Vector&
  operator=(const DataType& value) noexcept
  {
    m_values.fill(value);
    return *this;
  }

  template <typename DataType2>
  PUGS_INLINE Vector&
  operator=(const Vector<DataType2>& x) noexcept
  {
    // ensures that DataType is the same as source DataType2
    static_assert(std::is_same<std::remove_const_t<DataType>, std::remove_const_t<DataType2>>(),
                  "Cannot assign Vector of different type");
    // ensures that const is not lost through copy
    static_assert(((std::is_const<DataType2>() and std::is_const<DataType>()) or not std::is_const<DataType2>()),
                  "Cannot assign Vector of const to Vector of non-const");

    if (m_deep_copies) {
      copy_to(x.m_values, m_values);
    } else {
      m_values = x.m_values;
    }
    return *this;
  }

  PUGS_INLINE
  Vector&
  operator=(const Vector& x)
  {
    if (m_deep_copies) {
      copy_to(x.m_values, m_values);
    } else {
      m_values = x.m_values;
    }
    return *this;
  }

  PUGS_INLINE
  Vector&
  operator=(Vector&& x)
  {
    if (m_deep_copies) {
      copy_to(x.m_values, m_values);
    } else {
      m_values = x.m_values;
    }
    return *this;
  }

  template <typename DataType2>
  Vector(const Vector<DataType2>& x) : m_deep_copies{x.m_deep_copies}
  {
    // ensures that DataType is the same as source DataType2
    static_assert(std::is_same<std::remove_const_t<DataType>, std::remove_const_t<DataType2>>(),
                  "Cannot assign Vector of different type");
    // ensures that const is not lost through copy
    static_assert(((std::is_const<DataType2>() and std::is_const<DataType>()) or not std::is_const<DataType2>()),
                  "Cannot assign Vector of const to Vector of non-const");

    m_values = x.m_values;
  }

  explicit Vector(const Array<DataType>& values) : m_deep_copies{true}
  {
    m_values = values;
  }

  Vector(const Vector&) = default;

  Vector(Vector&&) = default;

  Vector(size_t size) : m_values{size}, m_deep_copies{false} {}
  ~Vector() = default;
};

#endif   // VECTOR_HPP