#include "altai_cuda_backend.h"
#include "altai_cuda_runtime.h"
#include <knp/meta/stringify.h>
#include <knp/meta/variant_helpers.h>
#include <algorithm>
#include <array>
#include <cmath>
#include <cstdlib>
#include <iostream>
#include <limits>

namespace knp_altai_cuda {
using Neuron = knp::neuron_traits::AltAILIF;
using Synapse = knp::synapse_traits::DeltaSynapse;
using Population = knp::core::Population<Neuron>;
using Projection = knp::core::Projection<Synapse>;
using Spike = knp::core::messaging::SpikeMessage;
using Impact = knp::core::messaging::SynapticImpactMessage;
using NeuronList = boost::mp11::mp_list<Neuron>;
using SynapseList = boost::mp11::mp_list<Synapse>;
using PopulationList = boost::mp11::mp_list<Population>;
using ProjectionList = boost::mp11::mp_list<Projection>;

class Device final : public knp::core::Device {
public:
    int index;
    explicit Device(int i) : index(i) {
        std::array<std::uint8_t,16> bytes{};
        bytes[0]=0xA1;bytes[1]=0x7A;bytes[12]=(i>>24)&255;bytes[13]=(i>>16)&255;bytes[14]=(i>>8)&255;bytes[15]=i&255;
        base_.uid_=knp::core::UID(bytes);
    }
    knp::core::DeviceType get_type() const override { return knp::core::DeviceType::GPU; }
    const std::string get_name() const override {
        CUdevice device;checked(cuDeviceGet(&device,index));
        char name[256];checked(cuDeviceGetName(name,sizeof(name),device));
        return std::string("CUDA ")+name;
    }
    float get_power() const override { return std::numeric_limits<float>::quiet_NaN(); }
};
struct Backend::State {
    struct Pop {
        Population *host;
        int n,slots=2;
        CUdeviceptr neurons=0,pending=0,absolute=0,events=0,spikes=0;
        std::vector<NeuronState> states;
        std::vector<unsigned char> flags;
    };
    struct Proj {
        Projection *host;
        size_t target;
        int sources=0,max_row=0;
        CUdeviceptr rows=0,source=0,weights=0,delays=0,counts=0;
        std::vector<int> mask;
    };
    Runtime runtime;
    std::vector<Pop> pops;
    std::vector<Proj> projections;
    CUdeviceptr error=0;
};
Backend::Backend() {
    if (const char *v=std::getenv("KNP_CUDA_DEVICE")) {
        std::string value(v);size_t end=0;device_index_=std::stoi(value,&end);
        if (end!=value.size() || device_index_<0) throw std::runtime_error("Invalid KNP_CUDA_DEVICE");
    }
}
Backend::~Backend() = default;
std::shared_ptr<knp::core::Backend> Backend::create() { return std::make_shared<Backend>(); }
std::vector<std::string> Backend::get_supported_neurons() const {
    return knp::meta::get_supported_type_names<knp::neuron_traits::AllNeurons,NeuronList>(knp::neuron_traits::neurons_names);
}
std::vector<std::string> Backend::get_supported_synapses() const {
    return knp::meta::get_supported_type_names<knp::synapse_traits::AllSynapses,SynapseList>(knp::synapse_traits::synapses_names);
}
std::vector<size_t> Backend::get_supported_population_indexes() const {
    return knp::meta::get_supported_type_indexes<knp::core::AllPopulations,PopulationList>();
}
std::vector<size_t> Backend::get_supported_projection_indexes() const {
    return knp::meta::get_supported_type_indexes<knp::core::AllProjections,ProjectionList>();
}
std::vector<std::unique_ptr<knp::core::Device>> Backend::get_devices() const {
    checked(cuInit(0));int count=0;checked(cuDeviceGetCount(&count));
    std::vector<std::unique_ptr<knp::core::Device>> result;
    for(int i=0;i<count;++i) result.push_back(std::make_unique<Device>(i));
    return result;
}
void Backend::select_devices(const std::set<knp::core::UID> &uids) {
    if(state_ || uids.size()!=1) throw std::runtime_error("Select exactly one CUDA device before initialization");
    for(auto &device:get_devices()) if(uids.count(device->get_uid())) {select_device(std::move(device));return;}
    throw std::runtime_error("Unknown CUDA device UID");
}
void Backend::select_device(std::unique_ptr<knp::core::Device> &&device) {
    auto *gpu=dynamic_cast<Device*>(device.get());
    if(state_ || !gpu) throw std::runtime_error("Select a device from this CUDA backend before initialization");
    device_index_=gpu->index;knp::core::Backend::select_device(std::move(device));
}
void Backend::load_all_populations(const std::vector<knp::core::AllPopulationsVariant> &populations) {
    if(state_) throw std::runtime_error("CUDA topology is immutable after initialization; create a new backend");
    for(const auto &p:populations) if(!std::holds_alternative<Population>(p))
        throw std::runtime_error("CUDA backend supports integer AltAILIF inference only");
    Base::load_all_populations(populations);
}
void Backend::load_all_projections(const std::vector<knp::core::AllProjectionsVariant> &projections) {
    if(state_) throw std::runtime_error("CUDA topology is immutable after initialization; create a new backend");
    for(const auto &p:projections) if(!std::holds_alternative<Projection>(p))
        throw std::runtime_error("CUDA backend supports DeltaSynapse only; no STDP");
    Base::load_all_projections(projections);
}
void Backend::remove_populations(const std::vector<knp::core::UID> &) {throw std::runtime_error("CUDA topology mutation unsupported");}
void Backend::remove_projections(const std::vector<knp::core::UID> &) {throw std::runtime_error("CUDA topology mutation unsupported");}
void Backend::start_learning() {throw std::runtime_error("CUDA AltAI inference backend has no plasticity");}
void Backend::_init() {
    if(state_) {state_->runtime.current();return;}
    auto state=std::make_unique<State>();
    state->runtime.initialize(device_index_);
    auto &runtime=state->runtime;
    state->error=runtime.allocate(sizeof(int));
    for(auto it=Base::begin_populations();it!=Base::end_populations();++it) {
        auto &population=std::get<Population>(*it);
        if(population.size()>10000000) throw std::runtime_error("CUDA population size limit exceeded");
        State::Pop p{};p.host=&population;p.n=static_cast<int>(population.size());
        for(auto &neuron:population) {
            if(!std::isfinite(neuron.potential_) || std::round(neuron.potential_)!=neuron.potential_ ||
               std::abs(neuron.potential_)>=16777216.F || neuron.additional_threshold_!=0)
                throw std::runtime_error("CUDA requires integer finite potential and zero additional threshold");
            unsigned flags=(neuron.is_diff_?1:0)|(neuron.is_reset_?2:0)|(neuron.leak_rev_?4:0)|
                           (neuron.saturate_?8:0)|(neuron.do_not_save_?16:0);
            p.states.push_back({neuron.potential_,neuron.pre_impact_potential_,neuron.activity_time_,
                               neuron.activation_threshold_,neuron.negative_activation_threshold_,neuron.potential_leak_,
                               neuron.potential_reset_value_,flags});
        }
        p.flags.resize(p.n);state->pops.push_back(std::move(p));
    }
    std::vector<std::vector<std::uint64_t>> absolute_bound;
    for(const auto &p:state->pops) absolute_bound.emplace_back(p.n,0);
    for(auto it=Base::begin_projections();it!=Base::end_projections();++it) {
        auto &projection=std::get<Projection>(it->arg_);
        size_t target=state->pops.size();
        for(size_t i=0;i<state->pops.size();++i) if(state->pops[i].host->get_uid()==projection.get_postsynaptic()) target=i;
        if(target==state->pops.size()) throw std::runtime_error("Projection target is not loaded");
        auto &population=state->pops[target];
        if(projection.size()>std::numeric_limits<std::uint32_t>::max()) throw std::runtime_error("CUDA projection too large");
        std::vector<unsigned> rows(population.n+1,0),sources(projection.size()),delays(projection.size());
        std::vector<int> weights(projection.size());
        int source_count=0;
        for(const auto &synapse:projection) {
            const auto &parameters=std::get<knp::core::synapse_data>(synapse);
            auto source=std::get<knp::core::source_neuron_id>(synapse),post=std::get<knp::core::target_neuron_id>(synapse);
            if(source>=10000000 || post>=static_cast<size_t>(population.n)) throw std::runtime_error("CUDA synapse index outside limits");
            if(!std::isfinite(parameters.weight_) || std::round(parameters.weight_)!=parameters.weight_ ||
               std::abs(parameters.weight_)>65535 || (parameters.weight_==0 && std::signbit(parameters.weight_)))
                throw std::runtime_error("CUDA requires finite integer weights with abs <=65535; negative zero unsupported");
            if(parameters.delay_<1 || parameters.delay_>64) throw std::runtime_error("CUDA supports delays 1..64");
            using Output=knp::synapse_traits::OutputType;
            if(parameters.output_type_!=Output::EXCITATORY && parameters.output_type_!=Output::INHIBITORY_CURRENT)
                throw std::runtime_error("CUDA supports excitatory and inhibitory-current Delta only");
            auto &bound=absolute_bound[target][post];
            bound+=static_cast<std::uint64_t>(std::abs(parameters.weight_));
            if(bound>static_cast<std::uint64_t>(std::numeric_limits<std::int64_t>::max())/(1024*65))
                throw std::runtime_error("CUDA integer queue accumulation bound exceeded");
            ++rows[post+1];source_count=std::max(source_count,static_cast<int>(source+1));
            population.slots=std::max(population.slots,static_cast<int>(parameters.delay_+1));
        }
        if(projection.get_presynaptic()) {
            auto pre=std::find_if(state->pops.begin(),state->pops.end(),[&](const auto &p){return p.host->get_uid()==projection.get_presynaptic();});
            if(pre==state->pops.end() || source_count>pre->n) throw std::runtime_error("CUDA invalid presynaptic population/index");
        }
        int max_row=0;
        for(int i=0;i<population.n;++i) {max_row=std::max(max_row,static_cast<int>(rows[i+1]));rows[i+1]+=rows[i];}
        auto cursor=rows;
        for(const auto &synapse:projection) {
            const auto &parameters=std::get<knp::core::synapse_data>(synapse);
            auto post=std::get<knp::core::target_neuron_id>(synapse);auto index=cursor[post]++;
            sources[index]=std::get<knp::core::source_neuron_id>(synapse);
            weights[index]=static_cast<int>(parameters.weight_)*(parameters.output_type_==knp::synapse_traits::OutputType::INHIBITORY_CURRENT?-1:1);
            delays[index]=parameters.delay_;
        }
        State::Proj p{};p.host=&projection;p.target=target;p.sources=source_count;p.max_row=max_row;p.mask.resize(source_count);
        p.rows=runtime.allocate(rows.size()*sizeof(unsigned),rows.data());
        p.source=runtime.allocate(sources.size()*sizeof(unsigned),sources.data());
        p.weights=runtime.allocate(weights.size()*sizeof(int),weights.data());
        p.delays=runtime.allocate(delays.size()*sizeof(unsigned),delays.data());
        p.counts=runtime.allocate(p.mask.size()*sizeof(int));state->projections.push_back(std::move(p));
    }
    for(auto &p:state->pops) {
        p.neurons=runtime.allocate(p.states.size()*sizeof(NeuronState),p.states.data());
        auto bytes=static_cast<size_t>(p.n)*p.slots*sizeof(std::int64_t);
        p.pending=runtime.allocate(bytes);p.absolute=runtime.allocate(bytes);p.events=runtime.allocate(bytes);
        p.spikes=runtime.allocate(p.flags.size());
    }
    Base::_init();
    state_=std::move(state);
    char name[256];checked(cuDeviceGetName(name,sizeof(name),state_->runtime.device));
    std::cerr<<"KNP AltAI CUDA integer inference initialized on "<<name<<"\n";
}
void Backend::_step() {
    if(!state_) throw std::runtime_error("CUDA backend not initialized");
    auto &s=*state_;auto &runtime=s.runtime;runtime.current();
    auto step=static_cast<unsigned long long>(get_step());
    auto &bus=get_message_bus();auto &endpoint=get_message_endpoint();
    bus.route_messages();endpoint.receive_all_messages();
    for(auto &p:s.pops) {
        if(!endpoint.unload_messages<Impact>(p.host->get_uid()).empty())
            throw std::runtime_error("Direct external impact messages unsupported by CUDA integer backend");
        void *arguments[]={&p.neurons,&p.pending,&p.absolute,&p.events,&p.spikes,&s.error,&p.n,&p.slots,&step};
        runtime.launch(runtime.population_kernel,p.n,arguments);
        if(p.n) checked(cuMemcpyDtoH(p.flags.data(),p.spikes,p.flags.size()));
        Spike message{{p.host->get_uid(),get_step()},{}};
        for(int i=0;i<p.n;++i) if(p.flags[i]) message.neuron_indexes_.push_back(i);
        if(!message.neuron_indexes_.empty()) endpoint.send_message(message);
    }
    int error=0;checked(cuMemcpyDtoH(&error,s.error,sizeof(error)));
    if(error) throw std::runtime_error("CUDA exactness guard: potential + absolute incoming impacts reaches 2^24");
    bus.route_messages();endpoint.receive_all_messages();
    for(auto &p:s.projections) {
        auto messages=endpoint.unload_messages<Spike>(p.host->get_uid());
        if(messages.empty()) continue;
        std::fill(p.mask.begin(),p.mask.end(),0);
        bool active=false;
        for(const auto &message:messages) for(auto index:message.neuron_indexes_) {
            // Spikes beyond this projection's used source range have no connections.
            if(index>=p.mask.size()) continue;
            if(++p.mask[index]>1024) throw std::runtime_error("CUDA input multiplicity exceeds 1024 per projection/step");
            active=true;
        }
        if(!active) continue;
        checked(cuMemcpyHtoD(p.counts,p.mask.data(),p.mask.size()*sizeof(int)));
        auto &post=s.pops[p.target];
        void *arguments[]={&p.rows,&p.source,&p.weights,&p.delays,&p.counts,&post.pending,&post.absolute,&post.events,&post.n,&post.slots,&step};
        if(p.max_row>512 && post.n) checked(cuLaunchKernel(runtime.wide_kernel,post.n,1,1,256,1,1,0,nullptr,arguments,nullptr));
        else runtime.launch(runtime.projection_kernel,post.n,arguments);
    }
    bus.route_messages();endpoint.receive_all_messages();
    (void)gad_step();
}
void Backend::synchronize() {
    if(!state_) return;
    state_->runtime.current();
    for(auto &p:state_->pops) {
        if(p.n) checked(cuMemcpyDtoH(p.states.data(),p.neurons,p.states.size()*sizeof(NeuronState)));
        for(int i=0;i<p.n;++i) {
            (*p.host)[i].potential_=p.states[i].potential;(*p.host)[i].pre_impact_potential_=p.states[i].pre;
            (*p.host)[i].activity_time_=p.states[i].activity;
        }
    }
}
knp::core::Backend::DataRanges Backend::get_network_data() const {
    const_cast<Backend*>(this)->synchronize();return Base::get_network_data();
}
void Backend::_uninit() {synchronize();Base::_uninit();}
BOOST_DLL_ALIAS(knp_altai_cuda::Backend::create, create_knp_backend)
}