// Speed check: a convolution-like feature layer built from AltAI neurons in KNP.
// usage: bench_conv <K> <stride> <images> <backend.so>
// K prototypes × positions (6×6×3 patches of a 32×32×3 image with the given stride) AltAI neurons; random integer
// weights; random input spikes (pixel rate ~0.4 per active step); 15 steps per image, 10 of them active.

#include <knp/framework/model_executor.h>
#include <knp/framework/network.h>
#include <knp/neuron-traits/altai_lif.h>
#include <knp/synapse-traits/delta.h>

#include <chrono>
#include <iostream>
#include <random>
#include <string>


using DeltaProjection = knp::core::Projection<knp::synapse_traits::DeltaSynapse>;
using AltAIPopulation = knp::core::Population<knp::neuron_traits::AltAILIF>;

constexpr int side = 32, patch = 6, channels = 3, steps_per_image = 15, active_steps = 10;


int main(int argc, const char *const argv[])
{
    const int k = std::stoi(argv[1]), stride = std::stoi(argv[2]), images = std::stoi(argv[3]);
    const int per_side = (side - patch) / stride + 1, positions = per_side * per_side;
    const size_t neurons = static_cast<size_t>(k) * positions, fan_in = patch * patch * channels;

    AltAIPopulation population(
        [](size_t)
        {
            knp::neuron_traits::neuron_parameters<knp::neuron_traits::AltAILIF> p;
            p.activation_threshold_ = 1000;
            p.negative_activation_threshold_ = 0;
            p.potential_reset_value_ = 0;
            return p;
        },
        neurons);

    std::mt19937 rng(1);
    std::uniform_int_distribution<int> weight(-60, 60);
    std::vector<int> w(fan_in * k);
    for (auto &v : w) v = weight(rng);

    // Synapse index → (neuron, patch element); the source is the input channel index of that pixel.
    DeltaProjection input_projection(
        knp::core::UID{false}, population.get_uid(),
        [&](size_t index) -> std::optional<DeltaProjection::Synapse>
        {
            const size_t neuron = index / fan_in, e = index % fan_in;
            const int proto = static_cast<int>(neuron / positions), pos = static_cast<int>(neuron % positions);
            const int py = (pos / per_side) * stride, px = (pos % per_side) * stride;
            const int c = static_cast<int>(e / (patch * patch)), dy = static_cast<int>(e % (patch * patch)) / patch,
                      dx = static_cast<int>(e % patch);
            const size_t source = static_cast<size_t>(c * side * side + (py + dy) * side + px + dx);
            return DeltaProjection::Synapse{
                {static_cast<float>(w[proto * fan_in + e]), 1, knp::synapse_traits::OutputType::EXCITATORY}, source,
                neuron};
        },
        neurons * fan_in);

    const knp::core::UID input_uid = input_projection.get_uid(), output_uid = population.get_uid();
    knp::framework::Network network;
    network.add_population(std::move(population));
    network.add_projection(std::move(input_projection));
    knp::core::UID in_channel, out_channel;
    knp::framework::Model model(std::move(network));
    model.add_input_channel(in_channel, input_uid);
    model.add_output_channel(out_channel, output_uid);

    std::bernoulli_distribution fire(0.4);
    auto input_gen = [&](knp::core::Step step) -> knp::core::messaging::SpikeData
    {
        knp::core::messaging::SpikeData s;
        if (step % steps_per_image >= active_steps) return s;
        for (size_t i = 0; i < side * side * channels; ++i)
            if (fire(rng)) s.push_back(i);
        return s;
    };

    knp::framework::BackendLoader loader;
    knp::framework::ModelExecutor executor(model, loader.load(argv[4]), {{in_channel, input_gen}});
    auto &out = executor.get_loader().get_output_channel(out_channel);
    const auto t0 = std::chrono::steady_clock::now();
    executor.start([&](size_t step) { return step < static_cast<size_t>(images * steps_per_image); });
    const double sec = std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count();
    size_t spikes = 0;
    for (const auto &m : out.update()) spikes += m.neuron_indexes_.size();
    std::cout << "K=" << k << " stride=" << stride << " neurons=" << neurons << " synapses=" << neurons * fan_in
              << " images=" << images << " seconds=" << sec << " per_image=" << sec / images
              << " output_spikes_per_image=" << spikes / images << std::endl;
}