// Experiment-only cache for the fixed-topology CPU training backend.
#pragma once
#include <knp/core/projection.h>
#include <cstdint>
#include <functional>
#include <map>
#include <utility>
#include <vector>

namespace knp::backends::cpu::populations::impl::training::stdp
{
template <class Synapse>
struct ConnectionCacheEntry
{
    using Projection = knp::core::Projection<Synapse>;
    using Params = knp::synapse_traits::synapse_parameters<Synapse>;
    std::vector<Projection*> projections;
    std::vector<size_t> sizes;
    std::vector<const void*> storage;
    std::vector<std::reference_wrapper<Params>> values;
};

template <class Synapse>
auto& connection_cache()
{
    using Key = std::pair<std::uintptr_t, size_t>;
    static thread_local std::map<Key, ConnectionCacheEntry<Synapse>> cache;
    return cache;
}

template <class Synapse>
void clear_connection_cache()
{
    connection_cache<Synapse>().clear();
}

template <class Synapse>
const std::vector<std::reference_wrapper<knp::synapse_traits::synapse_parameters<Synapse>>>&
get_all_connected_synapses(
    std::vector<std::reference_wrapper<knp::core::Projection<Synapse>>>& projections, size_t neuron)
{
    using Projection = knp::core::Projection<Synapse>;
    using Values = std::vector<std::reference_wrapper<knp::synapse_traits::synapse_parameters<Synapse>>>;
    static const Values empty;
    if (projections.empty()) return empty;
    auto key = std::make_pair(reinterpret_cast<std::uintptr_t>(&projections.front().get()), neuron);
    auto& entry = connection_cache<Synapse>()[key];
    bool valid = entry.projections.size() == projections.size();
    for (size_t i = 0; valid && i < projections.size(); ++i)
    {
        auto& p = projections[i].get();
        const void* storage = p.size() ? static_cast<const void*>(&p[0]) : nullptr;
        valid = entry.projections[i] == &p && entry.sizes[i] == p.size() && entry.storage[i] == storage;
    }
    if (!valid)
    {
        entry = ConnectionCacheEntry<Synapse>{};
        for (auto& ref : projections)
        {
            auto& p = ref.get();
            entry.projections.push_back(&p);
            entry.sizes.push_back(p.size());
            entry.storage.push_back(p.size() ? static_cast<const void*>(&p[0]) : nullptr);
            // Preserve the original index order, hence the floating-point update order.
            for (auto index : p.find_synapses(neuron, Projection::Search::by_postsynaptic))
                entry.values.emplace_back(std::get<knp::core::synapse_data>(p[index]));
        }
    }
    return entry.values;
}
}