// Copyright 2024 The Manifold Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once

#include <iterator>
#include <type_traits>

namespace manifold {

template <typename F, typename Iter>
struct TransformIterator {
 private:
  Iter iter;
  F f;

 public:
  using pointer = void;
  using reference = std::invoke_result_t<
      F, typename std::iterator_traits<std::remove_const_t<Iter>>::value_type>;
  using difference_type =
      typename std::iterator_traits<std::remove_const_t<Iter>>::difference_type;
  using value_type = reference;
  using iterator_category = typename std::iterator_traits<
      std::remove_const_t<Iter>>::iterator_category;

  constexpr TransformIterator(Iter iter, F f) : iter(iter), f(f) {}

  TransformIterator& operator=(const TransformIterator& other) {
    if (this == &other) return *this;
    // don't copy function, should be the same
    iter = other.iter;
    return *this;
  }

  constexpr reference operator*() const { return f(*iter); }

  constexpr reference operator[](size_t i) const { return f(iter[i]); }

  // prefix increment
  TransformIterator& operator++() {
    iter += 1;
    return *this;
  }

  // postfix
  TransformIterator operator++(int) {
    auto old = *this;
    operator++();
    return old;
  }

  // prefix increment
  TransformIterator& operator--() {
    iter -= 1;
    return *this;
  }

  // postfix
  TransformIterator operator--(int) {
    auto old = *this;
    operator--();
    return old;
  }

  constexpr TransformIterator operator+(size_t n) const {
    return TransformIterator(iter + n, f);
  }

  TransformIterator& operator+=(size_t n) {
    iter += n;
    return *this;
  }

  constexpr TransformIterator operator-(size_t n) const {
    return TransformIterator(iter - n, f);
  }

  TransformIterator& operator-=(size_t n) {
    iter -= n;
    return *this;
  }

  constexpr bool operator==(TransformIterator other) const {
    return iter == other.iter;
  }

  constexpr bool operator!=(TransformIterator other) const {
    return !(iter == other.iter);
  }

  constexpr bool operator<(TransformIterator other) const {
    return iter < other.iter;
  }

  constexpr difference_type operator-(TransformIterator other) const {
    return iter - other.iter;
  }

  constexpr operator TransformIterator<F, const Iter>() const {
    return TransformIterator(f, iter);
  }
};

template <typename T>
struct CountingIterator {
 private:
  T counter;

 public:
  using pointer = void;
  using reference = T;
  using difference_type = std::make_signed_t<T>;
  using value_type = T;
  using iterator_category = std::random_access_iterator_tag;

  constexpr CountingIterator(T counter) : counter(counter) {}

  constexpr value_type operator*() const { return counter; }
  constexpr value_type operator[](T i) const { return counter + i; }

  // prefix increment
  CountingIterator& operator++() {
    counter += 1;
    return *this;
  }

  // postfix
  CountingIterator operator++(int) {
    auto old = *this;
    operator++();
    return old;
  }

  // prefix increment
  CountingIterator& operator--() {
    counter -= 1;
    return *this;
  }

  // postfix
  CountingIterator operator--(int) {
    auto old = *this;
    operator--();
    return old;
  }

  constexpr CountingIterator operator+(T n) const {
    return CountingIterator(counter + n);
  }

  CountingIterator& operator+=(T n) {
    counter += n;
    return *this;
  }

  constexpr CountingIterator operator-(T n) const {
    return CountingIterator(counter - n);
  }

  CountingIterator& operator-=(T n) {
    counter -= n;
    return *this;
  }

  constexpr friend bool operator==(CountingIterator a, CountingIterator b) {
    return a.counter == b.counter;
  }

  constexpr friend bool operator!=(CountingIterator a, CountingIterator b) {
    return a.counter != b.counter;
  }

  constexpr friend bool operator<(CountingIterator a, CountingIterator b) {
    return a.counter < b.counter;
  }

  constexpr friend difference_type operator-(CountingIterator a,
                                             CountingIterator b) {
    return a.counter - b.counter;
  }

  constexpr operator CountingIterator<const T>() const {
    return CountingIterator(counter);
  }
};

constexpr CountingIterator<size_t> countAt(size_t i) {
  return CountingIterator(i);
}

template <typename Iter>
struct StridedRange {
 private:
  struct StridedRangeIter {
   private:
    Iter iter;
    size_t stride;

   public:
    using pointer =
        typename std::iterator_traits<std::remove_const_t<Iter>>::pointer;
    using reference =
        typename std::iterator_traits<std::remove_const_t<Iter>>::reference;
    using difference_type = typename std::iterator_traits<
        std::remove_const_t<Iter>>::difference_type;
    using value_type =
        typename std::iterator_traits<std::remove_const_t<Iter>>::value_type;
    using iterator_category = typename std::iterator_traits<
        std::remove_const_t<Iter>>::iterator_category;

    constexpr StridedRangeIter(Iter iter, int stride)
        : iter(iter), stride(stride) {}

    constexpr reference operator*() { return *iter; }

    constexpr std::add_const_t<reference> operator*() const { return *iter; }

    constexpr reference operator[](size_t i) { return iter[i * stride]; }

    constexpr std::add_const_t<reference> operator[](size_t i) const {
      return iter[i * stride];
    }

    // prefix increment
    StridedRangeIter& operator++() {
      iter += stride;
      return *this;
    }

    // postfix
    StridedRangeIter operator++(int) {
      auto old = *this;
      operator++();
      return old;
    }

    // prefix increment
    StridedRangeIter& operator--() {
      iter -= stride;
      return *this;
    }

    // postfix
    StridedRangeIter operator--(int) {
      auto old = *this;
      operator--();
      return old;
    }

    constexpr StridedRangeIter operator+(size_t n) const {
      return StridedRangeIter(iter + n * stride, stride);
    }

    StridedRangeIter& operator+=(size_t n) {
      iter += n * stride;
      return *this;
    }

    constexpr StridedRangeIter operator-(size_t n) const {
      return StridedRangeIter(iter - n * stride, stride);
    }

    StridedRangeIter& operator-=(size_t n) {
      iter -= n * stride;
      return *this;
    }

    constexpr friend bool operator==(StridedRangeIter a, StridedRangeIter b) {
      return a.iter == b.iter;
    }

    constexpr friend bool operator!=(StridedRangeIter a, StridedRangeIter b) {
      return !(a.iter == b.iter);
    }

    constexpr friend bool operator<(StridedRangeIter a, StridedRangeIter b) {
      return a.iter < b.iter;
    }

    constexpr friend difference_type operator-(StridedRangeIter a,
                                               StridedRangeIter b) {
      // note that this is not well-defined if a.stride != b.stride...
      return (a.iter - b.iter) / a.stride;
    }
  };
  Iter _start, _end;
  const size_t stride;

 public:
  constexpr StridedRange(Iter start, Iter end, size_t stride)
      : _start(start), _end(end), stride(stride) {}

  constexpr StridedRangeIter begin() const {
    return StridedRangeIter(_start, stride);
  }

  constexpr StridedRangeIter end() const {
    return StridedRangeIter(_start, stride) +
           ((std::distance(_start, _end) + (stride - 1)) / stride);
  }
};

}  // namespace manifold