// Teaching snapshot extracted from research/slime_position_residue_cuda.cu.
// Shared Java arithmetic/profile solvers remain in ../research; see README.md for scope.
#include "slime_solver_core.hpp"

#include <cuda_runtime.h>

#include <algorithm>
#include <chrono>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>

namespace sr = slime_research;

namespace {

constexpr int kResidueBits = 18;
constexpr uint32_t kPeriod = 1U << kResidueBits;
constexpr uint32_t kResidueMask = kPeriod - 1U;
constexpr int kMaxSize = 8;

struct ResiduePair {
    uint32_t x;
    uint32_t z;
};

struct Position {
    int32_t x;
    int32_t z;
};

struct Args {
    uint64_t seed = 61692151016558ULL;
    int32_t x0 = -1875000;
    int32_t x1 = 1874999;
    int32_t z0 = -1875000;
    int32_t z1 = 1874999;
    int size = 5;
    int threads = 256;
    uint32_t maxResidues = 1U << 20;
};

void cudaCheck(cudaError_t status, const char* what) {
    if (status != cudaSuccess) {
        throw std::runtime_error(std::string(what) + ": " + cudaGetErrorString(status));
    }
}

__global__ void residueKernel(
    uint64_t seed,
    const uint32_t* xTerms,
    const uint32_t* zTerms,
    int n,
    ResiduePair* output,
    unsigned long long* count,
    uint32_t capacity
) {
    constexpr uint64_t total = uint64_t(kPeriod) * kPeriod;
    const uint64_t stride = uint64_t(gridDim.x) * blockDim.x;
    for (uint64_t index = uint64_t(blockIdx.x) * blockDim.x + threadIdx.x;
         index < total;
         index += stride) {
        const uint32_t xr = static_cast<uint32_t>(index >> kResidueBits);
        const uint32_t zr = static_cast<uint32_t>(index) & kResidueMask;
        const uint32_t* xs = xTerms + uint64_t(xr) * n;
        const uint32_t* zs = zTerms + uint64_t(zr) * n;
        bool possible = true;
        for (int dx = 0; dx < n && possible; ++dx) {
            for (int dz = 0; dz < n; ++dz) {
                const uint64_t term = (xs[dx] + zs[dz]) & kResidueMask;
                // A multiple of ten is even, so bit 17 of the first next(31)
                // result must be zero.  This bit depends only on seed, x, and z
                // modulo 2^18.
                if (((sr::nextStateFromTerm(seed, term) >> 17) & 1ULL) != 0ULL) {
                    possible = false;
                    break;
                }
            }
        }
        if (possible) {
            const unsigned long long slot = atomicAdd(count, 1ULL);
            if (slot < capacity) output[slot] = ResiduePair{xr, zr};
        }
    }
}

Args parseArgs(int argc, char** argv) {
    Args args;
    for (int i = 1; i < argc; ++i) {
        const std::string key = argv[i];
        if (key == "--help") {
            std::cout
                << "Usage: slime_position_residue_cuda [options]\n"
                << "  --seed N --x0 N --x1 N --z0 N --z1 N\n"
                << "  -n N --threads N --max-residues N\n";
            std::exit(0);
        }
        if (i + 1 >= argc) throw std::runtime_error("missing value after " + key);
        const std::string value = argv[++i];
        if (key == "--seed") args.seed = std::stoull(value) & sr::kMask48;
        else if (key == "--x0") args.x0 = std::stoi(value);
        else if (key == "--x1") args.x1 = std::stoi(value);
        else if (key == "--z0") args.z0 = std::stoi(value);
        else if (key == "--z1") args.z1 = std::stoi(value);
        else if (key == "-n" || key == "--size") args.size = std::stoi(value);
        else if (key == "--threads") args.threads = std::stoi(value);
        else if (key == "--max-residues") args.maxResidues = std::stoul(value);
        else throw std::runtime_error("unknown option " + key);
    }
    if (args.x1 < args.x0 || args.z1 < args.z0) throw std::runtime_error("invalid bounds");
    if (args.size <= 0 || args.size > kMaxSize) throw std::runtime_error("size must be 1..8");
    if (args.threads <= 0 || args.threads > 1024) throw std::runtime_error("invalid threads");
    if (args.maxResidues == 0) throw std::runtime_error("max-residues must be positive");
    if (int64_t(args.x1) + args.size - 1 > INT32_MAX ||
        int64_t(args.z1) + args.size - 1 > INT32_MAX) {
        throw std::runtime_error("coordinate footprint exceeds int32");
    }
    return args;
}

std::vector<uint32_t> buildAxisTerms(bool xAxis, int n) {
    std::vector<uint32_t> result(uint64_t(kPeriod) * n);
    for (uint32_t residue = 0; residue < kPeriod; ++residue) {
        for (int offset = 0; offset < n; ++offset) {
            const int32_t coordinate = static_cast<int32_t>(residue + offset);
            const uint64_t term = xAxis ? sr::coordinateXTerm(coordinate)
                                        : sr::coordinateZTerm(coordinate);
            result[uint64_t(residue) * n + offset] = static_cast<uint32_t>(term) & kResidueMask;
        }
    }
    return result;
}

int64_t firstCongruent(int32_t lower, uint32_t residue) {
    int64_t lowerResidue = int64_t(lower) % kPeriod;
    if (lowerResidue < 0) lowerResidue += kPeriod;
    const int64_t delta = (int64_t(residue) - lowerResidue + kPeriod) % kPeriod;
    return int64_t(lower) + delta;
}

bool exactSolid(uint64_t seed, int32_t x, int32_t z, int n) {
    for (int dx = 0; dx < n; ++dx) {
        for (int dz = 0; dz < n; ++dz) {
            if (!sr::exactIsSlimeFromTerm(seed, sr::coordinateTerm(x + dx, z + dz))) {
                return false;
            }
        }
    }
    return true;
}

}  // namespace

int main(int argc, char** argv) {
    try {
        const Args args = parseArgs(argc, argv);
        const auto prepareBegin = std::chrono::steady_clock::now();
        const std::vector<uint32_t> hostX = buildAxisTerms(true, args.size);
        const std::vector<uint32_t> hostZ = buildAxisTerms(false, args.size);

        uint32_t* deviceX = nullptr;
        uint32_t* deviceZ = nullptr;
        ResiduePair* devicePairs = nullptr;
        unsigned long long* deviceCount = nullptr;
        cudaCheck(cudaMalloc(&deviceX, hostX.size() * sizeof(uint32_t)), "allocate X terms");
        cudaCheck(cudaMalloc(&deviceZ, hostZ.size() * sizeof(uint32_t)), "allocate Z terms");
        cudaCheck(cudaMalloc(&devicePairs, uint64_t(args.maxResidues) * sizeof(ResiduePair)),
                  "allocate residue output");
        cudaCheck(cudaMalloc(&deviceCount, sizeof(unsigned long long)), "allocate residue count");
        cudaCheck(cudaMemcpy(deviceX, hostX.data(), hostX.size() * sizeof(uint32_t),
                             cudaMemcpyHostToDevice), "copy X terms");
        cudaCheck(cudaMemcpy(deviceZ, hostZ.data(), hostZ.size() * sizeof(uint32_t),
                             cudaMemcpyHostToDevice), "copy Z terms");
        cudaCheck(cudaMemset(deviceCount, 0, sizeof(unsigned long long)), "reset residue count");
        const auto prepareEnd = std::chrono::steady_clock::now();

        cudaEvent_t begin, end;
        cudaCheck(cudaEventCreate(&begin), "create begin event");
        cudaCheck(cudaEventCreate(&end), "create end event");
        cudaCheck(cudaEventRecord(begin), "record begin");
        residueKernel<<<65535, args.threads>>>(args.seed, deviceX, deviceZ, args.size,
                                               devicePairs, deviceCount, args.maxResidues);
        cudaCheck(cudaGetLastError(), "launch residue kernel");
        cudaCheck(cudaEventRecord(end), "record end");
        cudaCheck(cudaEventSynchronize(end), "wait for residue kernel");
        float kernelMs = 0;
        cudaCheck(cudaEventElapsedTime(&kernelMs, begin, end), "measure residue kernel");

        unsigned long long residueCount = 0;
        cudaCheck(cudaMemcpy(&residueCount, deviceCount, sizeof(residueCount), cudaMemcpyDeviceToHost),
                  "copy residue count");
        if (residueCount > args.maxResidues) {
            throw std::runtime_error("residue output truncated; raise --max-residues (needed " +
                                     std::to_string(residueCount) + ")");
        }
        std::vector<ResiduePair> pairs(residueCount);
        if (!pairs.empty()) {
            cudaCheck(cudaMemcpy(pairs.data(), devicePairs, pairs.size() * sizeof(ResiduePair),
                                 cudaMemcpyDeviceToHost), "copy residue pairs");
        }
        cudaEventDestroy(begin);
        cudaEventDestroy(end);
        cudaFree(deviceCount);
        cudaFree(devicePairs);
        cudaFree(deviceZ);
        cudaFree(deviceX);

        const auto exactBegin = std::chrono::steady_clock::now();
        uint64_t expanded = 0;
        std::vector<Position> hits;
        for (const ResiduePair pair : pairs) {
            for (int64_t x = firstCongruent(args.x0, pair.x); x <= args.x1; x += kPeriod) {
                for (int64_t z = firstCongruent(args.z0, pair.z); z <= args.z1; z += kPeriod) {
                    ++expanded;
                    if (exactSolid(args.seed, static_cast<int32_t>(x), static_cast<int32_t>(z),
                                   args.size)) {
                        hits.push_back(Position{static_cast<int32_t>(x), static_cast<int32_t>(z)});
                    }
                }
            }
        }
        const auto exactEnd = std::chrono::steady_clock::now();
        const double prepareMs = std::chrono::duration<double, std::milli>(prepareEnd - prepareBegin).count();
        const double exactMs = std::chrono::duration<double, std::milli>(exactEnd - exactBegin).count();
        constexpr uint64_t residueSpace = uint64_t(kPeriod) * kPeriod;
        const uint64_t worldPositions = uint64_t(int64_t(args.x1) - args.x0 + 1) *
                                        uint64_t(int64_t(args.z1) - args.z0 + 1);
        std::cout << "residue_space=" << residueSpace
                  << " residue_survivors=" << residueCount
                  << " survival_rate=" << double(residueCount) / residueSpace << '\n'
                  << "prepare_ms=" << prepareMs << " residue_kernel_ms=" << kernelMs << '\n'
                  << "world_positions=" << worldPositions
                  << " expanded_candidates=" << expanded
                  << " exact_ms=" << exactMs << " hits=" << hits.size() << '\n';
        std::sort(hits.begin(), hits.end(), [](const Position& a, const Position& b) {
            return a.x < b.x || (a.x == b.x && a.z < b.z);
        });
        for (const Position hit : hits) std::cout << "hit " << hit.x << ' ' << hit.z << '\n';
        return 0;
    } catch (const std::exception& error) {
        std::cerr << "error: " << error.what() << '\n';
        return 1;
    }
}