#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 kMaxSize = 8;

struct Hit {
    int32_t x;
    int32_t z;
};

bool operator==(const Hit& a, const Hit& b) {
    return a.x == b.x && a.z == b.z;
}

struct Args {
    uint64_t seed = 61692151016558ULL;
    int32_t x0 = 46973;
    int32_t x1 = 55165;
    int32_t z0 = 415028;
    int32_t z1 = 423220;
    int size = 5;
    int threads = 256;
    int repeats = 3;
    uint32_t maxHits = 1024;
};

__device__ bool solidDirect(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;
}

__device__ void emitHit(Hit* hits, uint32_t* count, uint32_t capacity, int32_t x, int32_t z) {
    const uint32_t index = atomicAdd(count, 1U);
    if (index < capacity) {
        hits[index] = Hit{x, z};
    }
}

__global__ void directKernel(
    uint64_t seed,
    int32_t x0,
    int32_t z0,
    uint64_t width,
    uint64_t height,
    int n,
    Hit* hits,
    uint32_t* count,
    uint32_t capacity
) {
    const uint64_t stride = uint64_t(gridDim.x) * blockDim.x;
    for (uint64_t index = uint64_t(blockIdx.x) * blockDim.x + threadIdx.x;
         index < width * height;
         index += stride) {
        const int32_t x = x0 + static_cast<int32_t>(index / height);
        const int32_t z = z0 + static_cast<int32_t>(index % height);
        if (solidDirect(seed, x, z, n)) {
            emitHit(hits, count, capacity, x, z);
        }
    }
}

__global__ void outerKernel(
    uint64_t seed,
    int32_t x0,
    int32_t z0,
    uint64_t width,
    uint64_t height,
    int n,
    Hit* hits,
    uint32_t* count,
    uint32_t capacity
) {
    // One block covers a 16x16 position tile.  The expanded one-dimensional
    // axis terms are shared by all 256 candidate positions in the tile.
    __shared__ uint64_t xTerms[16 + kMaxSize - 1];
    __shared__ uint64_t zTerms[16 + kMaxSize - 1];
    const int expanded = 16 + n - 1;
    const int linearThread = threadIdx.y * blockDim.x + threadIdx.x;
    if (linearThread < expanded) {
        xTerms[linearThread] = sr::coordinateXTerm(
            x0 + static_cast<int32_t>(uint64_t(blockIdx.x) * 16 + linearThread));
        zTerms[linearThread] = sr::coordinateZTerm(
            z0 + static_cast<int32_t>(uint64_t(blockIdx.y) * 16 + linearThread));
    }
    __syncthreads();

    const uint64_t xi = uint64_t(blockIdx.x) * 16 + threadIdx.x;
    const uint64_t zi = uint64_t(blockIdx.y) * 16 + threadIdx.y;
    if (xi >= width || zi >= height) {
        return;
    }
    for (int dx = 0; dx < n; ++dx) {
        for (int dz = 0; dz < n; ++dz) {
            if (!sr::exactIsSlimeFromTerm(seed,
                                          xTerms[threadIdx.x + dx] +
                                          zTerms[threadIdx.y + dz])) {
                return;
            }
        }
    }
    emitHit(hits, count, capacity, x0 + static_cast<int32_t>(xi),
            z0 + static_cast<int32_t>(zi));
}

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

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_cuda [options]\n"
                << "  --seed N --x0 N --x1 N --z0 N --z1 N\n"
                << "  -n N --threads N --repeats N --max-hits N\n"
                << "Bounds are inclusive. Both kernels run and their hits must agree.\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 == "--repeats") args.repeats = std::stoi(value);
        else if (key == "--max-hits") args.maxHits = 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.repeats <= 0 || args.maxHits == 0) throw std::runtime_error("invalid repeat/hit count");
    return args;
}

enum class Mode { Direct, Outer };

struct RunResult {
    float milliseconds;
    uint32_t totalHits;
    std::vector<Hit> storedHits;
};

RunResult run(const Args& args, Mode mode, Hit* deviceHits, uint32_t* deviceCount) {
    const uint64_t width = uint64_t(int64_t(args.x1) - args.x0 + 1);
    const uint64_t height = uint64_t(int64_t(args.z1) - args.z0 + 1);
    cudaCheck(cudaMemset(deviceCount, 0, sizeof(uint32_t)), "reset count");
    cudaEvent_t begin, end;
    cudaCheck(cudaEventCreate(&begin), "create begin event");
    cudaCheck(cudaEventCreate(&end), "create end event");
    cudaCheck(cudaEventRecord(begin), "record begin");
    if (mode == Mode::Direct) {
        const uint64_t work = width * height;
        const int blocks = static_cast<int>(std::min<uint64_t>((work + args.threads - 1) / args.threads, 65535));
        directKernel<<<blocks, args.threads>>>(args.seed, args.x0, args.z0, width, height,
                                              args.size, deviceHits, deviceCount, args.maxHits);
    } else {
        const dim3 threads(16, 16);
        const dim3 blocks(static_cast<unsigned>((width + 15) / 16),
                          static_cast<unsigned>((height + 15) / 16));
        outerKernel<<<blocks, threads>>>(args.seed, args.x0, args.z0, width, height,
                                         args.size, deviceHits, deviceCount, args.maxHits);
    }
    cudaCheck(cudaGetLastError(), "launch kernel");
    cudaCheck(cudaEventRecord(end), "record end");
    cudaCheck(cudaEventSynchronize(end), "wait for kernel");
    float milliseconds = 0;
    cudaCheck(cudaEventElapsedTime(&milliseconds, begin, end), "measure kernel");
    cudaEventDestroy(begin);
    cudaEventDestroy(end);

    uint32_t total = 0;
    cudaCheck(cudaMemcpy(&total, deviceCount, sizeof(total), cudaMemcpyDeviceToHost), "copy count");
    std::vector<Hit> hits(std::min(total, args.maxHits));
    if (!hits.empty()) {
        cudaCheck(cudaMemcpy(hits.data(), deviceHits, hits.size() * sizeof(Hit), cudaMemcpyDeviceToHost),
                  "copy hits");
    }
    std::sort(hits.begin(), hits.end(), [](const Hit& a, const Hit& b) {
        return a.x < b.x || (a.x == b.x && a.z < b.z);
    });
    return RunResult{milliseconds, total, std::move(hits)};
}

}  // namespace

int main(int argc, char** argv) {
    try {
        const Args args = parseArgs(argc, argv);
        Hit* deviceHits = nullptr;
        uint32_t* deviceCount = nullptr;
        cudaCheck(cudaMalloc(&deviceHits, args.maxHits * sizeof(Hit)), "allocate hits");
        cudaCheck(cudaMalloc(&deviceCount, sizeof(uint32_t)), "allocate count");

        // Warm up the context before timing.
        cudaCheck(cudaFree(nullptr), "initialize CUDA");
        std::vector<float> directTimes, outerTimes;
        RunResult direct{}, outer{};
        for (int repeat = 0; repeat < args.repeats; ++repeat) {
            direct = run(args, Mode::Direct, deviceHits, deviceCount);
            outer = run(args, Mode::Outer, deviceHits, deviceCount);
            if (direct.totalHits != outer.totalHits || direct.storedHits != outer.storedHits) {
                throw std::runtime_error("direct and outer-sum results disagree");
            }
            directTimes.push_back(direct.milliseconds);
            outerTimes.push_back(outer.milliseconds);
        }
        cudaFree(deviceCount);
        cudaFree(deviceHits);
        std::sort(directTimes.begin(), directTimes.end());
        std::sort(outerTimes.begin(), outerTimes.end());
        const float directMs = directTimes[directTimes.size() / 2];
        const float outerMs = outerTimes[outerTimes.size() / 2];
        const uint64_t positions = uint64_t(int64_t(args.x1) - args.x0 + 1) *
                                   uint64_t(int64_t(args.z1) - args.z0 + 1);
        std::cout << "positions=" << positions << " size=" << args.size
                  << " hits=" << direct.totalHits << "\n"
                  << "direct_kernel_ms=" << directMs << " positions_per_second="
                  << positions / (directMs * 1e-3) << "\n"
                  << "outer_kernel_ms=" << outerMs << " positions_per_second="
                  << positions / (outerMs * 1e-3) << "\n"
                  << "outer_speedup=" << directMs / outerMs << "\n";
        for (const Hit& hit : direct.storedHits) {
            std::cout << "hit " << hit.x << ' ' << hit.z << "\n";
        }
        if (direct.totalHits > direct.storedHits.size()) {
            std::cout << "truncated_hits=" << direct.totalHits - direct.storedHits.size() << "\n";
        }
        return 0;
    } catch (const std::exception& error) {
        std::cerr << "error: " << error.what() << '\n';
        return 1;
    }
}