#include <knp/backends/altai/backend.h>
#include <knp/devices/altai_npu.h>
#include <knp/framework/model_executor.h>
#include <knp/framework/backend_loader.h>
#include <knp/framework/network.h>
#include <fstream>
#include <iostream>
#include <filesystem>
#include <chrono>
#include <algorithm>
#include <numeric>
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>;
namespace fs=std::filesystem;
template<class T>std::vector<T> read(std::ifstream&f,size_t n){std::vector<T>a(n);f.read(reinterpret_cast<char*>(a.data()),n*sizeof(T));if(!f)throw std::runtime_error("Truncated model/data");return a;}
void progress(const fs::path&dir,int cur,int total,bool done=false,bool failed=false){std::ofstream f(dir/"progress.json.tmp");f<<"{\"current\":"<<cur<<",\"total\":"<<total<<",\"done\":"<<(done?"true":"false")<<",\"failed\":"<<(failed?"true":"false")<<"}\n";f.close();fs::rename(dir/"progress.json.tmp",dir/"progress.json");}
int main(int argc,char**argv){fs::path out;int completed=0,total=0;try{
 if(argc<8)throw std::runtime_error("modeldir outdir first count replicas gm|cpu|cuda indices-file-or-dash");
 fs::path root=argv[1];out=argv[2];fs::create_directories(out);int first=std::stoi(argv[3]);total=std::stoi(argv[4]);int reps=std::stoi(argv[5]);std::string mode=argv[6];if(reps<1||reps>64||total<1||first<0)throw std::runtime_error("Invalid range");
 std::vector<int>ids(total);std::iota(ids.begin(),ids.end(),first);if(std::string(argv[7])!="-"){std::ifstream f(argv[7]);for(auto&i:ids)if(!(f>>i))throw std::runtime_error("Missing id");}
 std::ifstream mf(root/"model.bin",std::ios::binary);auto hdr=read<int32_t>(mf,8);int pix=hdr[0],hidden=hdr[1],classes=hdr[2],groups=hdr[3],active=hdr[4],hthr=hdr[5],othr=hdr[6],clock=hdr[7],outs=2*groups*classes;
 if(pix!=784||hidden!=256||classes!=10||groups!=8||clock>255)throw std::runtime_error("Unsupported model dimensions");
 auto w1=read<int16_t>(mf,pix*hidden),b1=read<int16_t>(mf,hidden),w2=read<int16_t>(mf,hidden*classes),b2=read<int16_t>(mf,classes);
 std::ifstream imf(root/"images.bin",std::ios::binary),lf(root/"labels.bin",std::ios::binary);auto labels=read<int32_t>(lf,70000);std::vector<uint8_t>images(total*pix);
 for(int i=0;i<total;i++){if(ids[i]<0||ids[i]>=70000)throw std::runtime_error("image id out of bounds");imf.seekg(size_t(ids[i])*pix);imf.read(reinterpret_cast<char*>(images.data()+i*pix),pix);if(!imf)throw std::runtime_error("Truncated pixels");}
 auto make=[&](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 make(hthr);},hidden*reps),op([&](size_t){return make(othr);},outs*reps);auto huid=hp.get_uid(),ouid=op.get_uid();
 int hs=pix+clock+128,os=clock+128;std::vector<R::Synapse>hin,oin,mid;hin.reserve(reps*(pix*hidden+128*hidden));oin.reserve(reps*128*outs);mid.reserve(reps*hidden*classes);
 auto add=[](auto&v,int pre,int post,int weight,bool inhibitory=false){if(!weight)return;if(std::abs(weight)>255)throw std::runtime_error("Weight outside native range");if(inhibitory)weight=-std::abs(weight);v.push_back(R::Synapse{{float(weight),1,knp::synapse_traits::OutputType::EXCITATORY},size_t(pre),size_t(post)});};
 for(int r=0;r<reps;r++){
  for(int h=0;h<hidden;h++){
   for(int p=0;p<pix;p++)add(hin,r*hs+p,r*hidden+h,w1[p*hidden+h]);
   for(int c=0;c<std::abs(b1[h]);c++)add(hin,r*hs+pix+c,r*hidden+h,b1[h]>0?1:-1);
   for(int c=0;c<128;c++)add(hin,r*hs+pix+clock+c,r*hidden+h,255,true);
   for(int k=0;k<classes;k++){int w=w2[h*classes+k];int o=(w>=0?0:groups*classes)+(h/32)*classes+k;add(mid,r*hidden+h,r*outs+o,std::abs(w));}
  }
  for(int k=0;k<classes;k++)for(int c=0;c<std::abs(b2[k]);c++)add(oin,r*os+c,r*outs+(b2[k]>=0?0:groups*classes)+k,1);
  for(int o=0;o<outs;o++)for(int c=0;c<128;c++)add(oin,r*os+clock+c,r*outs+o,255,true);
 }
 std::vector<bool>hused(reps*hs),oused(reps*os);for(auto &z:hin)hused[std::get<1>(z)]=true;for(auto &z:oin)oused[std::get<1>(z)]=true;
 auto projection=[](knp::core::UID pre,knp::core::UID post,auto&vec){return R(pre,post,[&vec](size_t i)->std::optional<R::Synapse>{return vec[i];},vec.size());};
 auto ih=projection(knp::core::UID{false},huid,hin),io=projection(knp::core::UID{false},ouid,oin),ho=projection(huid,ouid,mid);auto ihid=ih.get_uid(),ioid=io.get_uid();
 knp::framework::Network net;net.add_population(std::move(hp));net.add_population(std::move(op));net.add_projection(std::move(ih));net.add_projection(std::move(io));net.add_projection(std::move(ho));knp::framework::Model model(std::move(net));knp::core::UID inch,incout,outch,outh;model.add_input_channel(inch,ihid);model.add_input_channel(incout,ioid);model.add_output_channel(outch,ouid);model.add_output_channel(outh,huid);
 std::shared_ptr<knp::core::Backend>backend;knp::framework::BackendLoader loader;
 if(mode=="gm"){auto gm=knp::backends::altai::AltAIBackend::create();auto dev=gm->get_devices();bool selected=false;for(auto&d:dev)if(dynamic_cast<knp::devices::altai::AltAI_GM*>(d.get())){std::cout<<"DEVICE "<<d->get_name()<<std::endl;gm->select_device(std::move(d));selected=true;break;}if(!selected)throw std::runtime_error("No official GoldenModel");backend=gm;}
 else backend=loader.load(mode=="cuda"?"/home/user/KNP/build/bin/libknp-altai-cuda-backend.so":"/home/user/KNP/build/bin/libknp-cpu-single-threaded-backend.so");
 std::ofstream result(out/"predictions.csv"),hcfile(out/"hidden_counts.bin",std::ios::binary),ocfile(out/"output_counts.bin",std::ios::binary);if(!result||!hcfile||!ocfile)throw std::runtime_error("Cannot open outputs");result<<"row,image_id,label,prediction\n";
 std::vector<int32_t>hc(reps*hidden),oc(reps*outs);std::vector<int>raster(reps*pix);int batch=0,phase=0,tick=0,quiet=0;size_t last=size_t(-1);bool done=false;int correct=0;knp::framework::ModelExecutor*ep=nullptr;
 auto update=[&](size_t step){
  if(last==step)return;last=step;bool any=false;
  if(step&&ep){for(auto&m:ep->get_loader().get_output_channel(outch).read_some_from_buffer(step-1,step-1))for(auto i:m.neuron_indexes_){if(i>=oc.size())throw std::runtime_error("Output index overflow");if(phase<=2){++oc[i];any=true;}}
   for(auto&m:ep->get_loader().get_output_channel(outh).read_some_from_buffer(step-1,step-1))for(auto i:m.neuron_indexes_)if(phase<=2){if(i>=hc.size())throw std::runtime_error("Hidden index overflow");++hc[i];}
  }
  if(step==0)return;
  if(phase==0){if(++tick>=active){phase=1;tick=0;}}
  else if(phase==1){phase=2;tick=0;quiet=0;}
  else if(phase==2){++tick;quiet=any?0:quiet+1;if(tick>10000)throw std::runtime_error("Output failed to drain");if(tick>=3&&quiet>=2){
    int n=std::min(reps,total-batch);for(int r=0;r<n;r++){int best=0;int64_t bs=INT64_MIN;for(int k=0;k<classes;k++){int64_t score=0;for(int g=0;g<groups;g++)score+=oc[r*outs+g*classes+k]-oc[r*outs+groups*classes+g*classes+k];if(score>bs){bs=score;best=k;}}
     int ix=batch+r;result<<ix<<","<<ids[ix]<<","<<labels[ids[ix]]<<","<<best<<"\n";correct+=best==labels[ids[ix]];
    }hcfile.write(reinterpret_cast<char*>(hc.data()),n*hidden*sizeof(int32_t));ocfile.write(reinterpret_cast<char*>(oc.data()),n*outs*sizeof(int32_t));result.flush();completed=batch+n;progress(out,completed,total);std::cout<<"BATCH "<<completed<<"/"<<total<<" correct "<<correct<<" steps "<<step<<std::endl;phase=3;tick=0;
  }}else if(phase==3){phase=4;tick=0;}
  else if(phase==4){if(++tick>=3){batch+=reps;if(batch>=total){done=true;return;}std::fill(hc.begin(),hc.end(),0);std::fill(oc.begin(),oc.end(),0);std::fill(raster.begin(),raster.end(),0);phase=0;tick=0;}}
 };
 auto hidden_input=[&](size_t step){update(step);knp::core::messaging::SpikeData data;if(done)return data;int n=std::min(reps,total-batch);for(int r=0;r<n;r++){
  if(phase==0){for(int p=0;p<pix;p++){auto&v=raster[r*pix+p];v+=images[(batch+r)*pix+p];if(v>=255){v-=255;if(hused[r*hs+p])data.push_back(r*hs+p);}}for(int c=0;c<clock;c++)if(hused[r*hs+pix+c])data.push_back(r*hs+pix+c);}
  else if(phase==1)for(int c=0;c<128;c++)if(hused[r*hs+pix+clock+c])data.push_back(r*hs+pix+clock+c);
 }return data;};
 auto output_input=[&](size_t step){update(step);knp::core::messaging::SpikeData data;if(done)return data;int n=std::min(reps,total-batch);for(int r=0;r<n;r++){if(phase==0){for(int c=0;c<clock;c++)if(oused[r*os+c])data.push_back(r*os+c);}else if(phase==3){for(int c=0;c<128;c++)if(oused[r*os+clock+c])data.push_back(r*os+clock+c);}}return data;};
 progress(out,0,total);auto start=std::chrono::steady_clock::now();knp::framework::ModelExecutor exec(model,backend,{{inch,hidden_input},{incout,output_input}});ep=&exec;exec.start([&](size_t step){return !done&&step<size_t(total)*(active+10016);});
 hcfile.close();ocfile.close();result.close();if(!hcfile||!ocfile||!result||!done)throw std::runtime_error("Run incomplete or write failed");double sec=std::chrono::duration<double>(std::chrono::steady_clock::now()-start).count();std::ofstream summary(out/"summary.json");summary<<"{\"backend\":\""<<mode<<"\",\"total\":"<<total<<",\"correct\":"<<correct<<",\"wall_seconds\":"<<sec<<",\"replicas\":"<<reps<<"}\n";summary.close();if(!summary)throw std::runtime_error("Summary write failed");progress(out,total,total,true);std::cout<<"DONE "<<correct<<"/"<<total<<" wall "<<sec<<std::endl;return 0;
 }catch(const std::exception&e){std::cerr<<"NATIVE_ERROR "<<e.what()<<std::endl;if(!out.empty())progress(out,completed,total,false,true);return 1;}}