// Object recognition inference entirely in KNP on AltAI neurons: pixels → feature detectors → class neurons.
//
// altai_infer <images.bin> <first> <count> <model_dir> <backend.so> <out_prefix>
//
// images.bin: uint8 images 3×32×32. Pixels spike with the incrementing converter (value/255 per step, 10 active
// steps of 15). model_dir (written by build_model.py / train_readout.py):
//   meta.txt        K stride det_threshold classes out_threshold
//   det_w.bin       int32 [K][108]  detector weights on a 6×6×3 pixel patch (channel, y, x)
//   det_b.bin       int32 [K]       detector bias, delivered on every active step by an always-on input
//   readout_w.bin   int32 [classes][4·K] (optional) class weights per (quadrant, detector type); index q·K + k
//   readout_b.bin   int32 [classes] (optional) class bias per active step
// Detectors: one AltAI neuron per detector type and position (stride), threshold det_threshold, reset to 0, the
// potential does not go below 0. Class neurons: AltAI, threshold out_threshold.
// Output: <out_prefix>_classes.csv (index, spike count per class) if a readout is given, and
// <out_prefix>_counts.bin (int32 [count][4·K], detector spikes per quadrant and type) for training the readout.

#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 <cstring>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <string>
#include <stdexcept>
#include <vector>


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;
constexpr int fan_in = patch * patch * channels, in_n = channels * side * side;


template <class T>
std::vector<T> read_vec(const std::filesystem::path &path)
{
    std::ifstream f(path, std::ios::binary);
    if (!f) return {};
    std::vector<char> raw((std::istreambuf_iterator<char>(f)), std::istreambuf_iterator<char>());
    if (raw.size() % sizeof(T)) throw std::runtime_error("Invalid weight file byte length");
    std::vector<T> out(raw.size() / sizeof(T));
    std::memcpy(out.data(), raw.data(), out.size() * sizeof(T));
    return out;
}


knp::neuron_traits::neuron_parameters<knp::neuron_traits::AltAILIF> neuron(int threshold)
{
    knp::neuron_traits::neuron_parameters<knp::neuron_traits::AltAILIF> p;
    p.activation_threshold_ = static_cast<uint16_t>(threshold);
    p.negative_activation_threshold_ = 0;
    p.potential_reset_value_ = 0;
    p.potential_leak_ = 0;
    return p;
}


int main(int argc, const char *const argv[]) try
{
    if (argc != 7)
    {
        std::cerr << "usage: altai_infer <images.bin> <first> <count> <model_dir> <backend.so> <out_prefix>\n";
        return 1;
    }
    const std::filesystem::path model_dir = argv[4];
    const size_t first = std::stoul(argv[2]), count = std::stoul(argv[3]);
    const std::string out_prefix = argv[6];
    int k = 0, stride = 0, det_thr = 0, classes = 0, out_thr = 0;
    {
        std::ifstream meta(model_dir / "meta.txt");
        if (!(meta >> k >> stride >> det_thr >> classes >> out_thr)) throw std::runtime_error("Invalid model metadata");
    }
    if (k <= 0 || k > 10000 || stride <= 0 || stride > side - patch + 1 ||
        det_thr < 1 || det_thr > 65535 || classes < 1 || classes > 1000 || count < 1 || count > 100000)
        throw std::runtime_error("Invalid model dimensions, threshold or image count");
    const int per_side = (side - patch) / stride + 1, positions = per_side * per_side, half = per_side / 2;
    const size_t det_n = static_cast<size_t>(k) * positions;
    const auto det_w = read_vec<int32_t>(model_dir / "det_w.bin"), det_b = read_vec<int32_t>(model_dir / "det_b.bin");
    const auto ro_w = read_vec<int32_t>(model_dir / "readout_w.bin"), ro_b = read_vec<int32_t>(model_dir / "readout_b.bin");
    if (det_w.size() != static_cast<size_t>(k * fan_in) || det_b.size() != static_cast<size_t>(k))
        throw std::runtime_error("Missing or invalid detector weights/bias");
    const bool readout = !ro_w.empty();
    if (readout && (ro_w.size() != static_cast<size_t>(classes * 4 * k) ||
                    ro_b.size() != static_cast<size_t>(classes) || out_thr < 1 || out_thr > 65535))
        throw std::runtime_error("Invalid readout weights, bias or threshold");
    const auto file_size = std::filesystem::file_size(argv[1]);
    if (file_size % in_n || first > file_size / in_n || count > file_size / in_n - first)
        throw std::runtime_error("Invalid or out-of-range image file request");
    std::ifstream images(argv[1], std::ios::binary);
    std::vector<uint8_t> px(count * in_n);
    images.seekg(static_cast<std::streamoff>(first * in_n));
    images.read(reinterpret_cast<char *>(px.data()), static_cast<std::streamsize>(px.size()));
    if (!images || static_cast<size_t>(images.gcount()) != px.size())
        throw std::runtime_error("Input file does not contain the requested image range");
    if (!std::filesystem::is_directory(std::filesystem::path(out_prefix).parent_path().empty()
                                      ? std::filesystem::path(".") : std::filesystem::path(out_prefix).parent_path()))
        throw std::runtime_error("Output directory does not exist");

    auto quadrant = [&](size_t pos) { return (pos / per_side >= half) * 2 + (pos % per_side >= half); };

    AltAIPopulation detectors([&](size_t) { return neuron(det_thr); }, det_n);
    AltAIPopulation classes_pop([&](size_t) { return neuron(out_thr); }, readout ? classes : 0);
    // Detector index = type * positions + position; synapse index = detector * fan_in + patch element.
    DeltaProjection pixel_proj(
        knp::core::UID{false}, detectors.get_uid(),
        [&](size_t index) -> std::optional<DeltaProjection::Synapse>
        {
            const size_t d = index / fan_in, e = index % fan_in, j = d / positions, pos = d % positions;
            const int32_t w = det_w[j * fan_in + e];
            if (w == 0) return std::nullopt;
            const size_t c = e / (patch * patch), dy = (e % (patch * patch)) / patch, dx = e % patch;
            const size_t y = (pos / per_side) * stride + dy, x = (pos % per_side) * stride + dx;
            return DeltaProjection::Synapse{
                {static_cast<float>(w), 1, knp::synapse_traits::OutputType::EXCITATORY}, c * side * side + y * side + x,
                d};
        },
        det_n * fan_in);
    // Always-on input 0 → every detector (its bias) and every class neuron (its bias).
    DeltaProjection bias_det(
        knp::core::UID{false}, detectors.get_uid(),
        [&](size_t d) -> std::optional<DeltaProjection::Synapse>
        {
            const int32_t b = det_b[d / positions];
            if (b == 0) return std::nullopt;
            return DeltaProjection::Synapse{{static_cast<float>(b), 1, knp::synapse_traits::OutputType::EXCITATORY}, 0, d};
        },
        det_n);
    DeltaProjection readout_proj(
        detectors.get_uid(), classes_pop.get_uid(),
        [&](size_t index) -> std::optional<DeltaProjection::Synapse>
        {
            const size_t c = index / det_n, d = index % det_n, j = d / positions, pos = d % positions;
            const int32_t w = ro_w[c * 4 * k + quadrant(pos) * k + j];
            if (w == 0) return std::nullopt;
            return DeltaProjection::Synapse{{static_cast<float>(w), 1, knp::synapse_traits::OutputType::EXCITATORY}, d, c};
        },
        readout ? det_n * classes : 0);
    // The class bias arrives one step later than the detector bias, in step with the detector spikes.
    DeltaProjection bias_cls(
        knp::core::UID{false}, classes_pop.get_uid(),
        [&](size_t c) -> std::optional<DeltaProjection::Synapse>
        {
            if (ro_b[c] == 0) return std::nullopt;
            return DeltaProjection::Synapse{
                {static_cast<float>(ro_b[c]), 2, knp::synapse_traits::OutputType::EXCITATORY}, 0, c};
        },
        readout ? classes : 0);

    const auto pixel_uid = pixel_proj.get_uid(), bias_det_uid = bias_det.get_uid(), bias_cls_uid = bias_cls.get_uid(),
               det_uid = detectors.get_uid(), cls_uid = classes_pop.get_uid();
    knp::framework::Network network;
    network.add_population(std::move(detectors));
    network.add_projection(std::move(pixel_proj));
    network.add_projection(std::move(bias_det));
    if (readout)
    {
        network.add_population(std::move(classes_pop));
        network.add_projection(std::move(readout_proj));
        network.add_projection(std::move(bias_cls));
    }
    knp::framework::Model model(std::move(network));
    knp::core::UID pixel_channel, bias_channel, det_channel, cls_channel;
    model.add_input_channel(pixel_channel, pixel_uid);
    model.add_input_channel(bias_channel, bias_det_uid);
    if (readout) model.add_input_channel(bias_channel, bias_cls_uid);
    model.add_output_channel(det_channel, det_uid);
    if (readout) model.add_output_channel(cls_channel, cls_uid);

    std::vector<knp::core::messaging::SpikeData> current(active_steps);
    auto pixel_gen = [&](knp::core::Step step) -> knp::core::messaging::SpikeData
    {
        const size_t img = step / steps_per_image, local = step % steps_per_image;
        if (img >= count) return {};
        if (local == 0)
        {
            std::vector<float> st(in_n, 0.F);
            for (auto &s : current) s.clear();
            const uint8_t *p = px.data() + img * in_n;
            for (int t = 0; t < active_steps; ++t)
                for (int i = 0; i < in_n; ++i)
                {
                    st[i] += p[i] / 255.F;
                    if (st[i] >= 1.F)
                    {
                        st[i] -= 1.F;
                        current[t].push_back(static_cast<uint32_t>(i));
                    }
                }
        }
        if (local >= active_steps) return {};
        return current[local];
    };
    auto bias_gen = [&](knp::core::Step step) -> knp::core::messaging::SpikeData
    {
        if (step >= count * steps_per_image || step % steps_per_image >= active_steps) return {};
        return {0};
    };

    knp::framework::BackendLoader loader;
    knp::framework::ModelExecutor executor(
        model, loader.load(argv[5]), {{pixel_channel, pixel_gen}, {bias_channel, bias_gen}});
    auto &det_out = executor.get_loader().get_output_channel(det_channel);
    std::vector<int32_t> counts(count * 4 * k, 0);
    std::vector<int32_t> class_counts(count * std::max(classes, 1), 0);
    const auto t0 = std::chrono::steady_clock::now();
    auto collect_step = [&](size_t step)
    {
        for (const auto &m : det_out.read_some_from_buffer(step, step))
        {
            const size_t im = m.header_.send_time_ / steps_per_image;
            if (im >= count) continue;
            for (auto d : m.neuron_indexes_) ++counts[(im * 4 + quadrant(d % positions)) * k + d / positions];
        }
        if (readout)
            for (const auto &m : executor.get_loader().get_output_channel(cls_channel).read_some_from_buffer(step, step))
            {
                const size_t im = m.header_.send_time_ / steps_per_image;
                if (im < count)
                    for (auto c : m.neuron_indexes_) ++class_counts[im * classes + c];
            }
    };
    executor.start([&](size_t step)
    {
        if (step > 0) collect_step(step - 1);
        return step < count * steps_per_image;
    });
    std::ofstream binary;
    binary.exceptions(std::ios::failbit | std::ios::badbit);
    binary.open(out_prefix + "_counts.bin", std::ios::binary);
    binary.write(reinterpret_cast<const char *>(counts.data()), static_cast<std::streamsize>(counts.size() * 4));
    binary.close();
    if (readout)
    {
        std::ofstream csv;
        csv.exceptions(std::ios::failbit | std::ios::badbit);
        csv.open(out_prefix + "_classes.csv");
        csv << "index";
        for (int c = 0; c < classes; ++c) csv << ",c" << c;
        csv << "\n";
        for (size_t i = 0; i < count; ++i)
        {
            csv << first + i;
            for (int c = 0; c < classes; ++c) csv << ',' << class_counts[i * classes + c];
            csv << "\n";
        }
        csv.close();
    }
    std::cout << "images " << first << ".." << first + count - 1 << ", "
              << std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count() << " s" << std::endl;
}

catch (const std::exception &error)
{
    std::cerr << "altai_infer: " << error.what() << "\n";
    return 1;
}