// CPU emulator driver for the OpenCL pack: runs the generated kernel.cl (compiled as C++ through emu_opencl.h) on
// host threads with a configurable sub-group width, and checks cache, dataset and vectors against the pack.
//
// Model. A launch of global G work-items with work-group size L spawns L host threads; thread l plays work-item l
// of every work-group in turn (group 0, 1, 2, ...), exactly like proto-cuda/emu. barrier() synchronises the L
// threads of the current group. Sub-groups are consecutive runs of SG work-items inside the group (SG = --sg 32 or
// 64; the last run is clipped to the group size); sub_group_shuffle_xor exchanges through a slot array with a
// barrier among the sub-group's threads, so a 64-wide sub-group really carries two logical 32-lane units.
// Local memory is one arena per work-group instance, never shared between groups, as on hardware.
//
// Compile-time (as on the device): IGNEUM_GROUP (work-group size of igneum_hash) and IGNEUM_EXCHANGE (0 local memory,
// 1 sub-group shuffles). Runtime: --sg 32|64, --batch-log2 B (default 13), --dataset-mib N (default the pack size).
// Only PASS/FAIL matters here. Rates are meaningless and not printed.
#include "emu_opencl.h"
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <chrono>
#include <condition_variable>
#include <mutex>
#include <memory>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>

#define IGNEUM_NO_CUDA
#include "program.h"
#include "vectors.h"
#ifndef IGNEUM_DATASET_MODE
#define IGNEUM_DATASET_MODE 0
#endif
#if IGNEUM_DATASET_MODE == 1
#include "memhard.h"
#endif
#ifndef IGNEUM_GROUP
#define IGNEUM_GROUP 32
#endif
#ifndef IGNEUM_EXCHANGE
#define IGNEUM_EXCHANGE 0
#endif

// Kernels from kernel.cl (compiled as a separate C++ translation unit with the same defines).
void igneum_hash(const uint* ds, ulong* out, uint baseNonce, uint mask);
#if IGNEUM_DATASET_MODE == 1
void igneum_cache_fill(uint* cache, uint nSegments);
void igneum_build(uint* ds, const uint* cache, uint nItems);
#else
void igneum_fill(uint* ds, uint n, uint d0, uint d1);
#endif

// ---------------------------------------------------------------------------------------------
// Runtime

struct Barrier {
    std::mutex m;
    std::condition_variable cv;
    unsigned size = 0, arrived = 0, generation = 0;
    void wait() {
        std::unique_lock<std::mutex> lk(m);
        unsigned gen = generation;
        if (++arrived == size) { arrived = 0; ++generation; cv.notify_all(); }
        else cv.wait(lk, [&] { return gen != generation; });
    }
};

struct SubGroup {
    Barrier bar;
    uint slot[64];
    unsigned size = 0;
};

struct Launch {
    unsigned local = 0, groups = 0, sg = 32;
    Barrier groupBar;
    std::vector<std::unique_ptr<SubGroup>> subs;
    std::mutex arenaMutex;
    std::unordered_map<unsigned, std::unique_ptr<uint[]>> arenas;   // group index -> local memory words
    unsigned arenaWords = 0;
};

static thread_local Launch* tlLaunch = nullptr;
static thread_local unsigned tlLid = 0, tlGroup = 0;
static unsigned gSubGroupWidth = 32;

size_t get_global_id(uint) { return (size_t)tlGroup * tlLaunch->local + tlLid; }
size_t get_local_id(uint) { return tlLid; }
size_t get_group_id(uint) { return tlGroup; }
size_t get_local_size(uint) { return tlLaunch->local; }
size_t get_global_size(uint) { return (size_t)tlLaunch->local * tlLaunch->groups; }
uint get_sub_group_size(void) { return tlLaunch->subs[tlLid / tlLaunch->sg]->size; }
uint get_sub_group_local_id(void) { return tlLid % tlLaunch->sg; }
uint get_sub_group_id(void) { return tlLid / tlLaunch->sg; }
uint get_num_sub_groups(void) { return (uint)tlLaunch->subs.size(); }

void barrier(int) { tlLaunch->groupBar.wait(); }

uint sub_group_shuffle_xor(uint v, uint mask) {
    SubGroup* s = tlLaunch->subs[tlLid / tlLaunch->sg].get();
    unsigned lane = tlLid % tlLaunch->sg;
    unsigned partner = lane ^ mask;
    if (partner >= s->size) {
        std::fprintf(stderr, "emu: sub_group_shuffle_xor partner lane %u outside the sub-group of %u lanes (undefined on hardware)\n", partner, s->size);
        std::exit(3);
    }
    s->slot[lane] = v;
    s->bar.wait();
    uint r = s->slot[partner];
    s->bar.wait();
    return r;
}
uint intel_sub_group_shuffle_xor(uint v, uint mask) { return sub_group_shuffle_xor(v, mask); }

uint sub_group_broadcast(uint v, uint laneSrc) {
    SubGroup* s = tlLaunch->subs[tlLid / tlLaunch->sg].get();
    unsigned lane = tlLid % tlLaunch->sg;
    s->slot[lane] = v;
    s->bar.wait();
    uint r = s->slot[laneSrc % s->size];
    s->bar.wait();
    return r;
}

uint* emu_local_words(uint n) {
    Launch* L = tlLaunch;
    std::lock_guard<std::mutex> lk(L->arenaMutex);
    auto it = L->arenas.find(tlGroup);
    if (it == L->arenas.end()) {
        std::unique_ptr<uint[]> a(new uint[n]());
        it = L->arenas.emplace(tlGroup, std::move(a)).first;
        L->arenaWords = n;
    }
    return it->second.get();
}

template<class F, class... A> static void emu_launch(F f, size_t global, unsigned local, A... args) {
    if (global % local != 0) { std::fprintf(stderr, "emu: global %zu not a multiple of local %u\n", global, local); std::exit(3); }
    Launch L;
    L.local = local; L.groups = (unsigned)(global / local); L.sg = gSubGroupWidth;
    L.groupBar.size = local;
    unsigned nSub = (local + L.sg - 1) / L.sg;
    for (unsigned k = 0; k < nSub; ++k) {
        L.subs.emplace_back(new SubGroup());
        L.subs.back()->size = (local - k * L.sg) < L.sg ? (local - k * L.sg) : L.sg;
        L.subs.back()->bar.size = L.subs.back()->size;
    }
    std::vector<std::thread> ts;
    for (unsigned t = 0; t < local; ++t) {
        ts.emplace_back([&, t]() {
            tlLaunch = &L; tlLid = t;
            for (unsigned g = 0; g < L.groups; ++g) { tlGroup = g; f(args...); }
        });
    }
    for (auto& th : ts) th.join();
}

// ---------------------------------------------------------------------------------------------
// Checks

static uint64_t fnv1a64(const void* p, size_t n) {
    const uint8_t* b = (const uint8_t*)p;
    uint64_t h = 0xcbf29ce484222325ull;
    for (size_t i = 0; i < n; ++i) { h ^= b[i]; h *= 0x100000001b3ull; }
    return h;
}

#if IGNEUM_DATASET_MODE == 0
static uint32_t host_ds_elem(uint32_t i, uint32_t d0, uint32_t d1) {
    uint32_t x = i ^ d0;
    x *= 0x9E3779B1u; x ^= x >> 15;
    x += d1;
    x *= 0x85EBCA77u; x ^= x >> 13;
    x *= 0xC2B2AE3Du; x ^= x >> 16;
    return x;
}
#endif

static bool compareWarp(const uint64_t* got, const uint64_t* want, uint32_t base, const char* how) {
    int bad = 0, first = -1;
    for (int l = 0; l < 32; ++l) if (got[l] != want[l]) { if (first < 0) first = l; ++bad; }
    if (bad == 0) std::printf("verify warp base %u %s: PASS\n", base, how);
    else std::printf("verify warp base %u %s: FAIL  %d of 32 lanes differ, first lane %d: emu=%016llx expected=%016llx\n",
                     base, how, bad, first, (unsigned long long)got[first], (unsigned long long)want[first]);
    return bad == 0;
}

static double wallMs() {
    using namespace std::chrono;
    return duration<double, std::milli>(steady_clock::now().time_since_epoch()).count();
}

int main(int argc, char** argv) {
    int batchLog2 = 13;
    int datasetMib = (int)(((1ull << IGNEUM_DATASET_LOG2) * 4ull) >> 20);
    for (int i = 1; i < argc; ++i) {
        std::string a = argv[i];
        if (a == "--sg" && i + 1 < argc) gSubGroupWidth = (unsigned)std::atoi(argv[++i]);
        else if (a == "--batch-log2" && i + 1 < argc) batchLog2 = std::atoi(argv[++i]);
        else if (a == "--dataset-mib" && i + 1 < argc) datasetMib = std::atoi(argv[++i]);
        else { std::printf("usage: igneum-emu-cl [--sg 32|64] [--batch-log2 13] [--dataset-mib N]\n"); return 2; }
    }
    if (gSubGroupWidth != 32 && gSubGroupWidth != 64) { std::printf("--sg must be 32 or 64\n"); return 2; }
    if (IGNEUM_EXCHANGE != 0 && IGNEUM_GROUP != 32 && IGNEUM_GROUP != 64) { std::printf("sub-group exchange needs IGNEUM_GROUP 32 or 64 here\n"); return 2; }

    const uint32_t words = (uint32_t)(((uint64_t)datasetMib << 20) / 4ull);
    const uint32_t mask = words - 1u;
    const bool atPackSize = (words == (1u << IGNEUM_DATASET_LOG2));
    std::printf("igneum-emu-cl  pack \"%s\"  CPU EMULATION of kernel.cl (not a GPU; PASS/FAIL only)\n", IGNEUM_SEED_STRING);
    std::printf("configuration: IGNEUM_GROUP %d, IGNEUM_EXCHANGE %d (%s), emulated sub-group width %u%s, dataset %d MiB, batch 2^%d\n",
                IGNEUM_GROUP, IGNEUM_EXCHANGE, IGNEUM_EXCHANGE == 0 ? "local-memory exchange with barrier" : "sub_group_shuffle_xor",
                gSubGroupWidth, gSubGroupWidth == 64 ? " (two logical 32-lane units per wave)" : "", datasetMib, batchLog2);
    std::printf("hardware threads: %u\n", std::thread::hardware_concurrency());

    bool overall = true;
    std::vector<uint> ds(words);
#if IGNEUM_DATASET_MODE == 1
    const uint32_t cacheWords = 1u << IGNEUM_CACHE_LOG2_WORDS;
    std::vector<uint> cache(cacheWords), hostCache(cacheWords);
    double t0 = wallMs();
    emu_launch(igneum_cache_fill, (size_t)IGNEUM_CACHE_SEGMENTS, 256u, cache.data(), (uint)IGNEUM_CACHE_SEGMENTS);
    double t1 = wallMs();
    for (uint32_t seg = 0; seg < IGNEUM_CACHE_SEGMENTS; ++seg) mh_cache_segment(hostCache.data(), seg);
    double t2 = wallMs();
    bool same = std::memcmp(cache.data(), hostCache.data(), (size_t)cacheWords * 4u) == 0;
    uint64_t fnv = fnv1a64(cache.data(), (size_t)cacheWords * 4u);
    bool fnvOk = fnv == IGNEUM_CACHE_FNV64;
    bool headOk = std::memcmp(cache.data(), IGNEUM_CACHE_HEAD, 64) == 0;
    bool lastOk = std::memcmp(cache.data() + cacheWords - 16u, IGNEUM_CACHE_LAST, 64) == 0;
    std::printf("cache: emulated igneum_cache_fill %.0f ms (256 threads), host memhard.h one thread %.0f ms\n", t1 - t0, t2 - t1);
    std::printf("cache check: %s  (emulated kernel == host all %u words %s, FNV-1a 64 %016llx vs Mac %016llx %s, head %s, last line %s)\n",
                (same && fnvOk && headOk && lastOk) ? "PASS" : "FAIL", cacheWords, same ? "PASS" : "FAIL",
                (unsigned long long)fnv, (unsigned long long)IGNEUM_CACHE_FNV64, fnvOk ? "PASS" : "FAIL", headOk ? "PASS" : "FAIL", lastOk ? "PASS" : "FAIL");
    overall = overall && same && fnvOk && headOk && lastOk;
    t0 = wallMs();
    emu_launch(igneum_build, (size_t)(words / 16u), 256u, ds.data(), (const uint*)cache.data(), (uint)(words / 16u));
    std::printf("dataset: emulated igneum_build %.0f ms for 2^%u items\n", wallMs() - t0, (unsigned)(IGNEUM_DATASET_LOG2 - 4));
#else
    double t0 = wallMs();
    emu_launch(igneum_fill, (size_t)words, 256u, ds.data(), words, (uint)IGNEUM_DAY0, (uint)IGNEUM_DAY1);
    std::printf("dataset: emulated igneum_fill %.0f ms\n", wallMs() - t0);
#endif

    // Dataset self-test, same shape as host.c.
    {
        int badHead = 0, badRnd = 0, badSample = 0, nSample = 0;
        bool lastOk = true;
        for (int i = 0; i < 16; ++i) if (ds[i] != IGNEUM_DS_HEAD[i]) ++badHead;
        if (atPackSize) lastOk = (ds[IGNEUM_DS_LAST_INDEX] == IGNEUM_DS_LAST);
        uint64_t s = 0x9E3779B97F4A7C15ull ^ (uint64_t)words;
        for (int k = 0; k < 64; ++k) {
            s += 0x9E3779B97F4A7C15ull;
            uint64_t z = s;
            z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ull;
            z = (z ^ (z >> 27)) * 0x94D049BB133111EBull;
            z ^= z >> 31;
            uint32_t idx = (uint32_t)z & mask;
#if IGNEUM_DATASET_MODE == 1
            uint32_t want = mh_word(hostCache.data(), idx);
#else
            uint32_t want = host_ds_elem(idx, IGNEUM_DAY0, IGNEUM_DAY1);
#endif
            if (ds[idx] != want) ++badRnd;
        }
#ifdef IGNEUM_DS_SAMPLES
        for (int k = 0; k < IGNEUM_DS_SAMPLES; ++k) {
            if (IGNEUM_DS_SAMPLE_INDEX[k] > mask) continue;
            ++nSample;
            if (ds[IGNEUM_DS_SAMPLE_INDEX[k]] != IGNEUM_DS_SAMPLE_VALUE[k]) ++badSample;
        }
#endif
        bool dsPass = badHead == 0 && lastOk && badRnd == 0 && badSample == 0;
        std::printf("dataset self-test: %s  (head 16 vs Mac %s, element [MASK] %s, 64 random points vs host %s, %d Mac samples %s)\n",
                    dsPass ? "PASS" : "FAIL", badHead == 0 ? "PASS" : "FAIL", atPackSize ? (lastOk ? "PASS" : "FAIL") : "skipped",
                    badRnd == 0 ? "PASS" : "FAIL", nSample, badSample == 0 ? "PASS" : "FAIL");
        overall = overall && dsPass;
    }

    if (!atPackSize) { std::printf("vectors: skipped (not the pack size)\nOVERALL: %s\n", overall ? "PASS" : "FAIL"); return overall ? 0 : 1; }

    // Vectors standalone: one work-group of IGNEUM_GROUP items per base nonce (the first 32 are the vector warp).
    const uint32_t nonces = 1u << batchLog2;
    std::vector<ulong> out(nonces);
    for (int w = 0; w < IGNEUM_VEC_WARPS; ++w) {
        emu_launch(igneum_hash, (size_t)IGNEUM_GROUP, (unsigned)IGNEUM_GROUP, (const uint*)ds.data(), out.data(), (uint)IGNEUM_VEC_BASE[w], mask);
        char how[96];
        std::snprintf(how, sizeof(how), "standalone, work-group %d, sub-group width %u", IGNEUM_GROUP, gSubGroupWidth);
        overall = compareWarp((const uint64_t*)out.data(), IGNEUM_VEC_OUT[w], IGNEUM_VEC_BASE[w], how) && overall;
    }
    // In batch: every vector warp that fits in 2^batchLog2 nonces.
    emu_launch(igneum_hash, (size_t)nonces, (unsigned)IGNEUM_GROUP, (const uint*)ds.data(), out.data(), 0u, mask);
    for (int w = 0; w < IGNEUM_VEC_WARPS; ++w) {
        if ((uint64_t)IGNEUM_VEC_BASE[w] + 32ull > nonces) { std::printf("verify warp base %u in batch: skipped (batch has %u nonces)\n", IGNEUM_VEC_BASE[w], nonces); continue; }
        char how[96];
        std::snprintf(how, sizeof(how), "in batch of 2^%d, work-group %d, sub-group width %u", batchLog2, IGNEUM_GROUP, gSubGroupWidth);
        overall = compareWarp((const uint64_t*)out.data() + IGNEUM_VEC_BASE[w], IGNEUM_VEC_OUT[w], IGNEUM_VEC_BASE[w], how) && overall;
    }
    // Fingerprint of the whole batch so different configurations can be compared bit for bit.
    std::printf("batch fingerprint (FNV-1a 64 of 2^%d outputs): %016llx\n", batchLog2, (unsigned long long)fnv1a64(out.data(), (size_t)nonces * 8u));
    std::printf("OVERALL: %s\n", overall ? "PASS" : "FAIL");
    return overall ? 0 : 1;
}