// 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 = 20, 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 != 3)
    {
        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[1];
    const size_t first = 0, count = 1;
    
    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) throw std::runtime_error("Detector executable requires detector-only model");
    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");
    std::vector<uint8_t> px(in_n);
    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);

    DeltaProjection reset_det(knp::core::UID{false}, detectors.get_uid(),
        [&](size_t d) -> std::optional<DeltaProjection::Synapse> {
            return DeltaProjection::Synapse{{-65535.F, 1, knp::synapse_traits::OutputType::EXCITATORY}, 0, d};
        }, det_n);
    // Each detector's positive residual is below the largest possible single-step
    // input (108*127 + bias); one negative65535 pulse therefore clamps it to0.
    for (auto w : det_w) if (std::abs(w)>127) throw std::runtime_error("Detector weight outside reset proof");
    for (auto b : det_b) if (std::abs(b)>1024) throw std::runtime_error("Detector bias outside reset proof");
    const auto reset_uid=reset_det.get_uid();
    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(reset_det));
    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;
    knp::core::UID reset_channel;
    model.add_input_channel(reset_channel, reset_uid);
    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);
    std::vector<int32_t> counts(4*k),heat(positions);
    knp::framework::ModelExecutor*ep=nullptr;size_t last=size_t(-1);bool done=false;std::string prefix;
    auto write=[&](const std::string&name,const std::vector<int32_t>&v){std::ofstream f;f.exceptions(std::ios::badbit|std::ios::failbit);f.open(name,std::ios::binary);f.write(reinterpret_cast<const char*>(v.data()),v.size()*4);f.close();};
    auto advance=[&](knp::core::Step t){
        if(last==t)return;last=t;
        if(ep&&t){for(auto&msg:ep->get_loader().get_output_channel(det_channel).read_some_from_buffer(t-1,t-1)){
            if(msg.header_.send_time_%steps_per_image<4&&!msg.neuron_indexes_.empty())throw std::runtime_error("Unexpected reset spike");
            for(auto d:msg.neuron_indexes_){++counts[quadrant(d%positions)*k+d/positions];++heat[d%positions];}
        }}
        if(t%steps_per_image==0){
            if(t){write(prefix+"_counts.bin",counts);write(prefix+"_heat.bin",heat);std::cout<<"RESULT "<<prefix<<std::endl;}
            std::string path;if(!(std::cin>>path>>prefix)){done=true;return;}
            std::ifstream f(path,std::ios::binary);if(std::filesystem::file_size(path)!=in_n)throw std::runtime_error("Worker image size");f.read(reinterpret_cast<char*>(px.data()),px.size());if(!f)throw std::runtime_error("Worker image read");
            std::fill(counts.begin(),counts.end(),0);std::fill(heat.begin(),heat.end(),0);
        }
    };
    auto pixels=[&](knp::core::Step t)->knp::core::messaging::SpikeData{
        advance(t);if(done)return {};int phase=t%steps_per_image;
        if(phase==4){std::vector<float>raster(in_n,0.F);for(auto&a:current)a.clear();for(int tick=0;tick<active_steps;tick++)for(int p=0;p<in_n;p++){raster[p]+=px[p]/255.F;if(raster[p]>=1.F){raster[p]-=1.F;current[tick].push_back(p);}}}
        if(phase<4||phase>=14)return {};return current[phase-4];
    };
    auto bias=[&](knp::core::Step t)->knp::core::messaging::SpikeData{advance(t);int phase=t%steps_per_image;return !done&&phase>=4&&phase<14?knp::core::messaging::SpikeData{0}:knp::core::messaging::SpikeData{};};
    auto reset=[&](knp::core::Step t)->knp::core::messaging::SpikeData{advance(t);return !done&&t%steps_per_image==0?knp::core::messaging::SpikeData{0}:knp::core::messaging::SpikeData{};};
    knp::framework::BackendLoader loader;knp::framework::ModelExecutor e(model,loader.load(argv[2]),{{pixel_channel,pixels},{bias_channel,bias},{reset_channel,reset}});ep=&e;
    e.start([&](knp::core::Step t){advance(t);return !done;});
}catch(const std::exception&e){std::cerr<<"PETS_WORKER_ERROR "<<e.what()<<std::endl;return 1;}