//
// Distributed under the ITensor Library License, Version 1.2
//    (See accompanying LICENSE file.)
//
#ifndef __ITENSOR_INDEXSET_H
#define __ITENSOR_INDEXSET_H
#include <algorithm>
#include "itensor/util/readwrite.h"
#include "itensor/util/safe_ptr.h"
#include "itensor/tensor/range.h"
#include "itensor/tensor/types.h"
#include "itensor/tensor/permutation.h"
#include "itensor/index.h"


namespace itensor {


template<class IndexT>
class IndexSetT;

class IQIndex;

//
// IndexSetT
//
// Aliases:
using IndexSet = IndexSetT<Index>;
using IQIndexSet = IndexSetT<IQIndex>;

using IndexSetBuilder = RangeBuilderT<IndexSet>;
using IQIndexSetBuilder = RangeBuilderT<IQIndexSet>;

//
// When constructed from a collection of indices,
// (as an explicit set of arguments or via
// a container) puts the indices with m>1 to the
// front and those with m==1 at the back but otherwise
// keeps the indices in the order given.
//

template<typename index_type_> 
class IndexSetIter;

template <class index_type_>
class IndexSetT : public RangeT<index_type_>
    {
    public:
    using index_type = index_type_;
    using extent_type = index_type;
    using range_type = RangeT<index_type>;
    using parent = RangeT<index_type>;
    using size_type = typename range_type::size_type;
    using storage_type = typename range_type::storage_type;
    using value_type = index_type;
    using iterator = IndexSetIter<index_type>;
    using const_iterator = IndexSetIter<const index_type>;
    using indexval_type = typename index_type::indexval_type;

    public:

    IndexSetT() { }

    // construct from 1 or more indices
    template <typename... Inds>
    explicit
    IndexSetT(index_type const& i1, 
              Inds&&... rest)
      : parent(i1,std::forward<Inds>(rest)...)
        { }

    template<typename IndxContainer>
    explicit
    IndexSetT(IndxContainer && ii) 
      : parent(std::forward<IndxContainer>(ii)) { }

    IndexSetT(std::initializer_list<index_type> ii) : parent(ii) { }

    explicit
    IndexSetT(storage_type && store) 
      : parent(std::move(store)) 
        { }

    IndexSetT&
    operator=(storage_type&& store)
        {
        parent::operator=(std::move(store));
        return *this;
        }
    
    explicit operator bool() const { return !parent::empty(); }

    long
    extent(size_type i) const { return parent::extent(i); }

    size_type
    stride(size_type i) const { return parent::stride(i); }

    long
    r() const { return parent::r(); }
    
    // 0-indexed access
    index_type &
    operator[](size_type i)
        { 
#ifdef DEBUG
        if(i >= parent::size()) Error("IndexSetT[i] arg out of range");
#endif
        return parent::index(i);
        }

    // 1-indexed access
    index_type &
    index(size_type I)
        { 
#ifdef DEBUG
        if(I < 1 || I > parent::size()) Error("IndexSetT.index(i) arg out of range");
#endif
        return operator[](I-1);
        }

    // 0-indexed access
    index_type const&
    operator[](size_type i) const
        { 
#ifdef DEBUG
        if(i >= parent::size()) Error("IndexSetT[i] arg out of range");
#endif
        return parent::index(i);
        }

    // 1-indexed access
    index_type const&
    index(size_type I) const
        { 
#ifdef DEBUG
        if(I < 1 || I > parent::size()) Error("IndexSetT.index(i) arg out of range");
#endif
        return operator[](I-1);
        }


    index_type const&
    front() const { return parent::front().ind; }

    index_type const&
    back() const { return parent::back().ind; }

    iterator
    begin() { return iterator{*this}; }

    iterator
    end() { return iterator::makeEnd(*this); }

    const_iterator
    begin() const { return const_iterator{*this}; }

    const_iterator
    end() const { return const_iterator::makeEnd(*this); }

    const_iterator
    cbegin() const { return begin(); }

    const_iterator
    cend() const { return end(); }

    parent const&
    range() const { return *this; }

    void
    swap(IndexSetT& other) { parent::swap(other); }

    void
    dag() { for(auto& J : *this) J.dag(); }

    void
    read(std::istream& s);

    void
    write(std::ostream& s) const;
    };

template<typename index_type>
auto
rangeBegin(IndexSetT<index_type> const& is) -> decltype(is.range().begin())
    {
    return is.range().begin();
    }

template<typename index_type>
auto
rangeEnd(IndexSetT<index_type> const& is) -> decltype(is.range().end())
    {
    return is.range().end();
    }

//
// IndexSetT Primelevel Methods
//

// increment primelevel of all indices of
// type "type" by an amount "inc"
template<typename IndexT, typename... Types>
void 
prime(IndexSetT<IndexT>& is, 
      IndexType type,
      int inc = 1);

// same as above but for multiple types
// optionally, last argument can be 
// an increment amount
template<typename IndexT, typename... Types>
void 
prime(IndexSetT<IndexT>& is, 
      IndexType type1,
      Types&&... rest);

// increment primelevel of all
// indices by an amount "inc"
template<typename IndexT>
void 
prime(IndexSetT<IndexT>& is, 
      int inc = 1);

// Increment primelevels of the indices
// specified by 1, or an optional amount "inc"
// For example, to prime indices I and J by 2,
// prime(is,I,J,2);
template<typename IndexT, typename... Inds>
void 
prime(IndexSetT<IndexT>& is, 
      IndexT const& I1, 
      Inds&&... rest);

//
//Given a list of indices and an increment (an int)
//as the optional last argument (default is inc=1)
//increment all indices NOT listed in the arguments
//by the amout inc.
//
//For example, primeExcept(is,I1,I3,I4,I7,2);
//will increment all prime levels by 2 except for
//those of I1,I3,I4, and I7.
//
template<typename IndexT, typename... Inds>
void 
primeExcept(IndexSetT<IndexT>& is, 
            IndexT const& I1, 
            Inds&&... inds);

template<typename IndexT, typename... ITs>
void 
primeExcept(IndexSetT<IndexT>& is, 
            IndexType it,
            ITs&&... etc);

template<typename IndexT>
void 
noprime(IndexSetT<IndexT>& is, IndexType type = All);

template<typename IndexT, typename... ITs>
void 
noprime(IndexSetT<IndexT>& is,
        IndexType it1,
        IndexType it2,
        ITs&&... rest);

template<typename IndexT, typename... Inds>
void 
noprime(IndexSetT<IndexT>& is, 
        const IndexT& I1, 
        Inds&&... inds);

// This version of mapprime takes
// any number of triples: I,p1,p2
// where I is an index or an IndexType,
// p1 is the inital primelevel and
// p2 if the final primlevel
// For example,
// prime(is,j,0,2,Site,1,0);
// would change an Index "j" with
// primelevel 0 to have primelevel 2
// and any indices with type Site
// and primelevel 1 to have primelevel 0
// If two or more mappings match for a particular
// index, only the first mapping is used.
template<typename IndexT, typename... VArgs>
void 
mapprime(IndexSetT<IndexT>& is, 
         VArgs&&... vargs);

template<typename IndexT>
void 
mapprime(IndexSetT<IndexT>& is, 
         int plevold, 
         int plevnew, 
         IndexType type = All);

//
// IndexSetT helper methods
//


template<class IndexT>
Arrow
dir(const IndexSetT<IndexT>& is, const IndexT& I);


template <class IndexT>
IndexT const&
finddir(IndexSetT<IndexT> const& iset, Arrow dir);

//
// Given IndexSetT<IndexT> iset and IndexT I,
// return int j such that iset[j] == I.
// If not found, returns -1
//
template <class IndexT>
long
findindex(IndexSetT<IndexT> const& iset, 
          IndexT const& I);

template <class IndexT>
IndexT const&
findtype(IndexSetT<IndexT> const& iset, 
         IndexType t);

//
// Compute the permutation P taking an IndexSetT iset
// to oset (of type IndexSetT or array<IndexT,NMAX>)
//
template <class IndexT>
void
getperm(const IndexSetT<IndexT>& iset, 
        const typename IndexSetT<IndexT>::storage& oset, 
        Permutation& P);

template <class IndexT>
bool
hasindex(const IndexSetT<IndexT>& iset, 
         const IndexT& I);

template <class IndexT>
bool
hastype(const IndexSetT<IndexT>& iset, 
        IndexType t);

template <class IndexT>
long
minM(const IndexSetT<IndexT>& iset);

template <class IndexT>
long
maxM(const IndexSetT<IndexT>& iset);

template<class IndexT>
void
contractIS(IndexSetT<IndexT> const& Lis,
           IndexSetT<IndexT> const& Ris,
           IndexSetT<IndexT> & Nis,
           bool sortResult = false);

template<class IndexT, class LabelT>
void
contractIS(IndexSetT<IndexT> const& Lis,
           LabelT const& Lind,
           IndexSetT<IndexT> const& Ris,
           LabelT const& Rind,
           IndexSetT<IndexT> & Nis,
           LabelT & Nind,
           bool sortResult = false);

template<class IndexT, class LabelT>
void
ncprod(IndexSetT<IndexT> const& Lis,
       LabelT const& Lind,
       IndexSetT<IndexT> const& Ris,
       LabelT const& Rind,
       IndexSetT<IndexT> & Nis,
       LabelT & Nind);

template <class IndexT>
std::ostream&
operator<<(std::ostream& s, IndexSetT<IndexT> const& is);

template<typename index_type_> 
class IndexSetIter
    { 
    public:
    using index_type = stdx::remove_const_t<index_type_>;
    using value_type = index_type;
    using reference = index_type_&;
    using difference_type = std::ptrdiff_t;
    using pointer = index_type_*;
    using iterator_category = std::random_access_iterator_tag;
    using indexset_type = stdx::conditional_t<std::is_const<index_type_>::value,
                                             const IndexSetT<index_type>,
                                             IndexSetT<index_type>>;
    using range_ptr = typename RangeT<index_type>::value_type*;
    using const_range_ptr = const typename RangeT<index_type>::value_type*;
    using data_ptr = stdx::conditional_t<std::is_const<index_type_>::value,
                                      const_range_ptr,
                                      range_ptr>;
    private:
    size_t off_ = 0;
    indexset_type* p_; 
    public: 

    IndexSetIter() : p_(nullptr) { }

    explicit
    IndexSetIter(indexset_type & is) : p_(&is) { }

    size_t
    offset() const { return off_; }

    IndexSetIter& 
    operator++() 
        { 
        ++off_; 
        return *this; 
        } 

    IndexSetIter 
    operator++(int) 
        { 
        auto tmp = *this; //save copy of this
        ++off_; 
        return tmp; 
        } 

    IndexSetIter& 
    operator+=(difference_type x) 
        { 
        off_ += x;
        return *this; 
        } 

    IndexSetIter& 
    operator--( ) 
        { 
        --off_;
        return *this; 
        } 

    IndexSetIter 
    operator--(int) 
        { 
        auto tmp = *this; //save copy of this
        --off_;
        return tmp; 
        } 

    IndexSetIter& 
    operator-=(difference_type x) 
        { 
        off_ -= x;
        return *this; 
        } 

    reference 
    operator[](difference_type n) { return p_->operator[](n); } 

    reference 
    operator*() { return p_->operator[](off_); }  

    pointer 
    operator->() { return &(p_->operator[](off_)); }

    IndexSetIter static
    makeEnd(indexset_type & is)
        {
        IndexSetIter end;
        end.p_ = &is;
        end.off_ = is.size();
        return end;
        }
    }; 

template <typename T>
bool 
operator==(const IndexSetIter<T>& x, const IndexSetIter<T>& y) 
    { 
    return x.offset() == y.offset(); 
    } 

template <typename T>
bool 
operator!=(const IndexSetIter<T>& x, const IndexSetIter<T>& y) 
    { 
    return x.offset() != y.offset(); 
    } 

template <typename T>
bool 
operator<(const IndexSetIter<T>& x, const IndexSetIter<T>& y) 
    { 
    return x.offset() < y.offset(); 
    } 

template <typename T>
IndexSetIter<T>
operator+(IndexSetIter<T> x, 
          typename IndexSetIter<T>::difference_type d) 
    { 
    return x += d;
    } 

template <typename T>
IndexSetIter<T>
operator+(typename IndexSetIter<T>::difference_type d, 
          IndexSetIter<T> x) 
    { 
    return x += d;
    } 

} //namespace itensor

#include "itensor/indexset.ih"

#endif