#ifndef PYTHONIC_NUMPY_ARGMINMAX_HPP
#define PYTHONIC_NUMPY_ARGMINMAX_HPP

#include "pythonic/utils/functor.hpp"
#include "pythonic/types/ndarray.hpp"
#include "pythonic/numpy/asarray.hpp"
#include "pythonic/builtins/ValueError.hpp"

PYTHONIC_NS_BEGIN

namespace numpy
{
  namespace details
  {
    template <class P, size_t... Is>
    P iota(utils::index_sequence<Is...>)
    {
      return {static_cast<typename P::value_type>(Is)...};
    }

    template <class P>
    P iota()
    {
      return iota<P>(utils::make_index_sequence<P::size>());
    }
  }
  template <class Op, class E, class T>
  long _argminmax_seq(E const &elts, T &minmax_elts)
  {
    long index = 0;
    long res = -1;
    for (auto const &elt : elts) {
      if (Op::value(elt, minmax_elts)) {
        minmax_elts = elt;
        res = index;
      }
      ++index;
    }

    return res;
  }
  template <class Op, class E, class T>
#ifdef USE_XSIMD
  typename std::enable_if<
      !E::is_vectorizable ||
          !types::is_vector_op<typename Op::op, T, T>::value ||
          std::is_same<typename E::dtype, bool>::value,
      long>::type
#else
  long
#endif
  _argminmax(E const &elts, T &minmax_elts, utils::int_<1>)
  {
    return _argminmax_seq<Op>(elts, minmax_elts);
  }

  template <class Op, class E, class T, class... Indices>
  std::tuple<long, long> _argminmax_fast(E const &elts, T &minmax_elts,
                                         long current_pos, utils::int_<1>,
                                         Indices... indices)
  {
    long res = -1;
    long n = elts.template shape<std::decay<E>::type::value - 1>();
    for (long i = 0; i < n; ++i) {
      auto elt = elts.load(indices..., i);
      if (Op::value(elt, minmax_elts)) {
        minmax_elts = elt;
        res = current_pos + i;
      }
    }

    return std::make_tuple(res, current_pos + n);
  }

#ifdef USE_XSIMD
  template <bool IsInt>
  struct bool_caster;
  template <>
  struct bool_caster<true> {
    template <class T>
    auto operator()(T const &value) -> decltype(xsimd::bool_cast(value))
    {
      return xsimd::bool_cast(value);
    }
  };
  template <>
  struct bool_caster<false> {
    template <class T>
    T operator()(T const &value)
    {
      return value;
    }
  };

  template <class Op, class E, class T>
  typename std::enable_if<
      E::is_vectorizable && types::is_vector_op<typename Op::op, T, T>::value &&
          !std::is_same<typename E::dtype, bool>::value,
      long>::type
  _argminmax(E const &elts, T &minmax_elts, utils::int_<1>)
  {
    using vT = xsimd::simd_type<T>;
    using iT = xsimd::as_integer_t<T>;
    static const size_t vN = vT::size;
    const long n = elts.size();
    if (n >= std::numeric_limits<iT>::max()) {
      return _argminmax_seq<Op>(elts, minmax_elts);
    }

    auto viter = types::vectorizer_nobroadcast::vbegin(elts),
         vend = types::vectorizer_nobroadcast::vend(elts);

    const long bound = std::distance(viter, vend);
    long minmax_index = -1;
    if (bound > 0) {
      auto vacc = *viter;
      iT iota[vN] = {0};
      for (long i = 0; i < (long)vN; ++i)
        iota[i] = i;
      auto curr = xsimd::load_unaligned(iota);
      xsimd::simd_type<iT> indices = curr;
      xsimd::simd_type<iT> step{vN};

      for (++viter; viter != vend; ++viter) {
        curr += step;
        auto c = *viter;
        vacc = typename Op::op{}(vacc, c);
        auto mask = c == vacc;
        indices =
            xsimd::select(bool_caster<std::is_floating_point<T>::value>{}(mask),
                          curr, indices);
      }

      alignas(sizeof(vT)) T stored[vN];
      vacc.store_aligned(&stored[0]);
      alignas(sizeof(vT)) long indexed[vN];
      indices.store_aligned(&indexed[0]);

      for (size_t j = 0; j < vN; ++j) {
        if (Op::value(stored[j], minmax_elts)) {
          minmax_elts = stored[j];
          minmax_index = indexed[j];
        }
      }
    }
    auto iter = elts.begin() + bound * vN;

    for (long i = bound * vN; i < n; ++i, ++iter) {
      if (Op::value(*iter, minmax_elts)) {
        minmax_elts = *iter;
        minmax_index = i;
      }
    }
    return minmax_index;
  }
#endif

  template <class Op, class E, size_t N, class T>
  long _argminmax(E const &elts, T &minmax_elts, utils::int_<N>)
  {
    long current_pos = 0;
    long current_minmaxarg = 0;
    for (auto &&elt : elts) {
      long v = _argminmax<Op>(elt, minmax_elts, utils::int_<N - 1>());
      if (v >= 0)
        current_minmaxarg = current_pos + v;
      current_pos += elt.flat_size();
    }
    return current_minmaxarg;
  }
  template <class Op, class E, size_t N, class T, class... Indices>
  typename std::enable_if<N != 1, std::tuple<long, long>>::type
  _argminmax_fast(E const &elts, T &minmax_elts, long current_pos,
                  utils::int_<N>, Indices... indices)
  {
    long current_minmaxarg = 0;
    for (long i = 0, n = elts.template shape<std::decay<E>::type::value - N>();
         i < n; ++i) {
      long v;
      std::tie(v, current_pos) = _argminmax_fast<Op>(
          elts, minmax_elts, current_pos, utils::int_<N - 1>(), indices..., i);
      if (v >= 0)
        current_minmaxarg = v;
    }
    return std::make_tuple(current_minmaxarg, current_pos);
  }

  template <class Op, class E>
  long argminmax(E const &expr)
  {
    if (!expr.flat_size())
      throw types::ValueError("empty sequence");
    using elt_type = typename E::dtype;
    elt_type argminmax_value = Op::limit();
#ifndef USE_XSIMD
    if (utils::no_broadcast_ex(expr)) {
      return std::get<0>(_argminmax_fast<Op>(expr, argminmax_value, 0,
                                             utils::int_<E::value>()));
    } else
#endif
      return _argminmax<Op>(expr, argminmax_value, utils::int_<E::value>());
  }

  template <class Op, size_t Dim, size_t Axis, class T, class E, class V>
  void _argminmax_tail(T &&out, E const &expr, long curr, V &&curr_minmax,
                       std::integral_constant<size_t, 0>)
  {
    if (Op::value(expr, curr_minmax)) {
      out = curr;
      curr_minmax = expr;
    }
  }

  template <class Op, size_t Dim, size_t Axis, class T, class E, class V,
            size_t N>
  typename std::enable_if<Axis != (Dim - N), void>::type
  _argminmax_tail(T &&out, E const &expr, long curr, V &&curr_minmax,
                  std::integral_constant<size_t, N>)
  {
    static_assert(N >= 1, "specialization ok");
    long i = 0;
    for (auto &&elt : expr) {
      _argminmax_tail<Op, Dim, Axis>(out.fast(i), elt, curr,
                                     curr_minmax.fast(i),
                                     std::integral_constant<size_t, N - 1>());
      ++i;
    }
  }

  template <class Op, size_t Dim, size_t Axis, class T, class E>
  typename std::enable_if<Axis == (Dim - 1), void>::type
  _argminmax_head(T &&out, E const &expr, std::integral_constant<size_t, 1>)
  {
    typename E::dtype val = Op::limit();
    long i = 0;
    for (auto &&elt : expr)
      _argminmax_tail<Op, Dim, Axis>(out, elt, i++, val,
                                     std::integral_constant<size_t, 0>());
  }

  template <class Op, size_t Dim, size_t Axis, class T, class E, size_t N>
  typename std::enable_if<Axis == (Dim - N), void>::type
  _argminmax_head(T &&out, E const &expr, std::integral_constant<size_t, N>)
  {
    static_assert(N > 1, "specialization ok");
    types::ndarray<typename E::dtype, types::array<long, N - 1>> val{
        sutils::getshape(out), Op::limit()};
    long i = 0;
    for (auto &&elt : expr) {
      _argminmax_tail<Op, Dim, Axis>(out, elt, i++, val,
                                     std::integral_constant<size_t, N - 1>());
    }
  }

  template <class Op, size_t Dim, size_t Axis, class T, class E, size_t N>
  typename std::enable_if<Axis != (Dim - N), void>::type
  _argminmax_head(T &&out, E const &expr, std::integral_constant<size_t, N>)
  {
    static_assert(N >= 1, "specialization ok");
    auto out_iter = out.begin();
    for (auto &&elt : expr) {
      _argminmax_head<Op, Dim, Axis>(*out_iter, elt,
                                     std::integral_constant<size_t, N - 1>());
      ++out_iter;
    }
  }

  template <class Op, size_t N, class T, class E, size_t... Axis>
  void _argminmax_pick_axis(long axis, T &&out, E const &expr,
                            utils::index_sequence<Axis...>)
  {
    (void)std::initializer_list<bool>{
        ((Axis == axis) && (_argminmax_head<Op, N, Axis>(
                                out, expr, std::integral_constant<size_t, N>()),
                            true))...};
  }

  template <class Op, class E>
  types::ndarray<long, types::array<long, E::value - 1>>
  argminmax(E const &array, long axis)
  {
    if (axis < 0)
      axis += E::value;
    if (axis < 0 || size_t(axis) >= E::value)
      throw types::ValueError("axis out of bounds");
    auto shape = sutils::getshape(array);
    types::array<long, E::value - 1> shp;
    auto next = std::copy(shape.begin(), shape.begin() + axis, shp.begin());
    std::copy(shape.begin() + axis + 1, shape.end(), next);
    types::ndarray<long, types::array<long, E::value - 1>> out{shp,
                                                               builtins::None};
    _argminmax_pick_axis<Op, E::value>(axis, out, array,
                                       utils::make_index_sequence<E::value>());
    return out;
  }
}
PYTHONIC_NS_END

#endif