#ifndef PYTHONIC_INCLUDE_NUMPY_ARANGE_HPP
#define PYTHONIC_INCLUDE_NUMPY_ARANGE_HPP

#include "pythonic/include/utils/functor.hpp"
#include "pythonic/include/operator_/pos.hpp"
#include "pythonic/include/types/ndarray.hpp"

PYTHONIC_NS_BEGIN

namespace numpy
{
  namespace details
  {
#ifdef USE_XSIMD
    template <class T>
    struct arange_simd_iterator {
      using vector_type = xsimd::simd_type<T>;
      vector_type curr_;
      vector_type step_;
      long index_;
      arange_simd_iterator(T start, T step, long n)
          : curr_{}, step_{static_cast<T>(vector_type::size * step)},
            index_{static_cast<long>(n / vector_type::size)}
      {
        T from[vector_type::size];
        for (size_t i = 0; i < vector_type::size; ++i)
          from[i] = start + i * step;
        curr_ = xsimd::load_unaligned(from);
      }
      vector_type operator*() const
      {
        return curr_;
      }
      arange_simd_iterator &operator++()
      {
        curr_ += step_;
        ++index_;
        return *this;
      }
      arange_simd_iterator &operator+=(long n)
      {
        curr_ += n * step_;
        index_ += n;
        return *this;
      }
      arange_simd_iterator operator+(long n) const
      {
        arange_simd_iterator other{*this};
        return other += n;
      }
      arange_simd_iterator &operator--()
      {
        curr_ -= step_;
        --index_;
        return *this;
      }
      long operator-(arange_simd_iterator const &other) const
      {
        return index_ - other.index_;
      }
      bool operator!=(arange_simd_iterator const &other) const
      {
        return index_ != other.index_;
      }
      bool operator==(arange_simd_iterator const &other) const
      {
        return index_ == other.index_;
      }
      bool operator<(arange_simd_iterator const &other) const
      {
        return index_ < other.index_;
      }
      arange_simd_iterator &
      operator=(arange_simd_iterator const &other) = default;
    };
#endif
    template <class T>
    struct arange_index {
      T start, step;
      long size;
      using iterator = types::nditerator<arange_index>;
      using const_iterator = types::const_nditerator<arange_index>;
      using dtype = T;
      using value_type = dtype;
      using shape_t = types::pshape<long>;
#ifdef USE_XSIMD
      using simd_iterator = arange_simd_iterator<T>;
      using simd_iterator_nobroadcast = simd_iterator;
      template <class vectorizer>
      simd_iterator vbegin(vectorizer) const
      {
        return {start, step, 0};
      }
      template <class vectorizer>
      simd_iterator vend(vectorizer) const
      {
        return {static_cast<T>(start + size * step), step, size};
      }
#endif
      static constexpr size_t value = 1;
      static constexpr bool is_strided = false;
      static constexpr bool is_vectorizable = types::is_vectorizable<T>::value;

      T fast(long i) const
      {
        return start + i * step;
      }

      dtype load(long i) const
      {
        return fast(i);
      }

      template <size_t I>
      long shape() const
      {
        return size;
      }
      types::ndarray<dtype, shape_t> operator[](types::slice s) const
      {
        auto ns = s.normalize(size);
        arange_index r{start + s.lower * step, step * ns.step, ns.size()};
        return {
            types::numpy_expr<pythonic::operator_::functor::pos, arange_index>{
                r}};
      }
      types::ndarray<dtype, shape_t> operator()(types::slice s) const
      {
        return operator[](s);
      }
      types::ndarray<dtype, shape_t> operator[](types::contiguous_slice s) const
      {
        auto ns = s.normalize(size);
        arange_index r{start + s.lower * step, step, ns.size()};
        return {
            types::numpy_expr<pythonic::operator_::functor::pos, arange_index>{
                r}};
      }
      types::ndarray<dtype, shape_t> operator()(types::contiguous_slice s) const
      {
        return operator[](s);
      }
      template <class... S>
      auto operator()(S const &... s) const -> typename std::enable_if<
          (sizeof...(S) > 1),
          decltype(std::declval<types::ndarray<dtype, shape_t>>()(s...))>::type
      {
        return types::ndarray<dtype, shape_t>{
            types::numpy_expr<pythonic::operator_::functor::pos, arange_index>{
                *this}}(s...);
      }
      const_iterator begin() const
      {
        return {*this, 0};
      }
      const_iterator end() const
      {
        return {*this, size};
      }

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

  template <class T, class U, class S = long,
            class dtype = types::dtype_t<typename __combined<T, U, S>::type>>
  types::numpy_expr<pythonic::operator_::functor::pos,
                    details::arange_index<typename dtype::type>>
  arange(T begin, U end, S step = S(1), dtype d = dtype());

  template <class T>
  types::numpy_expr<pythonic::operator_::functor::pos,
                    details::arange_index<typename types::dtype_t<T>::type>>
  arange(T end);

  DEFINE_FUNCTOR(pythonic::numpy, arange);
}
PYTHONIC_NS_END

#endif