#include "gpu_stdp_bridge.h"
#include "gpu_stdp_types.h"
#include <knp/backends/cpu-library/impl/populations/training/stdp.h>
#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <iostream>
#include <map>
#include <memory>
#include <stdexcept>

namespace knp_gpu_experiment
{
using Params = knp::synapse_traits::synapse_parameters<Synapse>;
using NeuronParams = knp::neuron_traits::neuron_parameters<Neuron>;
struct Context
{
    Population* population;
    std::vector<std::reference_wrapper<Params>> synapses;
    std::vector<std::size_t> offsets;
    std::vector<GpuStdpSynapse> state;
    std::vector<GpuStdpNeuron> neurons;
    std::vector<std::uint64_t> last;
    std::vector<float> weights;
    GpuStdpContext* gpu = nullptr;
    ~Context() { gpu_stdp_destroy(gpu); }
};
static thread_local std::map<const Population*, std::unique_ptr<Context>> contexts;
static thread_local std::size_t checked_steps = 0, gpu_steps = 0;
static thread_local double max_error = 0;

static GpuStdpSynapse pack(const Params& p)
{
    const auto& r = p.rule_;
    return {r.synaptic_resource_, r.w_min_, r.w_max_, r.d_u_, p.weight_,
            r.dopamine_plasticity_period_, p.delay_, r.has_contributed_, r.had_hebbian_update_};
}
static GpuStdpNeuron pack(const NeuronParams& n, bool spike)
{
    return {n.additional_threshold_, n.dopamine_value_, n.synapse_sum_threshold_coefficient_,
            n.free_synaptic_resource_, n.synaptic_resource_threshold_, n.stability_,
            n.stability_change_parameter_, n.stability_change_at_isi_, n.d_h_,
            n.last_step_, n.last_spike_step_, n.first_isi_spike_, n.isi_max_,
            n.dopamine_plasticity_time_, n.resource_drain_coefficient_,
            static_cast<int>(n.isi_status_), n.is_being_forced_, spike};
}
static void unpack(const GpuStdpNeuron& s, NeuronParams& n)
{
    n.additional_threshold_ = s.additional;
    n.free_synaptic_resource_ = s.free_resource;
    n.stability_ = s.stability;
    n.last_step_ = s.last_step; n.last_spike_step_ = s.last_spike; n.first_isi_spike_ = s.first_isi;
    n.isi_status_ = static_cast<knp::neuron_traits::ISIPeriodType>(s.status);
}
static void unpack(const GpuStdpSynapse& s, Params& p)
{
    p.weight_ = s.weight;
    p.rule_.synaptic_resource_ = s.resource;
    p.rule_.has_contributed_ = s.contributed;
    p.rule_.had_hebbian_update_ = s.hebbian;
    // last_spike_step_ remains authoritative on CPU (the projection phase updates it).
}
static bool near(double a, double b)
{
    if (a == b) return true;
    const double error = std::abs(a - b);
    max_error = std::max(max_error, error);
    return std::isfinite(a) && std::isfinite(b) && error <= 1e-4 + 1e-6 * std::max(std::abs(a), std::abs(b));
}
static void compare(const GpuStdpNeuron& actual, const GpuStdpNeuron& expected)
{
    if (!near(actual.additional, expected.additional) || !near(actual.free_resource, expected.free_resource) ||
        !near(actual.stability, expected.stability) || actual.last_step != expected.last_step ||
        actual.last_spike != expected.last_spike || actual.first_isi != expected.first_isi ||
        actual.status != expected.status)
        throw std::runtime_error("CUDA STDP neuron-state parity failed");
}
static void synchronize(Context& c)
{
    gpu_stdp_download(c.gpu, c.state.data());
    for (std::size_t j = 0; j < c.synapses.size(); ++j) unpack(c.state[j], c.synapses[j]);
}
void flush()
{
    for (auto& entry : contexts) synchronize(*entry.second);
    if (!contexts.empty())
        std::cerr << "CUDA_STDP gpu_steps=" << gpu_steps << " checked_steps=" << checked_steps
                  << " max_state_abs_error=" << max_error << '\n';
}
void clear()
{
    contexts.clear(); checked_steps = gpu_steps = 0; max_error = 0;
}
void train(Population& population, Projections& projections,
           const knp::core::messaging::SpikeMessage& message, knp::core::Step step, Reference reference)
{
    if (projections.empty()) { reference(population, projections, message, step); return; }
    bool active = !message.neuron_indexes_.empty();
    for (const auto& n : population)
        active |= n.dopamine_value_ > 0 || (n.dopamine_value_ < 0 && n.isi_status_ != knp::neuron_traits::ISIPeriodType::is_forced)
            || (!(step - n.last_step_ <= n.isi_max_ && n.isi_status_ != knp::neuron_traits::ISIPeriodType::is_forced)
                && std::fabs(n.free_synaptic_resource_) >= n.synaptic_resource_threshold_);
    if (!active) return;
    auto& owned = contexts[&population];
    if (!owned)
    {
        owned = std::make_unique<Context>();
        auto& c = *owned;
        c.population = &population;
        c.offsets.push_back(0);
        for (std::size_t i = 0; i < population.size(); ++i)
        {
            const auto& synapses = knp::backends::cpu::populations::impl::training::stdp::
                get_all_connected_synapses<Synapse>(projections, i);
            c.synapses.insert(c.synapses.end(), synapses.begin(), synapses.end());
            c.offsets.push_back(c.synapses.size());
        }
        if (c.synapses.empty()) { owned.reset(); reference(population, projections, message, step); return; }
        c.neurons.resize(population.size()); c.state.reserve(c.synapses.size());
        for (auto& s : c.synapses) c.state.push_back(pack(s));
        c.last.resize(c.synapses.size()); c.weights.resize(c.synapses.size());
        for (std::size_t j = 0; j < c.synapses.size(); ++j) c.weights[j] = c.synapses[j].get().weight_;
        c.gpu = gpu_stdp_create(c.state.data(), c.weights.data(), c.offsets.data(), population.size(), c.synapses.size());
    }
    auto& c = *owned;
    if (c.neurons.size() != population.size()) throw std::runtime_error("CUDA STDP requires fixed topology");
    const bool verify = std::getenv("KNP_CUDA_VERIFY") != nullptr;
    if (verify) synchronize(c);
    std::vector<bool> spiked(population.size());
    for (auto i : message.neuron_indexes_) spiked.at(i) = true;
    for (std::size_t i = 0; i < population.size(); ++i) c.neurons[i] = pack(population[i], spiked[i]);
    for (std::size_t j = 0; j < c.synapses.size(); ++j) c.last[j] = c.synapses[j].get().rule_.last_spike_step_;
    std::vector<GpuStdpNeuron> expected_neurons;
    std::vector<GpuStdpSynapse> expected_synapses;
    if (verify)
    {
        reference(population, projections, message, step);
        for (std::size_t i = 0; i < population.size(); ++i) expected_neurons.push_back(pack(population[i], spiked[i]));
        for (auto& s : c.synapses) expected_synapses.push_back(pack(s));
    }
    gpu_stdp_step(c.gpu, c.neurons.data(), c.last.data(), c.weights.data(), step);
    ++gpu_steps;
    for (std::size_t i = 0; i < population.size(); ++i)
    {
        if (verify) compare(c.neurons[i], expected_neurons[i]);
        unpack(c.neurons[i], population[i]);
    }
    for (std::size_t j = 0; j < c.synapses.size(); ++j) c.synapses[j].get().weight_ = c.weights[j];
    if (verify)
    {
        gpu_stdp_download(c.gpu, c.state.data());
        for (std::size_t j = 0; j < c.synapses.size(); ++j)
        {
            const auto& s = c.state[j]; const auto& e = expected_synapses[j];
            if (!near(s.resource, e.resource) || !near(s.weight, e.weight) ||
                s.contributed != e.contributed || s.hebbian != e.hebbian)
                throw std::runtime_error("CUDA STDP synapse-state parity failed at " + std::to_string(step));
            unpack(s, c.synapses[j]);
        }
        ++checked_steps;
    }
}
}