#pragma once

#include <mutex>
#include <shared_mutex>
#include <thread>
#include <tuple>
#include <unordered_map>

#include <dune/common/hash.hh>
#include <dune/common/std/type_traits.hh>

namespace AMDiS
{
  /// Store cache in instance.
  template <class Container>
  struct ConsecutivePolicy;

  /// Store cache thread local, requires no locking.
  template <class Container>
  struct ThreadLocalPolicy;

  /// Stores cache global static, requires locking on write access.
  template <class Container>
  struct StaticLockedPolicy;


  /// \brief The class template ConcurrentCache describes an associative static container that allows the
  /// concurrent access to the stored data.
  /**
   * Cache data of arbitray type that needs initialization on the first access. The data is thereby
   * initialized thread-wise or globally only once, and guarantees that you always get initialized data.
   *
   * \tparam Key        The type of key to access the data.
   * \tparam Data       The type of the data stored in the cache. The behaviur is undefined if Data is not
   *                    the same type as Container::mapped_type.
   * \tparam Policy     A policy class template implementing the method `get_or_init()`. Three implementations
   *                    are provided: \ref ConsecutivePolicy, \ref ThreadLocalPolicy and \ref StaticLockedPolicy.
   *                    By default, if no policy class template is specified, the `ThreadLocalPolicy` is used.
   *                    \see ConcurrentCachePolicy
   * \tparam Container  The type of the underlying associative container to use to store the data. The
   *                    container must satisfy the requirements of AssociativeContainer. The standard
   *                    containers `std::map` and `std::unordered_map` satisfie this requirement. By default,
   *                    if not container class is specified, the standard container `std::unordered_map<Key,Data>`
   *                    is used. Note, an unordered_map requires the key to be hashable.
   *
   * The `Policy` class template is a template parametrizable with the container type, that provides a static `get_or_init()`
   * method that is called with the key, and a functor for creation of new data elements.
   **/
  template <class Key,
            class Data,
            template <class> class Policy = ThreadLocalPolicy,
            class Container = std::unordered_map<Key, Data>>
  class ConcurrentCache;


#ifdef DOXYGEN
  /// \brief The class template ConcurrentCachePolicy describes a concrete policies for the use in \ref ConcurrentCache.
  /**
   * Provide a static cache and a `get_or_init()` static method that extracts the data from the cache if it exists or
   * creates a new extry by using an initialization functor.
   *
   * Realizations of this template are \ref ConsecutivePolicy, \ref ThreadLocalPolicy and \ref StaticLockedPolicy.
   *
   * \tparam Container  The Type of the associative container key->data to store the cached data.
   **/
  template <class Container>
  class ConcurrentCachePolicy;
#endif

  // implementation of the consecutive policy. Data is stored in instance variable.
  template <class Container>
  struct ConsecutivePolicy
  {
    using key_type = typename Container::key_type;
    using data_type = typename Container::mapped_type;
    using container_type = Container;

    template <class F, class... Args>
    data_type const& get_or_init(key_type const& key, F&& f, Args&&... args) const
    {
      return impl(std::is_default_constructible<data_type>{}, key, FWD(f), FWD(args)...);
    }

  private:
    // data_type is default_constructible
    template <class F, class... Args>
    data_type const& impl(std::true_type, key_type const& key, F&& f, Args&&... args) const
    {
      data_type empty;
      auto it = cachedData_.emplace(key, std::move(empty));
      if (it.second) {
        data_type data = f(key, FWD(args)...);
        it.first->second = std::move(data);
      }
      return it.first->second;
    }

    // data_type is not default_constructible
    template <class F, class... Args>
    data_type const& impl(std::false_type, key_type const& key, F&& f, Args&&... args) const
    {
      auto it = cachedData_.find(key);
      if (it != cachedData_.end())
        return it->second;
      else {
        data_type data = f(key, FWD(args)...);
        auto it = cachedData_.emplace(key, std::move(data));
        return it.first->second;
      }
    }

    mutable container_type cachedData_;
  };


  // implementation of the ThreadLocal policy. Data is stored in thread_local variable.
  template <class Container>
  struct ThreadLocalPolicy
  {
    using key_type = typename Container::key_type;
    using data_type = typename Container::mapped_type;
    using container_type = Container;

    template <class F, class... Args>
    static data_type const& get_or_init(key_type const& key, F&& f, Args&&... args)
    {
      return impl(std::is_default_constructible<data_type>{}, key, FWD(f), FWD(args)...);
    }

  private:
    // data_type is default_constructible
    template <class F, class... Args>
    static data_type const& impl(std::true_type, key_type const& key, F&& f, Args&&... args)
    {
      // Container to store the cached values
      thread_local container_type cached_data;

      data_type empty;
      auto it = cached_data.emplace(key, std::move(empty));
      if (it.second) {
        data_type data = f(key, FWD(args)...);
        it.first->second = std::move(data);
      }
      return it.first->second;
    }

    // data_type is not default_constructible
    template <class F, class... Args>
    static data_type const& impl(std::false_type, key_type const& key, F&& f, Args&&... args)
    {
      // Container to store the cached values
      thread_local container_type cached_data;

      auto it = cached_data.find(key);
      if (it != cached_data.end())
        return it->second;
      else {
        data_type data = f(key, FWD(args)...);
        auto it = cached_data.emplace(key, std::move(data));
        return it.first->second;
      }
    }
  };


  // implementation of the Shared policy. Data is stored in static variable.
  template <class Container>
  struct StaticLockedPolicy
  {
    using key_type = typename Container::key_type;
    using data_type = typename Container::mapped_type;
    using container_type = Container;

    template <class F, class... Args>
    static data_type const& get_or_init(key_type const& key, F&& f, Args&&... args)
    {
      // Container to store the cached values
      static container_type cached_data;

      // mutex used to access the data in the container, necessary since
      // access emplace is read-write.
      using mutex_type = std::shared_timed_mutex;
      static mutex_type access_mutex;

      // first try to lock for read-only, if an element for key is found, return it,
      // if not, obtain a unique_lock to insert a new element and initialize it.
      std::shared_lock<mutex_type> read_lock(access_mutex);
      auto it = cached_data.find(key);
      if (it != cached_data.end())
        return it->second;
      else {
        read_lock.unlock();
        data_type data = f(key, FWD(args)...);
        std::unique_lock<mutex_type> write_lock(access_mutex);
        auto new_it = cached_data.emplace(key, std::move(data));
        return new_it.first->second;
      }
    }
  };


  template <class Key, class Data, template <class> class Policy, class Container>
  class ConcurrentCache
      : protected Policy<Container>
  {
    using key_type = Key;
    using data_type = Data;

  public:

    /// \brief Return the data associated to the `key`.
    /**
     * Return the data associated to key. If no data is found, create a new entry in the container
     * with a value obtained from the functor, by calling `f(key, args...)`.
     *
     * \param f        A functor of signature data_type(key_type, Args...)
     * \param args...  Arguments passed additionally to the functor f
     **/
    template <class F, class... Args>
    data_type const& get(key_type const& key, F&& f, Args&&... args) const
    {
      static_assert(Dune::Std::is_callable<F(key_type, Args...), data_type>::value,
        "Functor F must have the signature data_type(key_type, Args...)");

      return ConcurrentCache::get_or_init(key, FWD(f), FWD(args)...);
    }
  };

} // end namespace AMDiS