#pragma once
#include <set>
#include <queue>
#include <cassert>

template <class T, class Compare> class ConstSafeMultisetAdaptor;
template <class T, class Compare> class SafeMultisetAdaptor;

/**
 * @brief Base class for safe multiset, whose only purpose is to declare the return structures
 *        so that they are accessible without going through the templated safe multiset class.
 */
class SafeMultisetBase
{
public:
  struct ConstIterationStep
  {
    ConstIterationStep()
      : breakIteration(false) {}
    ConstIterationStep(bool _breakIteration)
      : breakIteration(_breakIteration) {}
    bool breakIteration;
  };

  struct IterationStep : ConstIterationStep
  {
    IterationStep()
      : ConstIterationStep(), deleteCurrent(false) {}
    IterationStep(bool _deleteCurrent, bool _breakIteration)
      : ConstIterationStep(_breakIteration), deleteCurrent(_deleteCurrent) {}
    IterationStep(ConstIterationStep r)
      : ConstIterationStep(r), deleteCurrent(false) {}
    bool deleteCurrent;
  };
};

/**
 * @brief A safe multiset is like a regular multiset storing sorted values, but in addition it allows one
 *        to iterate over its elements while some are being deleted and new ones are being added.
 *        Multiple iterations can be performed on the safe multiset at the same time. To allow this,
 *        deleted elements are kept in the multiset until the last iteration has finished, which is
 *        when they are removed.
 */
template <class T, class Compare = std::less<T>>
class SafeMultiset : public SafeMultisetBase
{
public:
  /**
   * @brief Wraps an element of the multiset by including other data, in particular 
   *        a flag indicating whether this element has been deleted.
   */
  class Wrapper
  {
  public:
    Wrapper(T t)
      : m_t(t), m_deleted(false) {}
    /**
     * @return Whether this element has been deleted.
     */
    bool isDeleted() const { return m_deleted; }
    operator T&() { return m_t; }
    operator const T&() const { return m_t; }
  protected:
    friend class SafeMultiset;
    /**
     * @brief The wrapped element.
     */
    T m_t;
    /**
     * @brief Whether the element has been deleted.
     */
    bool m_deleted;
  };

protected:
  friend class ConstSafeMultisetAdaptor<T, Compare>;
  friend class SafeMultisetAdaptor<T, Compare>;

  /**
   * @brief Trivial functor comparing two wrappers.
   */
  class CompareWrapper
  {
    CompareWrapper & operator=(const CompareWrapper &);
  public:
    CompareWrapper(const Compare& compare)
      : m_compare(compare) {}
    bool operator()(const Wrapper& wrapper1, const Wrapper& wrapper2) const
    {
      return m_compare(wrapper1.m_t, wrapper2.m_t); 
    }
  private:
    const Compare& m_compare;
  };
  typedef typename std::multiset<Wrapper, CompareWrapper>::const_iterator internal_const_iterator;
  typedef typename std::multiset<Wrapper, CompareWrapper>::iterator internal_iterator;
public:
  class iterator
  {
  public:
    iterator()
    {}
    iterator(internal_iterator it)
      : m_it(it)
    {}
    typename internal_iterator::reference operator*() const
    {
      return *m_it;
    }
    typename internal_iterator::reference operator->() const
    {
      return *m_it;
    }
    bool operator==(const iterator& it) const
    {
      return m_it == it.m_it;
    }
  private:
    friend class SafeMultiset;

    internal_iterator m_it;
  };

  /**
   * @brief Safe multiset constructor.
   * @param compare The comparator used to compare elements.
   */
  SafeMultiset(const Compare& compare = Compare())
    : m_compare(compare)
    , m_multiset(CompareWrapper(m_compare))
    , m_size(0)
    , m_first(NULL)
  {
  }

  /**
   * @brief Inserts a new element in the safe multiset. The element is also inserted (if necessary) 
   *        in each pending priority queue.
   * @param t The new element to insert.
   * @return The position of the new element in the multiset.
   */
  iterator insert(const T& t)
  {
    ++m_size;
    internal_iterator it = m_multiset.insert(t);
    for (Iteration* iteration = m_first; iteration != NULL; iteration = iteration->_getNext())
    {
      iteration->pushNewElement(&(*it));
    }
    return iterator(it);
  }
 
  /**
   * @brief Iterates on all the elements in the safe multiset, including potential new elements added up during the iteration.
   *        The elements are iterated over in the order specified by the comparator.
   * @tparam Function The type of function to apply to the elements. The function must take an element T as
   *                  unique parameter and return a IterationStep result;
   * @param function The function to apply to the elements. The function returns whether to break the iteration and
   *                 whether the current element iterated over should be deleted.
   * @return A copy of the function applied to the elements.
   */
  template <class Function>
  Function do_while(Function function)
  {
    Iteration iteration(*this);

    bool breakIteration = false;
    for (internal_iterator it = m_multiset.begin(); !breakIteration && it != m_multiset.end(); ++it)
    {
      Wrapper& current = *it;
      if (!current.m_deleted)
      {
        iteration.setElementsAdditionThreshold(&current);
        IterationStep result = function(current);
        if (result.deleteCurrent)
        {
          current.m_deleted = true;
          --m_size;
        }
        bool breakIteration = result.breakIteration;
        // Iterate over newly created higher priority elements if any
        while (!breakIteration && iteration.hasNewElements())
        {
          Wrapper* wrapper = iteration.popNewElement();
          assert(CompareWrapper(m_compare)(*wrapper, current));
          if (!wrapper->m_deleted)
          {
            IterationStep result = function(*wrapper);
            if (result.deleteCurrent)
            {
              wrapper->m_deleted = true;
              --m_size;
            }
            breakIteration = result.breakIteration;
          }
        }
      }
    }
    return function;
  }

  /**
   * @brief Iterates on all the elements in the safe multiset, including potential new elements added up during the iteration.
   *        The elements are iterated over in the order specified by the comparator.
   * @tparam Function The type of function to apply to the elements. The function must take an element T as
   *                  unique parameter and return a ConstIterationStep result;
   * @param function The function to apply to the elements. The function returns whether to break the iteration.
   * @return A copy of the function applied to the elements.
   */
  template <class Function>
  Function do_while(Function function) const
  {
    Iteration iteration(*this);

    bool breakIteration = false;
    for (internal_const_iterator it = m_multiset.begin(); !breakIteration && it != m_multiset.end(); ++it)
    {
      const Wrapper& current = *it;
      if (!current.m_deleted)
      {
        iteration.setElementsAdditionThreshold(&current);
        ConstIterationStep result = function(current);
        bool breakIteration = result.breakIteration;
        // Iterate over newly created higher priority elements if any
        while (!breakIteration && iteration.hasNewElements())
        {
          Wrapper* wrapper = iteration.popNewElement();
          assert(CompareWrapper(m_compare)(*wrapper, current));
          if (!wrapper->m_deleted)
          {
            ConstIterationStep result = function(*wrapper);
            breakIteration = result.breakIteration;
          }
        }
      }
    }
    return function;
  }

  /**
   * @brief Removes the element at the given position from the multiset.
   * @param position Iterator to the element to remove.
   */
  void erase(iterator position)
  {
    assert(!position.m_it->m_deleted);
    position.m_it->m_deleted = true;
    --m_size;
    if (m_first == NULL)
    {
      m_multiset.erase(position.m_it);
    }
  }

  /**
   * @return The number of elements present in the multiset.
   */
  size_t size() const
  {
    assert(m_first != NULL || m_size == m_multiset.size());
    return m_size;
  }

  /**
   * @return Whether the multiset is empty.
   */
  bool empty() const
  {
    return size() == 0;
  }

protected:
  SafeMultiset(const SafeMultiset& other);
  SafeMultiset& operator=(const SafeMultiset& other) const;

  /**
   * @brief Iteration objects provide an automatic (and exception-aware, although there's no need to be) mechanism for acquisition and release.
   *        When new elements are added to the multiset while one or more iterations are taking place, 
   *        these new elements above a certain threshold (the ones inserted before the current element iterated over) 
   *        are added to a priority queue, in order to iterate over them afterwards.
   *        Additionally, all iteration objects associated with the same safe multiset are linked to one another to create
   *        a linked list of iterations.
   */
  class Iteration
  {
  public:
    Iteration(const SafeMultiset& safeMultiset)
      : m_safeMultiset(safeMultiset)
      , m_compare(safeMultiset.m_compare)
      , m_newElements(ComparePtrWrapper(m_compare))
      , m_threshold(NULL)
      , m_next(NULL)
    {
      m_safeMultiset.iterationStarted(*this);
    }
    ~Iteration()
    {
      m_safeMultiset.iterationEnded(*this);
    }
    /**
     * @return The iteration object placed after this one in the linked list, possibly NULL.
     */
    Iteration* _getNext()
    {
      return m_next;
    }
    /**
     * @brief Links this iteration object with another one.
     * @param next The iteration object placed next in the linked list.
     */
    void setNext(Iteration* next)
    {
      m_next = next;
    }
    /**
     * @brief Sets the threshold below which elements are not stored.
     * @param threshold The threshold to store elements.
     */
    void setElementsAdditionThreshold(const Wrapper* threshold)
    {
      m_threshold = threshold;
    }
    /**
     * @return Whether new elements have been inserted in the priority queue.
     */
    bool hasNewElements() const
    {
      return !m_newElements.empty();
    }
    /**
     * @brief Pushes an element in the priority queue.
     * @param element The element to push in the priority queue.
     * @return Whether the element was added, which is true iff the element was above the threshold.
     */
    bool pushNewElement(Wrapper* element)
    {
      assert(m_threshold != NULL);
      bool add = m_compare(*element, *m_threshold); // if new >= current, the new element will be inserted after the current one
      if (add)
        m_newElements.push(element);
      return add;
    }
    /**
     * @brief Pops the highest element from the priority queue.
     * @return The highest element in the priority queue.
     */
    Wrapper* popNewElement()
    {
      Wrapper* element = m_newElements.top();
      m_newElements.pop();
      return element;
    }
  private:
    Iteration(const Iteration&);
    Iteration& operator=(const Iteration&);

  protected:
    class ComparePtrWrapper
    {
    public:
      /**
       * @brief Trivial functor comparing two wrapper pointers.
       */
      ComparePtrWrapper(const CompareWrapper& compare)
        : m_compare(compare) {}
      bool operator()(const Wrapper* wrapper1, const Wrapper* wrapper2) const
      {
        return !m_compare(*wrapper1, *wrapper2); 
      }
    private:
      const CompareWrapper& m_compare;

      ComparePtrWrapper & operator=(const ComparePtrWrapper &);
    };

    /**
     * @brief The safe multiset associated with this iteration object.
     */
    const SafeMultiset& m_safeMultiset;
    /**
     * @brief The comparator used to compare the elements.
     */
    const CompareWrapper m_compare;
    /**
     * @brief An inner priority queue storing the elements (above the threshold) which were
     *        added during the iteration.
     */
    std::priority_queue<Wrapper*, std::vector<Wrapper*>, ComparePtrWrapper> m_newElements;
    /**
     * @brief The threshold above which elements are added.
     */
    const Wrapper* m_threshold;
    /**
     * @brief A pointer on the next priority queue in the linked list of priority queues, possibly NULL.
     */
    Iteration* m_next;
  };
  friend class Iteration;

  /**
   * @brief This method is called by Iteration instances to signal that another iteration
   *        over the safe multiset is being performed.
   * @param iteration The new iteration being performed.
   */
  void iterationStarted(Iteration& iteration) const
  {
    iteration.setNext(m_first);
    m_first = &iteration;
  }

  /**
   * @brief This method is called by Iteration instances to signal that an iteration over
   *        the safe multiset has ended. If this was the last iteration, elements marked as 'deleted' 
   *        are removed.
   * @param iteration The iteration having ended.
   */
  void iterationEnded(Iteration& iteration) const
  {
    m_first = iteration._getNext();
    if (m_first == NULL)
    {
      internal_iterator it = m_multiset.begin();
      while (it != m_multiset.end())
      {
        internal_iterator next = it;
        ++next;
        if ((*it).m_deleted)
        {
          m_multiset.erase(it);
        }
        it = next;
      }
    }
  }

  /**
   * @brief The comparator used to compare the elements.
   */
  const Compare m_compare;
  /**
   * @brief The inner multiset which stores the elements.
   */
  mutable std::multiset<Wrapper, CompareWrapper> m_multiset;
  /**
   * @brief The beginning of a linked list of priority queues, one per iteration currently
   *        being performed on the safe multiset.
   */
  mutable Iteration* m_first;
  /**
   * @brief The number of elements in the safe multiset, excluding deleted ones.
   *        Note that this number should be equal to the number of elements in m_multiset, when no iterations
   *        are being performed on the safe multiset.
   */
  size_t m_size;
};
