#define main component_benchmark_main
#include "cuda_bench.cu"
#undef main

__global__ void expandJobs(const uint32_t* positionClass,const uint32_t* offsets,
                           const uint32_t* classOffsets,const uint32_t* roots,
                           uint32_t n,Job* jobs){
    uint32_t p=blockIdx.x*blockDim.x+threadIdx.x;
    if(p>=n)return;
    const uint32_t c=positionClass[p];
    for(uint32_t i=classOffsets[c];i<classOffsets[c+1];++i)
        jobs[offsets[p]+i-classOffsets[c]]={p,roots[i]};
}
struct Extra {
    Position* reps;uint32_t *positionClass,*offsets,*classOffsets,*roots;
    Extra(size_t n){
        CUDA(cudaMalloc(&reps,n*sizeof(Position)));
        CUDA(cudaMalloc(&positionClass,n*sizeof(uint32_t)));
        CUDA(cudaMalloc(&offsets,(n+1)*sizeof(uint32_t)));
        CUDA(cudaMalloc(&classOffsets,(n+1)*sizeof(uint32_t)));
        CUDA(cudaMalloc(&roots,JOB_CAPACITY*sizeof(uint32_t)));
    }
    ~Extra(){cudaFree(reps);cudaFree(positionClass);cudaFree(offsets);cudaFree(classOffsets);cudaFree(roots);}
};
struct Run {double hostMs,buildMs,rootMs,expandMs,searchMs;uint32_t classes,jobs,activePositions;std::vector<Match> matches;std::vector<Job> rootList;};
Run run(const std::vector<std::pair<int32_t,int32_t>>& coords,Device& d,Extra& extra,int mode){
    const auto start=Clock::now();
    auto ms=[](auto t){return std::chrono::duration<double,std::milli>(Clock::now()-t).count();};
    std::vector<Position> positions;positions.reserve(coords.size());
    if(mode!=3)for(auto [x,z]:coords)positions.push_back(makePosition(x,z));
    std::vector<Position> reps;std::vector<uint32_t> pc;
    if(mode>=2){
        std::map<uint64_t,uint32_t> map;
        for(auto [x,z]:coords){
            uint64_t key=(uint64_t(uint32_t(x)&65535)<<17)|(uint32_t(z)&131071);
            auto [it,inserted]=map.emplace(key,uint32_t(reps.size()));
            if(inserted)reps.push_back(makePosition(x,z));
            pc.push_back(it->second);
        }
    }
    if(!positions.empty())CUDA(cudaMemcpy(d.positions,positions.data(),positions.size()*sizeof(Position),cudaMemcpyHostToDevice));
    Run r{};r.classes=mode>=2?reps.size():positions.size();r.buildMs=ms(start);
    const auto rootStart=Clock::now();
    if(mode>=2){
        CUDA(cudaMemcpy(extra.reps,reps.data(),reps.size()*sizeof(Position),cudaMemcpyHostToDevice));
        CUDA(cudaMemset(d.total,0,sizeof(*d.total)));CUDA(cudaMemset(d.counts,0,reps.size()*sizeof(uint32_t)));
        rootKernel<0><<<unsigned(((uint64_t(reps.size())<<18)+255)/256),256>>>(extra.reps,reps.size(),nullptr,d.counts,d.jobs,d.total,true);
    }else d.roots(positions.size(),0,true);
    uint64_t n=d.totalCount();if(n>JOB_CAPACITY)throw std::runtime_error("root capacity exceeded");
    r.rootMs=ms(rootStart);
    const auto expandStart=Clock::now();
    std::vector<uint32_t> originalIndices;
    if(mode>=2){
        std::vector<Job> cj(n);if(n)CUDA(cudaMemcpy(cj.data(),d.jobs,n*sizeof(Job),cudaMemcpyDeviceToHost));
        std::vector<std::vector<uint32_t>> lists(reps.size());
        for(auto j:cj)lists[j.position].push_back(j.root);
        if(mode==3){
            std::vector<uint32_t> compactClasses;
            for(uint32_t i=0;i<coords.size();++i)if(!lists[pc[i]].empty()){
                auto [x,z]=coords[i];positions.push_back(makePosition(x,z));
                compactClasses.push_back(pc[i]);originalIndices.push_back(i);
            }
            pc=std::move(compactClasses);
            if(!positions.empty())CUDA(cudaMemcpy(d.positions,positions.data(),positions.size()*sizeof(Position),cudaMemcpyHostToDevice));
        }
        std::vector<uint32_t> roots,co{0},offsets{0};
        for(auto& list:lists){std::sort(list.begin(),list.end());roots.insert(roots.end(),list.begin(),list.end());co.push_back(roots.size());}
        uint64_t expanded=0;
        for(auto c:pc){expanded+=lists[c].size();if(expanded>JOB_CAPACITY)throw std::runtime_error("expanded capacity exceeded; use smaller batch");offsets.push_back(expanded);}
        n=expanded;
        CUDA(cudaMemcpy(extra.positionClass,pc.data(),pc.size()*sizeof(uint32_t),cudaMemcpyHostToDevice));
        CUDA(cudaMemcpy(extra.offsets,offsets.data(),offsets.size()*sizeof(uint32_t),cudaMemcpyHostToDevice));
        CUDA(cudaMemcpy(extra.classOffsets,co.data(),co.size()*sizeof(uint32_t),cudaMemcpyHostToDevice));
        if(!roots.empty())CUDA(cudaMemcpy(extra.roots,roots.data(),roots.size()*sizeof(uint32_t),cudaMemcpyHostToDevice));
        if(!positions.empty())expandJobs<<<unsigned((positions.size()+255)/256),256>>>(extra.positionClass,extra.offsets,extra.classOffsets,extra.roots,positions.size(),d.jobs);
        CUDA(cudaGetLastError());CUDA(cudaDeviceSynchronize());
    }
    r.expandMs=ms(expandStart);r.jobs=n;r.activePositions=positions.size();
    const auto searchStart=Clock::now();
    d.search(n,mode==0?3:7);r.matches=d.results();
    if(mode==3)for(auto& match:r.matches)match.position=originalIndices[match.position];
    r.searchMs=ms(searchStart);r.hostMs=ms(start);
    // Validation-only transfers/sorting are outside reported pipeline time.
    r.rootList.resize(n);if(n)CUDA(cudaMemcpy(r.rootList.data(),d.jobs,n*sizeof(Job),cudaMemcpyDeviceToHost));
    if(mode==3)for(auto& job:r.rootList)job.position=originalIndices[job.position];
    std::sort(r.rootList.begin(),r.rootList.end(),[](Job a,Job b){return std::tie(a.position,a.root)<std::tie(b.position,b.root);});
    return r;
}
int main(int argc,char**argv){try{
    if(argc<2)throw std::runtime_error("usage: pipeline_bench positions.txt [repeats=5]");
    const auto process=Clock::now();
    std::ifstream in(argv[1]);if(!in)throw std::runtime_error("cannot open input");
    std::vector<std::pair<int32_t,int32_t>> coords;int32_t x,z;while(in>>x>>z)coords.emplace_back(x,z);
    if(coords.empty()||coords.size()>65536)throw std::runtime_error("expected 1..65536 coordinates");
    int repeats=argc>2?std::stoi(argv[2]):5;if(repeats<1)throw std::runtime_error("invalid repeats");
    Device d(std::vector<Position>(coords.size()));Extra e(coords.size());
    std::cout<<std::setprecision(10);
    std::cout<<"{\"type\":\"setup\",\"host_ms\":"<<std::chrono::duration<double,std::milli>(Clock::now()-process).count()<<"}"<<std::endl;
    Run ref;
    for(int round=-1;round<repeats;++round)for(int k=0;k<4;++k){
        int mode=round<0?k:(round+k)%4;
        Run r=run(coords,d,e,mode);
        if(round==-1&&mode==0)ref=r;
        else {
            if(r.matches!=ref.matches||r.rootList.size()!=ref.rootList.size())throw std::runtime_error("pipeline mismatch");
            for(size_t i=0;i<r.rootList.size();++i)if(r.rootList[i].position!=ref.rootList[i].position||r.rootList[i].root!=ref.rootList[i].root)throw std::runtime_error("complete root list mismatch");
        }
        if(round>=0)std::cout<<"{\"type\":\"pipeline\",\"mode\":"<<mode<<",\"round\":"<<round<<",\"positions\":"<<coords.size()<<",\"classes\":"<<r.classes<<",\"active_positions\":"<<r.activePositions<<",\"jobs\":"<<r.jobs<<",\"matches\":"<<r.matches.size()<<",\"host_ms\":"<<r.hostMs<<",\"build_ms\":"<<r.buildMs<<",\"root_ms\":"<<r.rootMs<<",\"expand_ms\":"<<r.expandMs<<",\"search_ms\":"<<r.searchMs<<"}"<<std::endl;
    }
    std::cout<<"{\"type\":\"validation\",\"complete_roots_equal\":true,\"complete_matches_equal\":true,\"matches\":[";
    for(size_t i=0;i<ref.matches.size();++i){if(i)std::cout<<",";auto m=ref.matches[i];std::cout<<"["<<m.seed<<","<<coords[m.position].first<<","<<coords[m.position].second<<"]";}
    std::cout<<"]}"<<std::endl;
    return 0;
}catch(const std::exception& e){std::cerr<<e.what()<<"\n";return 1;}}