#include "interval_solver.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};
    }
};
template<int Mode>
__global__ void searchKernel(const Position* positions,const Job* jobs,uint32_t count,
                             Match* out,unsigned long long* total) {
    uint64_t idx=uint64_t(blockIdx.x)*blockDim.x+threadIdx.x;
    if(idx >= (uint64_t(count)<<17))return;
    const Job job=jobs[idx>>17];
    uint64_t low=job.root+((idx&((1U<<17)-1))<<18);
    const Position& p=positions[job.position];
    Emitter emit{job.position,p.anchor,out,total};
    if constexpr(Mode==0)sr::solveUpper13(low,p.cells,25,emit);
    else if constexpr(Mode==1)fs::solveIntervals(low,p.cells,25,emit);
    else if constexpr(Mode==2){if(fs::differenceFilter(low,p.cells,25))sr::solveUpper13(low,p.cells,25,emit);}
    else if constexpr(Mode==3){if(fs::differenceFilter(low,p.cells,25))fs::solveIntervals(low,p.cells,25,emit);}
}
__global__ void tailKernel(const Position* positions,uint32_t count,Match* out,unsigned long long* total,int cellCount){
    const uint64_t idx=uint64_t(blockIdx.x)*blockDim.x+threadIdx.x;
    if(idx >= uint64_t(count)*cellCount*(1U<<20))return;
    const uint64_t cell=idx>>20;
    const uint32_t pos=cell/cellCount,i=cell%cellCount;
    const uint64_t v=(uint64_t(2147483640U)<<17)+(idx&((1U<<20)-1));
    const Position& p=positions[pos];
    const uint64_t u=((v-11)*0xDFE05BCB1365ULL)&sr::kMask48;
    const uint64_t seed=((u^sr::kCombinedXor)-p.cells[i].term)&sr::kMask48;
    if(sr::verifyPattern(seed,p.cells,cellCount)){
        Emitter e{pos,p.anchor,out,total};e(seed);
    }
}

// If an output is odd, a true solid-pattern seed must reject that output.
// Pick the first odd cell: four tail states correspond to each seed-low18.
// When every output is even, enumerate all cells' four tail states instead.
// This covers multiple-rejection paths too, without assuming they are impossible.
__global__ void canonicalTailKernel(const Position* positions,uint32_t count,Match* out,unsigned long long* total,int cellCount){
    const uint64_t idx=uint64_t(blockIdx.x)*blockDim.x+threadIdx.x;
    if(idx >= (uint64_t(count)<<18))return;
    const uint32_t pos=idx>>18,root=idx&((1U<<18)-1);
    const Position& p=positions[pos];
    int firstOdd=-1;
    for(int i=0;i<cellCount;++i)if((sr::nextStateFromTerm(root,p.cells[i].term)>>17)&1U){firstOdd=i;break;}
    const int begin=firstOdd>=0?firstOdd:0,end=firstOdd>=0?firstOdd+1:cellCount;
    for(int i=begin;i<end;++i){
        const uint32_t vlow=uint32_t(sr::nextStateFromTerm(root,p.cells[i].term))&((1U<<18)-1);
        for(uint32_t k=0;k<4;++k){
            const uint64_t v=(uint64_t(2147483640U+(vlow>>17)+2*k)<<17)+(vlow&((1U<<17)-1));
            const uint64_t u=((v-11)*0xDFE05BCB1365ULL)&sr::kMask48;
            const uint64_t seed=((u^sr::kCombinedXor)-p.cells[i].term)&sr::kMask48;
            if(sr::verifyPattern(seed,p.cells,cellCount)){Emitter e{pos,p.anchor,out,total};e(seed);}
        }
    }
}

struct Device {
    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(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){
        CUDA(cudaMemset(total,0,sizeof(*total)));CUDA(cudaEventRecord(start));
        unsigned blocks=unsigned(((uint64_t(n)<<17)+255)/256);
        if(blocks){
            if(mode==0)searchKernel<0><<<blocks,256>>>(positions,jobs,n,matches,total);
            if(mode==1)searchKernel<1><<<blocks,256>>>(positions,jobs,n,matches,total);
            if(mode==2)searchKernel<2><<<blocks,256>>>(positions,jobs,n,matches,total);
            if(mode==3)searchKernel<3><<<blocks,256>>>(positions,jobs,n,matches,total);
        }
        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 root|compare|fast|tail|complete [repeats=5] [tail-cells=25]");
    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;
    int tailCells=argc>4?std::stoi(argv[4]):25;
    if(tailCells<1||tailCells>25)throw std::runtime_error("tail-cells must be 1..25");
    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;
    };
    if(task=="root"){
        collectRoots(0);collectRoots(1);
        for(int r=0;r<repeats;++r)for(int k=0;k<2;++k){int mode=(r+k)%2;float ms=collectRoots(mode);std::cout<<"{\"type\":\"root_timing\",\"mode\":"<<mode<<",\"round\":"<<r<<",\"kernel_ms\":"<<ms<<"}"<<std::endl;}
    }else if(task=="tail"){
        std::vector<Match> expected;
        for(int r=-1;r<repeats;++r)for(int k=0;k<2;++k){
            const int mode=r<0?k:(r+k)%2;
            CUDA(cudaMemset(d.total,0,sizeof(*d.total)));CUDA(cudaEventRecord(d.start));
            const uint64_t work=mode==0?uint64_t(positions.size())*tailCells*(1U<<20):(uint64_t(positions.size())<<18);
            if((work+255)/256>0x7fffffffULL)throw std::runtime_error("tail grid too large");
            unsigned blocks=unsigned((work+255)/256);
            if(mode==0)tailKernel<<<blocks,256>>>(d.positions,positions.size(),d.matches,d.total,tailCells);
            else canonicalTailKernel<<<blocks,256>>>(d.positions,positions.size(),d.matches,d.total,tailCells);
            float ms=d.finish();auto result=d.results();
            result.erase(std::unique(result.begin(),result.end()),result.end());
            if(r==-1&&mode==0)expected=result;else if(result!=expected)throw std::runtime_error("tail mismatch");
            if(r>=0)std::cout<<"{\"type\":\"tail_timing\",\"mode\":"<<mode<<",\"cells\":"<<tailCells<<",\"round\":"<<r<<",\"kernel_ms\":"<<ms<<",\"matches\":"<<result.size()<<"}"<<std::endl;
        }
        std::cout<<"{\"type\":\"process\",\"ms\":"<<std::chrono::duration<double,std::milli>(Clock::now()-processStart).count()<<"}"<<std::endl;return 0;
    }else if(task=="complete"){
        const float rootMs=collectRoots(0);
        CUDA(cudaMemcpy(d.jobs,referenceJobs.data(),referenceJobs.size()*sizeof(Job),cudaMemcpyHostToDevice));
        std::vector<Match> expected;
        for(int r=-1;r<repeats;++r){
            auto wallStart=Clock::now();
            const float firstMs=d.search(referenceJobs.size(),3);
            auto result=d.results();
            CUDA(cudaMemset(d.total,0,sizeof(*d.total)));CUDA(cudaEventRecord(d.start));
            const unsigned blocks=unsigned(((uint64_t(positions.size())<<18)+255)/256);
            canonicalTailKernel<<<blocks,256>>>(d.positions,positions.size(),d.matches,d.total,25);
            const float tailMs=d.finish();auto tail=d.results();
            result.insert(result.end(),tail.begin(),tail.end());std::sort(result.begin(),result.end());
            result.erase(std::unique(result.begin(),result.end()),result.end());
            if(r==-1)expected=result;else if(result!=expected)throw std::runtime_error("complete pipeline result mismatch");
            if(r>=0)std::cout<<"{\"type\":\"complete_timing\",\"round\":"<<r<<",\"first_ms\":"<<firstMs
                            <<",\"tail_ms\":"<<tailMs<<",\"compute_ms\":"<<(firstMs+tailMs)
                            <<",\"scan_ms\":"<<std::chrono::duration<double,std::milli>(Clock::now()-wallStart).count()
                            <<",\"matches\":"<<result.size()<<"}"<<std::endl;
        }
        std::cout<<"{\"type\":\"complete_info\",\"root_ms\":"<<rootMs<<",\"jobs\":"<<referenceJobs.size()<<",\"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;
    }else{
        float rootMs=collectRoots(0);
        CUDA(cudaMemcpy(d.jobs,referenceJobs.data(),referenceJobs.size()*sizeof(Job),cudaMemcpyHostToDevice));
        std::vector<Match> expected;
        const std::vector<int> modes=task=="compare"?std::vector<int>{0,1,2,3}:std::vector<int>{1,2,3};
        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;
 }catch(const std::exception&e){std::cerr<<e.what()<<"\n";return 1;}
}