// Toy insect controller. Hand-wired AltAILIF/Delta, NOT a fly connectome.
#include <knp/backends/altai/backend.h>
#include <knp/devices/altai_npu.h>
#include <knp/framework/model_executor.h>
#include <knp/framework/network.h>
#include <algorithm>
#include <array>
#include <iostream>
#include <sstream>
#include <vector>
using N=knp::neuron_traits::neuron_parameters<knp::neuron_traits::AltAILIF>;
using P=knp::core::Population<knp::neuron_traits::AltAILIF>;
using R=knp::core::Projection<knp::synapse_traits::DeltaSynapse>;
constexpr int FRAME=48,ACTIVE=12,FIRST=4,INPUTS=9,RESET=8;
// Input: foodL foodR foodFront dangerL dangerR dangerFront wanderL wanderR drive.
const int HW[6][INPUTS]={{4,0,0,0,0,0,0,0,0},{0,4,0,0,0,0,0,0,0},{0,0,0,0,8,4,0,0,0},{0,0,0,8,0,2,0,0,0},{0,0,0,0,0,0,4,0,0},{0,0,0,0,0,0,0,4,0}};
// Output: turnLeft turnRight drive. All non-reset magnitudes per cell <=3.
const int MW[3][6]={{4,-2,8,-2,2,-2},{-2,4,-2,8,-2,2},{0,0,0,0,0,0}};
const int DW[3][INPUTS]={{0,0,0,0,0,0,0,0,0},{0,0,0,0,0,0,0,0,0},{0,0,2,-2,-2,-8,0,0,4}};
int main(int argc,char**argv)try{
 std::string mode=argc>1?argv[1]:"cpu";if(mode!="cpu"&&mode!="cuda"&&mode!="gm")throw std::runtime_error("Invalid backend");
 auto neuron=[](int thr){N n;n.activation_threshold_=thr;n.negative_activation_threshold_=0;n.is_diff_=true;n.is_reset_=false;n.saturate_=true;n.potential_leak_=0;return n;};
 P hp([&](size_t){return neuron(4);},6),mp([&](size_t){return neuron(8);},3);auto hid=hp.get_uid(),mid=mp.get_uid();std::vector<R::Synapse>hs,ms,links;
 auto add=[](auto&v,int i,int o,int w){if(w)v.push_back({{float(w),1,knp::synapse_traits::OutputType::EXCITATORY},size_t(i),size_t(o)});};
 for(int h=0;h<6;h++){for(int i=0;i<INPUTS;i++)add(hs,i,h,HW[h][i]);for(int j=0;j<RESET;j++)add(hs,INPUTS+j,h,-255);}
 for(int o=0;o<3;o++){for(int i=0;i<INPUTS;i++)add(ms,i,o,DW[o][i]);for(int j=0;j<RESET;j++)add(ms,INPUTS+j,o,-255);for(int h=0;h<6;h++)add(links,h,o,MW[o][h]);}
 auto project=[](knp::core::UID pre,knp::core::UID post,auto&v){return R(pre,post,[&v](size_t i)->std::optional<R::Synapse>{return v[i];},v.size());};
 auto ih=project(knp::core::UID{false},hid,hs),im=project(knp::core::UID{false},mid,ms),hm=project(hid,mid,links);auto ihid=ih.get_uid(),imid=im.get_uid();std::array<bool,INPUTS+RESET>hu{},mu{};for(auto&s:hs)hu[std::get<1>(s)]=true;for(auto&s:ms)mu[std::get<1>(s)]=true;
 knp::framework::Network net;net.add_population(std::move(hp));net.add_population(std::move(mp));net.add_projection(std::move(ih));net.add_projection(std::move(im));net.add_projection(std::move(hm));knp::framework::Model model(std::move(net));knp::core::UID hi,mi,ho,mo;model.add_input_channel(hi,ihid);model.add_input_channel(mi,imid);model.add_output_channel(ho,hid);model.add_output_channel(mo,mid);
 knp::framework::BackendLoader loader;std::shared_ptr<knp::core::Backend>backend;
 if(mode=="gm"){auto gm=knp::backends::altai::AltAIBackend::create();bool found=false;for(auto&d:gm->get_devices())if(dynamic_cast<knp::devices::altai::AltAI_GM*>(d.get())){gm->select_device(std::move(d));found=true;break;}if(!found)throw std::runtime_error("No official GoldenModel");backend=gm;}
 else backend=loader.load("/home/user/KNP/build/bin/libknp-"+std::string(mode=="cuda"?"altai-cuda":"cpu-single-threaded")+"-backend.so");
 std::array<int,INPUTS>values{},acc{};std::array<int,9>counts{};std::vector<std::pair<int,int>>raster;knp::framework::ModelExecutor*ep=nullptr;size_t last=size_t(-1);bool done=false;
 auto advance=[&](knp::core::Step t){
  if(last==t)return;last=t;
  if(ep&&t){auto collect=[&](auto uid,int offset){for(auto&m:ep->get_loader().get_output_channel(uid).read_some_from_buffer(t-1,t-1)){
   int phase=m.header_.send_time_%FRAME;if(phase>=FIRST)for(auto i:m.neuron_indexes_){++counts[offset+i];raster.emplace_back(phase,offset+i);}
  }};collect(ho,0);collect(mo,6);}
  if(t%FRAME==0){
   if(t){std::cout<<"RESULT {\"hidden\":[";for(int i=0;i<6;i++)std::cout<<(i?",":"")<<counts[i];std::cout<<"],\"motors\":[";for(int i=0;i<3;i++)std::cout<<(i?",":"")<<counts[6+i];std::cout<<"],\"raster\":[";for(size_t i=0;i<raster.size();i++)std::cout<<(i?",":"")<<"["<<raster[i].first<<","<<raster[i].second<<"]";std::cout<<"]}"<<std::endl;}
   std::string line;if(!std::getline(std::cin,line)){done=true;return;}std::istringstream in(line);for(auto&v:values)if(!(in>>v)||v<0||v>15)throw std::runtime_error("Nine integers0..15 required");std::string extra;if(in>>extra)throw std::runtime_error("Extra input");counts.fill(0);acc.fill(0);raster.clear();
  }
 };
 size_t rasterstep=size_t(-1);std::vector<int>inputs;
 auto input=[&](knp::core::Step t,const auto&used){advance(t);knp::core::messaging::SpikeData data;if(done)return data;int phase=t%FRAME;
  if(phase==0){for(int i=0;i<RESET;i++)if(used[INPUTS+i])data.push_back(INPUTS+i);}
  if(phase>=FIRST&&phase<FIRST+ACTIVE){if(rasterstep!=t){rasterstep=t;inputs.clear();for(int i=0;i<INPUTS;i++){acc[i]+=values[i];if(acc[i]>=15){acc[i]-=15;inputs.push_back(i);}}}for(auto i:inputs)if(used[i])data.push_back(i);}
  return data;
 };
 knp::framework::ModelExecutor e(model,backend,{{hi,[&](auto t){return input(t,hu);}},{mi,[&](auto t){return input(t,mu);}}});ep=&e;e.start([&](knp::core::Step t){advance(t);return !done;});
}catch(const std::exception&e){std::cerr<<"FLY_ERROR "<<e.what()<<std::endl;return 1;}