#ifndef HeterogeneousCore_AlpakaInterface_interface_CachingAllocator_h
#define HeterogeneousCore_AlpakaInterface_interface_CachingAllocator_h

#include <cassert>
#include <exception>
#include <iomanip>
#include <iostream>
#include <map>
#include <mutex>
#include <optional>
#include <sstream>
#include <string>
#include <tuple>
#include <type_traits>

#include <alpaka/alpaka.hpp>

#include "HeterogeneousCore/AlpakaInterface/interface/devices.h"
#include "HeterogeneousCore/AlpakaInterface/interface/AllocatorConfig.h"
#include "HeterogeneousCore/AlpakaInterface/interface/AlpakaServiceFwd.h"

// Inspired by cub::CachingDeviceAllocator

namespace cms::alpakatools {

  namespace detail {

    inline constexpr unsigned int power(unsigned int base, unsigned int exponent) {
      unsigned int power = 1;
      while (exponent > 0) {
        if (exponent & 1) {
          power = power * base;
        }
        base = base * base;
        exponent = exponent >> 1;
      }
      return power;
    }

    // format a memory size in B/KiB/MiB/GiB/TiB
    inline std::string as_bytes(size_t value) {
      if (value == std::numeric_limits<size_t>::max()) {
        return "unlimited";
      } else if (value >= (1ul << 40) and value % (1ul << 40) == 0) {
        return std::to_string(value >> 40) + " TiB";
      } else if (value >= (1ul << 30) and value % (1ul << 30) == 0) {
        return std::to_string(value >> 30) + " GiB";
      } else if (value >= (1ul << 20) and value % (1ul << 20) == 0) {
        return std::to_string(value >> 20) + " MiB";
      } else if (value >= (1ul << 10) and value % (1ul << 10) == 0) {
        return std::to_string(value >> 10) + " KiB";
      } else {
        return std::to_string(value) + "   B";
      }
    }

  }  // namespace detail

  /*
   * The "memory device" identifies the memory space, i.e. the device where the memory is allocated.
   * A caching allocator object is associated to a single memory `Device`, set at construction time, and unchanged for
   * the lifetime of the allocator.
   *
   * Each allocation is associated to an event on a queue, that identifies the "synchronisation device" according to
   * which the synchronisation occurs.
   * The `Event` type depends only on the synchronisation `Device` type.
   * The `Queue` type depends on the synchronisation `Device` type and the queue properties, either `Sync` or `Async`.
   *
   * **Note**: how to handle different queue and event types in a single allocator ?  store and access type-punned
   * queues and events ?  or template the internal structures on them, but with a common base class ?
   * alpaka does rely on the compile-time type for dispatch.
   *
   * Common use case #1: accelerator's memory allocations
   *   - the "memory device" is the accelerator device (e.g. a GPU);
   *   - the "synchronisation device" is the same accelerator device;
   *   - the `Queue` type is usually always the same (either `Sync` or `Async`).
   *
   * Common use case #2: pinned host memory allocations
   *    - the "memory device" is the host device (e.g. system memory);
   *    - the "synchronisation device" is the accelerator device (e.g. a GPU) whose work queue will access the host;
   *      memory (direct memory access from the accelerator, or scheduling `alpaka::memcpy`/`alpaka::memset`), and can
   *      be different for each allocation;
   *    - the synchronisation `Device` _type_ could potentially be different, but memory pinning is currently tied to
   *      the accelerator's platform (CUDA, HIP, etc.), so the device type needs to be fixed to benefit from caching;
   *    - the `Queue` type can be either `Sync` _or_ `Async` on any allocation.
   */

  template <typename TDev, typename TQueue>
  class CachingAllocator {
  public:
#ifdef ALPAKA_ACC_GPU_CUDA_ENABLED
    friend class alpaka_cuda_async::AlpakaService;
#endif
#ifdef ALPAKA_ACC_GPU_HIP_ENABLED
    friend class alpaka_rocm_async::AlpakaService;
#endif
#ifdef ALPAKA_ACC_CPU_B_SEQ_T_SEQ_ENABLED
    friend class alpaka_serial_sync::AlpakaService;
#endif
#ifdef ALPAKA_ACC_CPU_B_TBB_T_SEQ_ENABLED
    friend class alpaka_tbb_async::AlpakaService;
#endif

    using Device = TDev;                 // the "memory device", where the memory will be allocated
    using Queue = TQueue;                // the queue used to submit the memory operations
    using Event = alpaka::Event<Queue>;  // the events used to synchronise the operations
    using Buffer = alpaka::Buf<Device, std::byte, alpaka::DimInt<1u>, size_t>;

    // The "memory device" type can either be the same as the "synchronisation device" type, or be the host CPU.
    static_assert(alpaka::isDevice<Device>, "TDev should be an alpaka Device type.");
    static_assert(alpaka::isQueue<Queue>, "TQueue should be an alpaka Queue type.");
    static_assert(std::is_same_v<Device, alpaka::Dev<Queue>> or std::is_same_v<Device, alpaka::DevCpu>,
                  "The \"memory device\" type can either be the same as the \"synchronisation device\" type, or be the "
                  "host CPU.");

    struct CachedBytes {
      size_t free = 0;       // total bytes freed and cached on this device
      size_t live = 0;       // total bytes currently in use oin this device
      size_t requested = 0;  // total bytes requested and currently in use on this device
    };

    explicit CachingAllocator(
        Device const& device,
        AllocatorConfig const& config,
        bool reuseSameQueueAllocations,  // Reuse non-ready allocations if they are in the same queue as the new one;
                                         // this is safe only if all memory operations are scheduled in the same queue.
                                         // In particular, this is not safe if the memory will be accessed without using
                                         // any queue, like host memory accessed directly or with immediate operations.
        bool debug = false)
        : device_(device),
          binGrowth_(config.binGrowth),
          minBin_(config.minBin),
          maxBin_(config.maxBin),
          minBinBytes_(detail::power(binGrowth_, minBin_)),
          maxBinBytes_(detail::power(binGrowth_, maxBin_)),
          maxCachedBytes_(cacheSize(config.maxCachedBytes, config.maxCachedFraction)),
          reuseSameQueueAllocations_(reuseSameQueueAllocations),
          debug_(debug),
          fillAllocations_(config.fillAllocations),
          fillAllocationValue_(config.fillAllocationValue),
          fillReallocations_(config.fillReallocations),
          fillReallocationValue_(config.fillReallocationValue),
          fillDeallocations_(config.fillDeallocations),
          fillDeallocationValue_(config.fillDeallocationValue),
          fillCaches_(config.fillCaches),
          fillCacheValue_(config.fillCacheValue) {
      if (debug_) {
        std::ostringstream out;
        out << "CachingAllocator settings\n"
            << "  bin growth " << binGrowth_ << "\n"
            << "  min bin    " << minBin_ << "\n"
            << "  max bin    " << maxBin_ << "\n"
            << "  resulting bins:\n";
        for (auto bin = minBin_; bin <= maxBin_; ++bin) {
          auto binSize = detail::power(binGrowth_, bin);
          out << "    " << std::right << std::setw(12) << detail::as_bytes(binSize) << '\n';
        }
        out << "  maximum amount of cached memory: " << detail::as_bytes(maxCachedBytes_);
        std::cout << out.str() << std::endl;
      }
    }

    ~CachingAllocator() {
      {
        // this should never be called while some memory blocks are still live
        std::scoped_lock lock(mutex_);
        assert(liveBlocks_.empty());
        assert(cachedBytes_.live == 0);
      }

      freeAllCached();
    }

    // return a copy of the cache allocation status, for monitoring purposes
    CachedBytes cacheStatus() const {
      std::scoped_lock lock(mutex_);
      return cachedBytes_;
    }

    // Fill a memory buffer with the specified bye value.
    // If the underlying device is the host and the allocator is configured to support immediate
    // (non queue-ordered) operations, fill the memory synchronously using std::memset.
    // Otherwise, let the alpaka queue schedule the operation.
    //
    // This is not used for deallocation/caching, because the memory may still be in use until the
    // corresponding event is reached.
    void immediateOrAsyncMemset(Queue queue, Buffer buffer, uint8_t value) {
      // host-only
      if (std::is_same_v<Device, alpaka::DevCpu> and not reuseSameQueueAllocations_) {
        std::memset(buffer.data(), value, alpaka::getExtentProduct(buffer) * sizeof(alpaka::Elem<Buffer>));
      } else {
        alpaka::memset(queue, buffer, value);
      }
    }

    // Allocate given number of bytes on the current device associated to given queue
    void* allocate(size_t bytes, Queue queue) {
      // create a block descriptor for the requested allocation
      BlockDescriptor block;
      block.queue = std::move(queue);
      block.requested = bytes;
      std::tie(block.bin, block.bytes) = findBin(bytes);

      // try to re-use a cached block, or allocate a new buffer
      if (tryReuseCachedBlock(block)) {
        // fill the re-used memory block with a pattern
        if (fillReallocations_) {
          immediateOrAsyncMemset(*block.queue, *block.buffer, fillReallocationValue_);
        } else if (fillAllocations_) {
          immediateOrAsyncMemset(*block.queue, *block.buffer, fillAllocationValue_);
        }
      } else {
        allocateNewBlock(block);
        // fill the newly allocated memory block with a pattern
        if (fillAllocations_) {
          immediateOrAsyncMemset(*block.queue, *block.buffer, fillAllocationValue_);
        }
      }

      return block.buffer->data();
    }

    // frees an allocation
    void free(void* ptr) {
      std::scoped_lock lock(mutex_);

      auto iBlock = liveBlocks_.find(ptr);
      if (iBlock == liveBlocks_.end()) {
        std::stringstream ss;
        ss << "Trying to free a non-live block at " << ptr;
        throw std::runtime_error(ss.str());
      }
      // remove the block from the list of live blocks
      BlockDescriptor block = std::move(iBlock->second);
      liveBlocks_.erase(iBlock);
      cachedBytes_.live -= block.bytes;
      cachedBytes_.requested -= block.requested;

      bool recache = (cachedBytes_.free + block.bytes <= maxCachedBytes_);
      if (recache) {
        // If enqueuing the event fails, very likely an error has
        // occurred in the asynchronous processing. In that case the
        // error will show up in all device API function calls, and
        // the free() will be called by destructors during stack
        // unwinding. In order to avoid terminate() being called
        // because of multiple exceptions it is best to ignore these
        // errors.
        try {
          // fill memory blocks with a pattern before caching them
          if (fillCaches_) {
            alpaka::memset(*block.queue, *block.buffer, fillCacheValue_);
          } else if (fillDeallocations_) {
            alpaka::memset(*block.queue, *block.buffer, fillDeallocationValue_);
          }
          // record in the block a marker associated to the work queue
          alpaka::enqueue(*(block.queue), *(block.event));
        } catch (std::exception& e) {
          if (debug_) {
            std::ostringstream out;
            out << "CachingAllocator::free() caught an alpaka error: " << e.what() << "\n";
            out << "\t" << deviceType_ << " " << alpaka::getName(device_) << " freed " << block.bytes << " bytes at "
                << ptr << " from associated queue " << block.queue->m_spQueueImpl.get() << ", event "
                << block.event->m_spEventImpl.get() << " .\n\t\t " << cachedBlocks_.size()
                << " available blocks cached (" << cachedBytes_.free << " bytes), " << liveBlocks_.size()
                << " live blocks (" << cachedBytes_.live << " bytes) outstanding." << std::endl;
            std::cout << out.str() << std::endl;
          }
          return;
        }
        cachedBytes_.free += block.bytes;
        // after the call to insert(), cachedBlocks_ shares ownership of the buffer
        // TODO use std::move ?
        cachedBlocks_.insert(std::make_pair(block.bin, block));

        if (debug_) {
          std::ostringstream out;
          out << "\t" << deviceType_ << " " << alpaka::getName(device_) << " returned " << block.bytes << " bytes at "
              << ptr << " from associated queue " << block.queue->m_spQueueImpl.get() << " , event "
              << block.event->m_spEventImpl.get() << " .\n\t\t " << cachedBlocks_.size() << " available blocks cached ("
              << cachedBytes_.free << " bytes), " << liveBlocks_.size() << " live blocks (" << cachedBytes_.live
              << " bytes) outstanding." << std::endl;
          std::cout << out.str() << std::endl;
        }
      } else {
        // If the memset fails, very likely an error has occurred in the
        // asynchronous processing. In that case the error will show up in all
        // device API function calls, and the free() will be called by
        // destructors during stack unwinding. In order to avoid terminate()
        // being called because of multiple exceptions it is best to ignore
        // these errors.
        try {
          // fill memory blocks with a pattern before freeing them
          if (fillDeallocations_) {
            alpaka::memset(*block.queue, *block.buffer, fillDeallocationValue_);
          }
        } catch (std::exception& e) {
          if (debug_) {
            std::ostringstream out;
            out << "CachingAllocator::free() caught an alpaka error: " << e.what() << "\n";
            out << "\t" << deviceType_ << " " << alpaka::getName(device_) << " freed " << block.bytes << " bytes at "
                << ptr << " from associated queue " << block.queue->m_spQueueImpl.get() << ", event "
                << block.event->m_spEventImpl.get() << " .\n\t\t " << cachedBlocks_.size()
                << " available blocks cached (" << cachedBytes_.free << " bytes), " << liveBlocks_.size()
                << " live blocks (" << cachedBytes_.live << " bytes) outstanding." << std::endl;
            std::cout << out.str() << std::endl;
          }
          return;
        }
        // if the buffer is not recached, it is automatically freed when block goes out of scope
        if (debug_) {
          std::ostringstream out;
          out << "\t" << deviceType_ << " " << alpaka::getName(device_) << " freed " << block.bytes << " bytes at "
              << ptr << " from associated queue " << block.queue->m_spQueueImpl.get() << ", event "
              << block.event->m_spEventImpl.get() << " .\n\t\t " << cachedBlocks_.size() << " available blocks cached ("
              << cachedBytes_.free << " bytes), " << liveBlocks_.size() << " live blocks (" << cachedBytes_.live
              << " bytes) outstanding." << std::endl;
          std::cout << out.str() << std::endl;
        }
      }
    }

  private:
    struct BlockDescriptor {
      std::optional<Buffer> buffer;
      std::optional<Queue> queue;
      std::optional<Event> event;
      size_t bytes = 0;
      size_t requested = 0;  // for monitoring only
      unsigned int bin = 0;

      // the "synchronisation device" for this block
      auto device() { return alpaka::getDev(*queue); }
    };

  private:
    // return the maximum amount of memory that should be cached on this device
    size_t cacheSize(size_t maxCachedBytes, double maxCachedFraction) const {
      // note that getMemBytes() returns 0 if the platform does not support querying the device memory
      size_t totalMemory = alpaka::getMemBytes(device_);
      size_t memoryFraction = static_cast<size_t>(maxCachedFraction * totalMemory);
      size_t size = std::numeric_limits<size_t>::max();
      if (maxCachedBytes > 0 and maxCachedBytes < size) {
        size = maxCachedBytes;
      }
      if (memoryFraction > 0 and memoryFraction < size) {
        size = memoryFraction;
      }
      return size;
    }

    // return (bin, bin size)
    std::tuple<unsigned int, size_t> findBin(size_t bytes) const {
      if (bytes < minBinBytes_) {
        return std::make_tuple(minBin_, minBinBytes_);
      }
      if (bytes > maxBinBytes_) {
        throw std::runtime_error("Requested allocation size " + std::to_string(bytes) +
                                 " bytes is too large for the caching detail with maximum bin " +
                                 std::to_string(maxBinBytes_) +
                                 " bytes. You might want to increase the maximum bin size");
      }
      unsigned int bin = minBin_;
      size_t binBytes = minBinBytes_;
      while (binBytes < bytes) {
        ++bin;
        binBytes *= binGrowth_;
      }
      return std::make_tuple(bin, binBytes);
    }

    bool tryReuseCachedBlock(BlockDescriptor& block) {
      std::scoped_lock lock(mutex_);

      // iterate through the range of cached blocks in the same bin
      const auto [begin, end] = cachedBlocks_.equal_range(block.bin);
      for (auto iBlock = begin; iBlock != end; ++iBlock) {
        if ((reuseSameQueueAllocations_ and (*block.queue == *(iBlock->second.queue))) or
            alpaka::isComplete(*(iBlock->second.event))) {
          // associate the cached buffer to the new queue
          auto queue = std::move(*(block.queue));
          // TODO cache (or remove) the debug information and use std::move()
          block = iBlock->second;
          block.queue = std::move(queue);

          // if the new queue is on different device than the old event, create a new event
          if (block.device() != alpaka::getDev(*(block.event))) {
            block.event = Event{block.device()};
          }

          // insert the cached block into the live blocks
          // TODO cache (or remove) the debug information and use std::move()
          liveBlocks_[block.buffer->data()] = block;

          // update the accounting information
          cachedBytes_.free -= block.bytes;
          cachedBytes_.live += block.bytes;
          cachedBytes_.requested += block.requested;

          if (debug_) {
            std::ostringstream out;
            out << "\t" << deviceType_ << " " << alpaka::getName(device_) << " reused cached block at "
                << block.buffer->data() << " (" << block.bytes << " bytes) for queue "
                << block.queue->m_spQueueImpl.get() << ", event " << block.event->m_spEventImpl.get()
                << " (previously associated with queue " << iBlock->second.queue->m_spQueueImpl.get() << " , event "
                << iBlock->second.event->m_spEventImpl.get() << ")." << std::endl;
            std::cout << out.str() << std::endl;
          }

          // remove the reused block from the list of cached blocks
          cachedBlocks_.erase(iBlock);
          return true;
        }
      }

      return false;
    }

    Buffer allocateBuffer(size_t bytes, Queue const& queue) {
      if constexpr (std::is_same_v<Device, alpaka::Dev<Queue>>) {
        // allocate device memory
        return alpaka::allocBuf<std::byte, size_t>(device_, bytes);
      } else if constexpr (std::is_same_v<Device, alpaka::DevCpu>) {
        // allocate pinned host memory accessible by the queue's platform
        using Platform = alpaka::Platform<alpaka::Dev<Queue>>;
        return alpaka::allocMappedBuf<std::byte, size_t>(device_, platform<Platform>(), bytes);
      } else {
        // unsupported combination
        static_assert(std::is_same_v<Device, alpaka::Dev<Queue>> or std::is_same_v<Device, alpaka::DevCpu>,
                      "The \"memory device\" type can either be the same as the \"synchronisation device\" type, or be "
                      "the host CPU.");
      }
    }

    void allocateNewBlock(BlockDescriptor& block) {
      try {
        block.buffer = allocateBuffer(block.bytes, *block.queue);
      } catch (std::runtime_error const& e) {
        // the allocation attempt failed: free all cached blocks on the device and retry
        if (debug_) {
          std::ostringstream out;
          out << "\t" << deviceType_ << " " << alpaka::getName(device_) << " failed to allocate " << block.bytes
              << " bytes for queue " << block.queue->m_spQueueImpl.get()
              << ", retrying after freeing cached allocations" << std::endl;
          std::cout << out.str() << std::endl;
        }
        // TODO implement a method that frees only up to block.bytes bytes
        freeAllCached();

        // throw an exception if it fails again
        block.buffer = allocateBuffer(block.bytes, *block.queue);
      }

      // create a new event associated to the "synchronisation device"
      block.event = Event{block.device()};

      {
        std::scoped_lock lock(mutex_);
        cachedBytes_.live += block.bytes;
        cachedBytes_.requested += block.requested;
        // TODO use std::move() ?
        liveBlocks_[block.buffer->data()] = block;
      }

      if (debug_) {
        std::ostringstream out;
        out << "\t" << deviceType_ << " " << alpaka::getName(device_) << " allocated new block at "
            << block.buffer->data() << " (" << block.bytes << " bytes associated with queue "
            << block.queue->m_spQueueImpl.get() << ", event " << block.event->m_spEventImpl.get() << "." << std::endl;
        std::cout << out.str() << std::endl;
      }
    }

    void freeAllCached() {
      std::scoped_lock lock(mutex_);

      while (not cachedBlocks_.empty()) {
        auto iBlock = cachedBlocks_.begin();
        cachedBytes_.free -= iBlock->second.bytes;

        if (debug_) {
          std::ostringstream out;
          out << "\t" << deviceType_ << " " << alpaka::getName(device_) << " freed " << iBlock->second.bytes
              << " bytes.\n\t\t  " << (cachedBlocks_.size() - 1) << " available blocks cached (" << cachedBytes_.free
              << " bytes), " << liveBlocks_.size() << " live blocks (" << cachedBytes_.live << " bytes) outstanding."
              << std::endl;
          std::cout << out.str() << std::endl;
        }

        cachedBlocks_.erase(iBlock);
      }
    }

    // TODO replace with a tbb::concurrent_multimap ?
    using CachedBlocks = std::multimap<unsigned int, BlockDescriptor>;  // ordered by the allocation bin
    // TODO replace with a tbb::concurrent_map ?
    using BusyBlocks = std::map<void*, BlockDescriptor>;  // ordered by the address of the allocated memory

    inline static const std::string deviceType_ = alpaka::core::demangled<Device>;

    mutable std::mutex mutex_;
    Device device_;  // the device where the memory is allocated

    CachedBytes cachedBytes_;
    CachedBlocks cachedBlocks_;  // Set of cached device allocations available for reuse
    BusyBlocks liveBlocks_;      // map of pointers to the live device allocations currently in use

    const unsigned int binGrowth_;  // Geometric growth factor for bin-sizes
    const unsigned int minBin_;
    const unsigned int maxBin_;

    const size_t minBinBytes_;
    const size_t maxBinBytes_;
    const size_t maxCachedBytes_;  // Maximum aggregate cached bytes per device

    const bool reuseSameQueueAllocations_;
    const bool debug_;

    const bool fillAllocations_;
    const uint8_t fillAllocationValue_;
    const bool fillReallocations_;
    const uint8_t fillReallocationValue_;
    const bool fillDeallocations_;
    const uint8_t fillDeallocationValue_;
    const bool fillCaches_;
    const uint8_t fillCacheValue_;
  };

}  // namespace cms::alpakatools

#endif  // HeterogeneousCore_AlpakaInterface_interface_CachingAllocator_h