#include "grouped.hpp"
#include <cuda_runtime.h>
#include <algorithm>
#include <array>
#include <chrono>
#include <cstdint>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <stdexcept>
#include <string>
#include <tuple>
#include <vector>

namespace sr = slime_research;
namespace fs = full_search_experiment;
using Clock = std::chrono::steady_clock;
#define CUDA(call) do { auto e = (call); if (e != cudaSuccess) throw std::runtime_error(std::string(#call)+": "+cudaGetErrorString(e)); } while (0)
struct Position { int32_t x, z; uint64_t anchor; sr::CellTerm cells[25]; };
struct Job { uint32_t position, root; };
struct Match { uint64_t position, seed; bool operator<(const Match& o)const{return std::tie(position,seed)<std::tie(o.position,o.seed);} bool operator==(const Match& o)const{return position==o.position&&seed==o.seed;} };
constexpr uint32_t MATCH_CAPACITY = 1U << 20;
constexpr uint32_t JOB_CAPACITY = 1U << 22;

Position makePosition(int32_t x,int32_t z) {
    Position p{}; p.x=x;p.z=z;p.anchor=sr::coordinateTerm(x,z);
    for(int j=0;j<5;++j)for(int i=0;i<5;++i)
        p.cells[j*5+i]={(sr::coordinateTerm(x+i,z+j)-p.anchor)&sr::kMask48,sr::kSlime};
    return p;
}

template<int Mode>
__global__ void rootKernel(const Position* positions,uint32_t count,const uint32_t* good,
                           uint32_t* counts,Job* jobs,unsigned long long* total,bool collect) {
    uint64_t idx=uint64_t(blockIdx.x)*blockDim.x+threadIdx.x;
    if(idx >= (uint64_t(count)<<18)) return;
    const uint32_t pos=uint32_t(idx>>18), root=uint32_t(idx)&((1U<<18)-1);
    const Position& p=positions[pos];
    for(int i=0;i<25;++i){
        const uint32_t in=root+uint32_t(p.cells[i].term);
        if constexpr(Mode==0){
            if((((in^uint32_t(sr::kCombinedXor))*uint32_t(sr::kJavaMultiplier)+11U)>>17)&1U)return;
        }else{
            const uint32_t k=in&((1U<<18)-1);
            if(((good[k>>5]>>(k&31))&1U)==0)return;
        }
    }
    atomicAdd(counts+pos,1U);
    const auto at=atomicAdd(total,1ULL);
    if(collect&&at<JOB_CAPACITY)jobs[at]={pos,root};
}

struct Emitter {
    uint32_t position; uint64_t anchor; Match* out; unsigned long long* count;
    __device__ void operator()(uint64_t s){
        auto at=atomicAdd(count,1ULL);
        if(at<MATCH_CAPACITY)out[at]={position,(s-anchor)&sr::kMask48};
    }
};

using normal_search::Prepared;
using normal_search::preparedQ;
__global__ void prepareKernel(const Position* positions,const Job* jobs,uint32_t count,Prepared* out){
    const uint32_t j=blockIdx.x*blockDim.x+threadIdx.x;
    if(j<count)normal_search::prepare(jobs[j].root,positions[jobs[j].position].cells,out[j]);
}
template<int Mode>
__global__ void searchKernel(const Position* positions,const Job* jobs,uint32_t count,
                             Match* out,unsigned long long* total,const Prepared* prepared) {
    uint64_t idx=uint64_t(blockIdx.x)*blockDim.x+threadIdx.x;
    constexpr int bits=Mode==7?13:17;
    if(idx >= (uint64_t(count)<<bits))return;
    const Job job=jobs[idx>>bits];
    const uint32_t m=idx&((1U<<bits)-1);
    uint64_t low=job.root+(uint64_t(m)<<18);
    const Position& p=positions[job.position];
    Emitter emit{job.position,p.anchor,out,total};
    if constexpr(Mode==7){
        const uint32_t base=((m>>4)<<8)+(m&15);
        uint32_t mask=normal_search::groupMask(prepared[idx>>13],base);
        while(mask){
            const uint32_t j=__ffs(mask)-1;mask &= mask-1;
            const uint64_t candidate=job.root+(uint64_t(base+16*j)<<18);
            fs::solveIntervals(candidate,p.cells,25,emit);
        }
    }
    else if constexpr(Mode==3){if(fs::differenceFilter(low,p.cells,25))fs::solveIntervals(low,p.cells,25,emit);}
    else {
        const Prepared& t=prepared[idx>>17];
        const uint32_t q0=preparedQ(t,0,m);
        const uint32_t u0=(m+t.h[0])^uint32_t(sr::kCombinedXor>>18);
        const int order[24]={24,4,20,12,2,10,14,22,6,8,16,18,1,3,5,9,15,19,21,23,7,11,13,17};
        for(int k=1;k<25;++k){
            const int i=Mode==6?order[k-1]:k;
            if constexpr(Mode==4){
                const uint32_t r=((preparedQ(t,i,m)-q0)&0x7fffffffU)%10;
                if(r!=0&&r!=8)return;
            }else{
                const uint32_t ui=(m+t.h[i])^uint32_t(sr::kCombinedXor>>18);
                const uint32_t d=((t.b[i]>>1)-(t.b[0]>>1)+uint32_t(sr::kJavaMultiplier)*(ui-u0))&((1U<<30)-1);
                const uint32_t r=d%5;
                if(r!=0&&r!=4)return;
            }
        }
        fs::solveIntervals(low,p.cells,25,emit);
    }
}
struct Device {
    Prepared* prepared=nullptr; Position* positions=nullptr; Job* jobs=nullptr; Match* matches=nullptr;
    uint32_t* counts=nullptr;uint32_t* good=nullptr;unsigned long long* total=nullptr;
    cudaEvent_t start,stop;
    Device(const std::vector<Position>& p){
        CUDA(cudaMalloc(&positions,p.size()*sizeof(Position)));CUDA(cudaMemcpy(positions,p.data(),p.size()*sizeof(Position),cudaMemcpyHostToDevice));
        CUDA(cudaMalloc(&jobs,JOB_CAPACITY*sizeof(Job)));CUDA(cudaMalloc(&matches,MATCH_CAPACITY*sizeof(Match)));
        CUDA(cudaMalloc(&counts,p.size()*sizeof(uint32_t)));CUDA(cudaMalloc(&total,sizeof(*total)));
        std::array<uint32_t,8192> bits{};
        for(uint32_t k=0;k<(1U<<18);++k)if(!((((k^uint32_t(sr::kCombinedXor))*uint32_t(sr::kJavaMultiplier)+11U)>>17)&1U))bits[k>>5]|=1U<<(k&31);
        CUDA(cudaMalloc(&good,sizeof(bits)));CUDA(cudaMemcpy(good,bits.data(),sizeof(bits),cudaMemcpyHostToDevice));
        CUDA(cudaEventCreate(&start));CUDA(cudaEventCreate(&stop));
    }
    ~Device(){cudaFree(prepared);cudaFree(positions);cudaFree(jobs);cudaFree(matches);cudaFree(counts);cudaFree(total);cudaFree(good);cudaEventDestroy(start);cudaEventDestroy(stop);}
    float finish(){CUDA(cudaGetLastError());CUDA(cudaEventRecord(stop));CUDA(cudaEventSynchronize(stop));float ms;CUDA(cudaEventElapsedTime(&ms,start,stop));return ms;}
    uint64_t totalCount(){unsigned long long n;CUDA(cudaMemcpy(&n,total,sizeof(n),cudaMemcpyDeviceToHost));return n;}
    float roots(uint32_t n,int mode,bool collect){
        CUDA(cudaMemset(total,0,sizeof(*total)));CUDA(cudaMemset(counts,0,n*sizeof(uint32_t)));CUDA(cudaEventRecord(start));
        unsigned blocks=unsigned(((uint64_t(n)<<18)+255)/256);
        if(mode==0)rootKernel<0><<<blocks,256>>>(positions,n,good,counts,jobs,total,collect);
        else rootKernel<1><<<blocks,256>>>(positions,n,good,counts,jobs,total,collect);
        return finish();
    }
    float search(uint32_t n,int mode){
        if(!prepared&&n)CUDA(cudaMalloc(&prepared,n*sizeof(Prepared)));
        CUDA(cudaMemset(total,0,sizeof(*total)));CUDA(cudaEventRecord(start));
        unsigned blocks=unsigned(((uint64_t(n)<<17)+255)/256);
        if(blocks){
            if(mode>=4)prepareKernel<<<(n+127)/128,128>>>(positions,jobs,n,prepared);
            if(mode==3)searchKernel<3><<<blocks,256>>>(positions,jobs,n,matches,total,prepared);
            if(mode==4)searchKernel<4><<<blocks,256>>>(positions,jobs,n,matches,total,prepared);
            if(mode==5)searchKernel<5><<<blocks,256>>>(positions,jobs,n,matches,total,prepared);
            if(mode==6)searchKernel<6><<<blocks,256>>>(positions,jobs,n,matches,total,prepared);
            if(mode==7)searchKernel<7><<<unsigned(((uint64_t(n)<<13)+255)/256),256>>>(positions,jobs,n,matches,total,prepared);
        }
        return finish();
    }
    std::vector<Match> results(){
        uint64_t n=totalCount();if(n>MATCH_CAPACITY)throw std::runtime_error("match overflow");
        std::vector<Match> r(n);CUDA(cudaMemcpy(r.data(),matches,n*sizeof(Match),cudaMemcpyDeviceToHost));std::sort(r.begin(),r.end());return r;
    }
};

int main(int argc,char**argv){
 try{
    if(argc<3)throw std::runtime_error("usage: cuda_bench positions.txt normal [repeats=7]");
    std::ifstream in(argv[1]);if(!in)throw std::runtime_error("cannot open positions");
    int32_t x,z;std::vector<Position> positions;while(in>>x>>z)positions.push_back(makePosition(x,z));
    if(positions.empty()||positions.size()>65536)throw std::runtime_error("position count must be 1..65536");
    std::string task=argv[2];int repeats=argc>3?std::stoi(argv[3]):5;
    if(task!="normal"||repeats<1)throw std::runtime_error("expected normal and positive repeat count");
    auto processStart=Clock::now();Device d(positions);
    std::cout<<std::setprecision(10);
    cudaDeviceProp prop{};CUDA(cudaGetDeviceProperties(&prop,0));std::cout<<"{\"type\":\"device\",\"name\":\""<<prop.name<<"\",\"positions\":"<<positions.size()<<"}"<<std::endl;
    std::vector<uint32_t> referenceCounts;
    std::vector<Job> referenceJobs;
    auto collectRoots=[&](int mode){
        float ms=d.roots(positions.size(),mode,true);uint64_t total=d.totalCount();if(total>JOB_CAPACITY)throw std::runtime_error("job overflow");
        std::vector<uint32_t> counts(positions.size());CUDA(cudaMemcpy(counts.data(),d.counts,counts.size()*sizeof(uint32_t),cudaMemcpyDeviceToHost));
        std::vector<Job> jobs(total);CUDA(cudaMemcpy(jobs.data(),d.jobs,total*sizeof(Job),cudaMemcpyDeviceToHost));
        std::sort(jobs.begin(),jobs.end(),[](Job a,Job b){return std::tie(a.position,a.root)<std::tie(b.position,b.root);});
        if(referenceCounts.empty()){referenceCounts=counts;referenceJobs=jobs;}
        else {
            if(counts!=referenceCounts||jobs.size()!=referenceJobs.size())throw std::runtime_error("root mismatch");
            for(size_t i=0;i<jobs.size();++i)if(jobs[i].position!=referenceJobs[i].position||jobs[i].root!=referenceJobs[i].root)throw std::runtime_error("root list mismatch");
        }
        return ms;
    };
    {
        float rootMs=collectRoots(0);
        CUDA(cudaMemcpy(d.jobs,referenceJobs.data(),referenceJobs.size()*sizeof(Job),cudaMemcpyHostToDevice));
        std::vector<Match> expected;
        const std::vector<int> modes=std::vector<int>{3,4,5,6,7};
        bool first=true;
        for(int m:modes){d.search(referenceJobs.size(),m);auto r=d.results();if(first){expected=r;first=false;}else if(r!=expected)throw std::runtime_error("warmup result mismatch");}
        for(int r=0;r<repeats;++r)for(size_t k=0;k<modes.size();++k){
            int mode=modes[(k+r)%modes.size()];auto wallStart=Clock::now();float ms=d.search(referenceJobs.size(),mode);auto result=d.results();
            if(result!=expected)throw std::runtime_error("search result mismatch");
            std::cout<<"{\"type\":\"search_timing\",\"mode\":"<<mode<<",\"round\":"<<r<<",\"kernel_ms\":"<<ms<<",\"scan_ms\":"<<std::chrono::duration<double,std::milli>(Clock::now()-wallStart).count()<<",\"matches\":"<<result.size()<<"}"<<std::endl;
        }
        std::cout<<"{\"type\":\"search_info\",\"root_ms\":"<<rootMs<<",\"jobs\":"<<referenceJobs.size()<<",\"lower35\":"<<(uint64_t(referenceJobs.size())<<17)<<",\"matches\":[";
        for(size_t i=0;i<expected.size();++i){if(i)std::cout<<",";const auto& p=positions[expected[i].position];std::cout<<"["<<expected[i].seed<<","<<p.x<<","<<p.z<<"]";}std::cout<<"]}"<<std::endl;
    }
    // CPU checks of complete low18 root lists, including varied coordinates.
    for(size_t n=0;n<std::min<size_t>(64,positions.size());++n){
        size_t pos=n*positions.size()/std::min<size_t>(64,positions.size());std::vector<uint32_t> cpu,gpu;
        for(uint32_t s=0;s<(1U<<18);++s)if(sr::partialPatternPossible(s,18,positions[pos].cells,25))cpu.push_back(s);
        for(auto j:referenceJobs)if(j.position==pos)gpu.push_back(j.root);
        if(cpu!=gpu)throw std::runtime_error("CPU/GPU roots mismatch");
    }
    std::map<uint64_t,std::vector<uint32_t>> classes;
    size_t offset=0;
    for(size_t pos=0;pos<positions.size();++pos){
        uint64_t key=(uint64_t(uint32_t(positions[pos].x)&65535)<<17)|(uint32_t(positions[pos].z)&131071);
        std::vector<uint32_t> roots;
        while(offset<referenceJobs.size()&&referenceJobs[offset].position==pos)roots.push_back(referenceJobs[offset++].root);
        auto it=classes.find(key);
        if(it==classes.end())classes.emplace(key,std::move(roots));
        else if(it->second!=roots)throw std::runtime_error("periodic class root list mismatch");
    }
    std::cout<<"{\"type\":\"periodic_validation\",\"classes\":"<<classes.size()<<",\"positions\":"<<positions.size()<<",\"equal\":true}"<<std::endl;
    std::cout<<"{\"type\":\"root_counts\",\"counts\":[";for(size_t i=0;i<referenceCounts.size();++i){if(i)std::cout<<",";std::cout<<referenceCounts[i];}std::cout<<"]}"<<std::endl;
    std::cout<<"{\"type\":\"process\",\"ms\":"<<std::chrono::duration<double,std::milli>(Clock::now()-processStart).count()<<"}"<<std::endl;
 return 0;
 }catch(const std::exception&e){std::cerr<<e.what()<<"\n";return 1;}
}