/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 * Copyright 2010-2018 Danny Robson <danny@nerdcruft.net>
 */

#pragma once

#include "iota.hpp"
#include "../concepts.hpp"
#include "../tuple/value.hpp"
#include "../variadic.hpp"
#include "../functor.hpp"

namespace cruft::iterator {
    ///////////////////////////////////////////////////////////////////////////
    namespace detail::zip {
        /// A container that holds multiple iterators and returns a tuple of
        /// their results when dereferenced.
        template <
            typename IteratorT,
            typename = std::make_index_sequence<std::tuple_size_v<IteratorT>>
        >
        class zipped_iterator;


        template <
            typename IteratorT,
            std::size_t ...IndexV
        >
        class zipped_iterator<
            IteratorT,
            std::index_sequence<IndexV...>
        > {
        public:
            // We can't declare ourselves as a forward_iterator because we're
            // unable to supply references to our value_type when we get
            // dereferenced unless we store the supplied values/references
            // inside ourself.
            //
            // This complicates implementation too much given we have no
            // pressing need for this functionality.
            using iterator_category = std::input_iterator_tag;
            using difference_type = std::ptrdiff_t;

            using value_type = std::tuple<
                typename std::iterator_traits<
                    std::tuple_element_t<IndexV, IteratorT>
                >::reference...
            >;

            using reference = value_type;
            using pointer = value_type*;

            zipped_iterator (IteratorT _iterators)
                : m_iterators (std::move (_iterators))
            { ; }


            zipped_iterator&
            operator++ (void)
            {
                (++std::get<IndexV> (m_iterators), ...);
                return *this;
            }


            zipped_iterator operator++ (int);


            // Do NOT use std::forward_as_tuple here. It favours rvalue
            // references which aren't useful after we've returned. decltype
            // gives us lvalues and lvalue-references.
            //
            // If you insist on using another approach then you should test
            // under a sanitizer (debug and release) before you commit the
            // changes.
            auto
            operator* (void)
            {
                return std::tuple<
                    decltype (*std::get<IndexV> (m_iterators))...
                > (
                    *std::get<IndexV> (m_iterators)...
                );
            }


            // Do NOT use std::forward_as_tuple here. It favours rvalue
            // references which aren't useful after we've returned. decltype
            // gives us lvalues and lvalue-references.
            //
            // If you insist on using another approach then you should test
            // under a sanitizer (debug and release) before you commit the
            // changes.
            auto
            operator* (void) const
            {
                return std::tuple<
                    std::add_const_t<
                        decltype(*std::get<IndexV> (m_iterators))
                    >...
                > (
                    *std::get<IndexV> (m_iterators)...
                );
            }


            bool
            operator== (zipped_iterator const &rhs) const
            {
                return m_iterators == rhs.m_iterators;
            }


            bool
            operator!= (zipped_iterator const &rhs) const
            {
                return !(*this == rhs);
            }


        private:
            IteratorT m_iterators;
        };


        template <typename TupleT>
        zipped_iterator (TupleT) -> zipped_iterator<TupleT>;


        /// A simple container for multiple collections that returns a
        /// wrapped multi-iterator for begin and end.
        ///
        /// It is up to the user to ensure StoreT does not contain dangling
        /// references.
        ///
        /// \tparam StoreT  A tuple-like object of collections (or references
        ///                 to collections).
        template <
            typename StoreT,
            typename = std::make_index_sequence<
                std::tuple_size_v<StoreT>
            >
        >
        class collection;

        template <
            typename StoreT,
            std::size_t ...IndexV
        > class collection<
            StoreT,
            std::index_sequence<IndexV...>
        > {
        public:
            collection (StoreT _store)
                : m_store (std::move (_store))
            { ; }


            auto begin (void)&
            {
                return zipped_iterator (
                    std::tuple (
                        std::begin (
                            std::get<IndexV> (m_store)
                        )...
                    )
                );
            }

            auto begin (void) const&
            {
                return zipped_iterator (
                    std::tuple (
                        std::begin (std::get<IndexV> (m_store))...
                    )
                );
            }


            auto end (void)&
            {
                return zipped_iterator (
                    std::tuple (
                        std::end (std::get<IndexV> (m_store))...
                    )
                );
            }


            auto end (void) const &
            {
                return zipped_iterator (
                    std::tuple (
                        std::end (std::get<IndexV> (m_store))...
                    )
                );
            }


            /// Returns the number of elements in the sequence.
            decltype (auto)
            size (void) const
            {
                // All stores should have the same size so we arbitrarily pick
                // the first to query.
                return std::size (std::get<0> (m_store));
            }

        private:
            StoreT m_store;
        };
    }


    ///------------------------------------------------------------------------
    /// Takes a variable number of container arguments and returns an interable
    /// object with a value_type that is a tuple of the each container's
    /// value_type.
    ///
    /// The returned iterator value_type is suitable for using in range-for
    /// and structured bindings (and really, that's the entire point here).
    ///
    /// eg, cruft::zip ({1,2,3}, {4,5,6}) ~= {{1,4},{2,5},{3,6}}
    template <typename ...ContainerT>
    requires (cruft::concepts::iterable<ContainerT> && ...)
    decltype(auto)
    zip (ContainerT&&... data)
    {
        CHECK (((std::size (data) == std::size (variadic::get<0> (data...))) && ...));

        return detail::zip::collection<std::tuple<ContainerT...>> (
            std::forward_as_tuple (
                std::forward<ContainerT> (data)...
            )
        );
    }


    ///------------------------------------------------------------------------
    /// Takes a variable number of containers and returns a zipped iterable
    /// object where the first of the iterator's value_types is the index of
    /// that iterator. ie, It combines container offsets with value_types.
    ///
    /// The supplied containers _must_ not change size while the izip object
    /// is live. The size of the containers may be cached for efficiency.
    ///
    /// eg, cruft::izip ("abc") ~= {{0,'a'},{1,'b'},{2,'c'}}
    template <typename ...ContainerT>
    requires (cruft::concepts::iterable<ContainerT> && ...)
    decltype(auto)
    izip (ContainerT&&... data)
    {
        // Assume that all containers are the same size and create a range
        // that will supply integral indices over the first container.
        //
        // The call to `zip` will peform more rigorous tests on the sizes
        // of the container collection.
        return zip (
            iota (
                std::size (::cruft::variadic::get<0> (data...))
            ),
            std::forward<ContainerT> (data)...
        );
    }
}