#include <CL/cl.h>

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <iomanip>
#include <iostream>
#include <random>
#include <string>
#include <vector>

static constexpr uint32_t DIM = 128;
static constexpr uint32_t TOPK = 32;
static constexpr uint32_t QUERIES = 256;

// Numero massimo di vettori elaborati da un singolo kernel launch.
static constexpr size_t CHUNK_SIZE = 50000;

static void check_cl(cl_int err, const char *what) {
    if (err != CL_SUCCESS) {
        std::cerr << "OpenCL error " << err << " at " << what << "\n";
        std::exit(1);
    }
}

static void normalize(float *v) {
    double s = 0.0;

    for (uint32_t i = 0; i < DIM; ++i)
        s += double(v[i]) * double(v[i]);

    float inv = 1.0f / std::sqrt(float(s));

    for (uint32_t i = 0; i < DIM; ++i)
        v[i] *= inv;
}

static std::vector<float> make_embeddings(size_t n) {
    std::vector<float> x(n * DIM);

    std::mt19937 rng(123456789);
    std::normal_distribution<float> d(0.0f, 1.0f);

    for (size_t i = 0; i < x.size(); ++i)
        x[i] = d(rng);

    for (size_t i = 0; i < n; ++i)
        normalize(&x[i * DIM]);

    return x;
}

struct Neighbor {
    float score;
    uint32_t index;
};

static void cpu_knn_one(
    const std::vector<float>& db,
    const float *query,
    size_t n,
    Neighbor *out)
{
    std::vector<Neighbor> all(n);

    for (size_t i = 0; i < n; ++i) {
        const float *v = &db[i * DIM];

        float s = 0.0f;

        for (uint32_t j = 0; j < DIM; ++j)
            s += v[j] * query[j];

        all[i] = {s, static_cast<uint32_t>(i)};
    }

    std::nth_element(
        all.begin(),
        all.begin() + TOPK,
        all.end(),
        [](const Neighbor& a, const Neighbor& b) {
            return a.score > b.score;
        });

    std::sort(
        all.begin(),
        all.begin() + TOPK,
        [](const Neighbor& a, const Neighbor& b) {
            return a.score > b.score;
        });

    for (uint32_t k = 0; k < TOPK; ++k)
        out[k] = all[k];
}

static double benchmark_cpu(
    const std::vector<float>& db,
    const std::vector<float>& queries,
    std::vector<Neighbor>& result,
    size_t n)
{
    result.resize(QUERIES * TOPK);


    auto t0 = std::chrono::steady_clock::now();

    for (uint32_t q = 0; q < QUERIES; ++q) {
        cpu_knn_one(
            db,
            &queries[q * DIM],
            n,
            &result[q * TOPK]);
    }

    auto t1 = std::chrono::steady_clock::now();

    return std::chrono::duration<double, std::milli>(t1 - t0).count();
}

static std::string cl_string(cl_device_id dev, cl_device_info info) {
    size_t size = 0;

    check_cl(
        clGetDeviceInfo(dev, info, 0, nullptr, &size),
        "clGetDeviceInfo");

    std::string s(size, '\0');

    check_cl(
        clGetDeviceInfo(dev, info, size, s.data(), nullptr),
        "clGetDeviceInfo");

    while (!s.empty() && s.back() == '\0')
        s.pop_back();

    return s;
}

static cl_device_id find_mali() {
    cl_uint np = 0;

    check_cl(
        clGetPlatformIDs(0, nullptr, &np),
        "clGetPlatformIDs");

    std::vector<cl_platform_id> platforms(np);

    check_cl(
        clGetPlatformIDs(np, platforms.data(), nullptr),
        "clGetPlatformIDs");

    for (auto p : platforms) {
        cl_uint nd = 0;

        if (clGetDeviceIDs(
                p,
                CL_DEVICE_TYPE_GPU,
                0,
                nullptr,
                &nd) != CL_SUCCESS)
            continue;

        std::vector<cl_device_id> devices(nd);

        check_cl(
            clGetDeviceIDs(
                p,
                CL_DEVICE_TYPE_GPU,
                nd,
                devices.data(),
                nullptr),
            "clGetDeviceIDs");

        for (auto d : devices) {
            std::string name = cl_string(d, CL_DEVICE_NAME);
            std::string vendor = cl_string(d, CL_DEVICE_VENDOR);

            std::cout
                << "OpenCL GPU: "
                << vendor
                << " / "
                << name
                << "\n";

            if (name.find("Mali-G610") != std::string::npos)
                return d;
        }
    }

    return nullptr;
}

static void merge_topk(
    const std::vector<uint32_t>& chunk_indices,
    const std::vector<float>& chunk_values,
    size_t chunk_offset,
    std::vector<Neighbor>& result)
{
    /*
     * Per ogni query abbiamo TOPK candidati dal chunk corrente.
     *
     * Il risultato globale viene mantenuto in:
     *
     *     result[q * TOPK ... q * TOPK + TOPK)
     *
     * Uniamo quindi al massimo 2*TOPK elementi per query.
     */

    for (uint32_t q = 0; q < QUERIES; ++q) {

        std::vector<Neighbor> candidates;
        candidates.reserve(TOPK * 2);

        Neighbor *dst = &result[q * TOPK];

        for (uint32_t k = 0; k < TOPK; ++k) {
            candidates.push_back(dst[k]);
        }

        for (uint32_t k = 0; k < TOPK; ++k) {
            uint32_t local_index =
                chunk_indices[q * TOPK + k];

            float score =
                chunk_values[q * TOPK + k];

            if (local_index == 0xffffffffu)
                continue;

            candidates.push_back({
                score,
                static_cast<uint32_t>(
                    chunk_offset + local_index)
            });
        }

        std::sort(
            candidates.begin(),
            candidates.end(),
            [](const Neighbor& a, const Neighbor& b) {
                return a.score > b.score;
            });

        for (uint32_t k = 0; k < TOPK; ++k)
            dst[k] = candidates[k];
    }
}

static double benchmark_opencl(
    const std::vector<float>& db,
    const std::vector<float>& queries,
    std::vector<Neighbor>& result,
    size_t n)
{
    cl_int err;

    cl_device_id dev = find_mali();

    if (!dev) {
        std::cerr << "Mali-G610 not found\n";
        std::exit(1);
    }

    cl_context ctx =
        clCreateContext(
            nullptr,
            1,
            &dev,
            nullptr,
            nullptr,
            &err);

    check_cl(err, "clCreateContext");

    cl_command_queue queue =
        clCreateCommandQueue(ctx, dev, 0, &err);

    check_cl(err, "clCreateCommandQueue");

const char *source = R"CLC(

#define DIM 128
#define TOPK 32

/*
 * Numero di work-item che collaborano per ogni query.
 */
#define LANES_PER_QUERY 8

/*
 * Ogni work-group contiene 8 work-item.
 *
 * Ogni work-item calcola un TOPK locale.
 * Poi il work-item 0 fonde gli 8 TOPK locali.
 */

__kernel void knn(
    __global const float *db,
    __global const float *queries,
    __global uint *indices,
    __global float *values,
    uint n)
{
    uint lane = get_local_id(0);
    uint q = get_group_id(0);

    if (q >= 256)
        return;

    /*
     * TOPK locale per ogni work-item.
     */
    float best_score[TOPK];
    uint best_index[TOPK];

    for (uint k = 0; k < TOPK; ++k) {
        best_score[k] = -3.402823466e+38f;
        best_index[k] = 0xffffffff;
    }

    __global const float *query =
        &queries[q * DIM];

    /*
     * Ogni lane analizza una parte del database.
     *
     * Esempio con 8 lane:
     *
     * lane 0 -> 0, 8, 16, 24, ...
     * lane 1 -> 1, 9, 17, 25, ...
     * ...
     * lane 7 -> 7, 15, 23, 31, ...
     */
    for (uint i = lane; i < n; i += LANES_PER_QUERY) {

        __global const float *v =
            &db[i * DIM];

        float s = 0.0f;

        for (uint j = 0; j < DIM; ++j)
            s += v[j] * query[j];

        if (s <= best_score[TOPK - 1])
            continue;

        int pos = TOPK - 1;

        while (pos > 0 &&
               s > best_score[pos - 1]) {
            --pos;
        }

        for (int k = TOPK - 1; k > pos; --k) {
            best_score[k] =
                best_score[k - 1];

            best_index[k] =
                best_index[k - 1];
        }

        best_score[pos] = s;
        best_index[pos] = i;
    }

    /*
     * TOPK locali nella local memory.
     *
     * 8 lane × 32 elementi.
     */
    __local float local_scores[LANES_PER_QUERY][TOPK];
    __local uint local_indices[LANES_PER_QUERY][TOPK];

    for (uint k = 0; k < TOPK; ++k) {
        local_scores[lane][k] = best_score[k];
        local_indices[lane][k] = best_index[k];
    }

    /*
     * Tutti i lane devono aver scritto
     * la local memory prima della riduzione.
     */
    barrier(CLK_LOCAL_MEM_FENCE);

    /*
     * Solo lane 0 esegue il merge finale.
     */
    if (lane == 0) {

        float final_score[TOPK];
        uint final_index[TOPK];

        for (uint k = 0; k < TOPK; ++k) {
            final_score[k] =
                -3.402823466e+38f;

            final_index[k] =
                0xffffffff;
        }

        /*
         * Fonde gli 8 TOPK locali.
         *
         * Al massimo 8 × 32 = 256 candidati.
         */
        for (uint l = 0;
             l < LANES_PER_QUERY;
             ++l) {

            for (uint k = 0;
                 k < TOPK;
                 ++k) {

                float s =
                    local_scores[l][k];

                uint idx =
                    local_indices[l][k];

                if (idx == 0xffffffff)
                    continue;

                if (s <= final_score[TOPK - 1])
                    continue;

                int pos = TOPK - 1;

                while (pos > 0 &&
                       s > final_score[pos - 1]) {
                    --pos;
                }

                for (int p = TOPK - 1;
                     p > pos;
                     --p) {

                    final_score[p] =
                        final_score[p - 1];

                    final_index[p] =
                        final_index[p - 1];
                }

                final_score[pos] = s;
                final_index[pos] = idx;
            }
        }

        /*
         * Risultato finale della query.
         */
        for (uint k = 0; k < TOPK; ++k) {
            indices[q * TOPK + k] =
                final_index[k];

            values[q * TOPK + k] =
                final_score[k];
        }
    }
}

)CLC";

    cl_program program =
        clCreateProgramWithSource(
            ctx,
            1,
            &source,
            nullptr,
            &err);

    check_cl(err, "clCreateProgramWithSource");

    err = clBuildProgram(
        program,
        1,
        &dev,
        "-cl-fast-relaxed-math",
        nullptr,
        nullptr);

    if (err != CL_SUCCESS) {

        size_t log_size = 0;

        clGetProgramBuildInfo(
            program,
            dev,
            CL_PROGRAM_BUILD_LOG,
            0,
            nullptr,
            &log_size);

        std::string log(log_size, '\0');

        clGetProgramBuildInfo(
            program,
            dev,
            CL_PROGRAM_BUILD_LOG,
            log_size,
            log.data(),
            nullptr);

        std::cerr << log << "\n";

        check_cl(err, "clBuildProgram");
    }

    cl_kernel kernel =
        clCreateKernel(program, "knn", &err);

    check_cl(err, "clCreateKernel");

    /*
     * Buffer GPU riutilizzato per ogni chunk.
     */
    const size_t chunk_capacity =
        std::min(CHUNK_SIZE, n);

    cl_mem db_buf =
        clCreateBuffer(
            ctx,
            CL_MEM_READ_ONLY,
            chunk_capacity * DIM * sizeof(float),
            nullptr,
            &err);

    check_cl(err, "db buffer");

    cl_mem query_buf =
        clCreateBuffer(
            ctx,
            CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
            queries.size() * sizeof(float),
            const_cast<float*>(queries.data()),
            &err);

    check_cl(err, "query buffer");

    std::vector<uint32_t> indices(
        QUERIES * TOPK);

    std::vector<float> values(
        QUERIES * TOPK);

    cl_mem index_buf =
        clCreateBuffer(
            ctx,
            CL_MEM_WRITE_ONLY,
            indices.size() * sizeof(uint32_t),
            nullptr,
            &err);

    check_cl(err, "index buffer");

    cl_mem value_buf =
        clCreateBuffer(
            ctx,
            CL_MEM_WRITE_ONLY,
            values.size() * sizeof(float),
            nullptr,
            &err);

    check_cl(err, "value buffer");

    check_cl(
        clSetKernelArg(
            kernel,
            0,
            sizeof(db_buf),
            &db_buf),
        "arg0");

    check_cl(
        clSetKernelArg(
            kernel,
            1,
            sizeof(query_buf),
            &query_buf),
        "arg1");

    check_cl(
        clSetKernelArg(
            kernel,
            2,
            sizeof(index_buf),
            &index_buf),
        "arg2");

    check_cl(
        clSetKernelArg(
            kernel,
            3,
            sizeof(value_buf),
            &value_buf),
        "arg3");

    /*
     * Risultato globale.
     *
     * Inizializziamo gli score al minimo possibile,
     * così il primo chunk può essere mergiato
     * senza una gestione speciale.
     */
    result.resize(QUERIES * TOPK);

    for (size_t i = 0; i < result.size(); ++i) {
        result[i] = {
            -3.402823466e+38f,
            0xffffffffu
        };
    }

    size_t num_chunks =
        (n + CHUNK_SIZE - 1) / CHUNK_SIZE;

    std::cout
        << "GPU CHUNK_SIZE : "
        << CHUNK_SIZE
        << "\n";

    std::cout
        << "GPU CHUNKS     : "
        << num_chunks
        << "\n";

    size_t global = QUERIES * 8;
    size_t local = 8;

double write_ms = 0.0;
double kernel_ms = 0.0;
double read_ms = 0.0;


    auto t0 =
        std::chrono::steady_clock::now();

    for (size_t chunk_offset = 0;
         chunk_offset < n;
         chunk_offset += CHUNK_SIZE) {

        size_t chunk_n =
            std::min(
                CHUNK_SIZE,
                n - chunk_offset);

        const float *chunk_ptr =
            &db[chunk_offset * DIM];

        /*
         * Copia del chunk nel buffer GPU.
         *
         * La coda è in-order, quindi il write,
         * kernel e read vengono eseguiti nell'ordine
         * corretto.
         */
auto tw0 = std::chrono::steady_clock::now();

check_cl(
    clEnqueueWriteBuffer(
        queue,
        db_buf,
        CL_TRUE,
        0,
        chunk_n * DIM * sizeof(float),
        chunk_ptr,
        0,
        nullptr,
        nullptr),
    "write db chunk");

auto tw1 = std::chrono::steady_clock::now();

write_ms +=
    std::chrono::duration<double, std::milli>(
        tw1 - tw0).count();

        uint32_t nn =
            static_cast<uint32_t>(chunk_n);

        check_cl(
            clSetKernelArg(
                kernel,
                4,
                sizeof(nn),
                &nn),
            "arg4");

auto tk0 = std::chrono::steady_clock::now();

check_cl(
    clEnqueueNDRangeKernel(
        queue,
        kernel,
        1,
        nullptr,
        &global,
        &local,
        0,
        nullptr,
        nullptr),
    "clEnqueueNDRangeKernel");

check_cl(
    clFinish(queue),
    "clFinish after kernel");

auto tk1 = std::chrono::steady_clock::now();

kernel_ms +=
    std::chrono::duration<double, std::milli>(
        tk1 - tk0).count();

auto tr0 = std::chrono::steady_clock::now();

check_cl(
    clEnqueueReadBuffer(
        queue,
        index_buf,
        CL_TRUE,
        0,
        indices.size() * sizeof(uint32_t),
        indices.data(),
        0,
        nullptr,
        nullptr),
    "read indices");

check_cl(
    clEnqueueReadBuffer(
        queue,
        value_buf,
        CL_TRUE,
        0,
        values.size() * sizeof(float),
        values.data(),
        0,
        nullptr,
        nullptr),
    "read values");

auto tr1 = std::chrono::steady_clock::now();

read_ms +=
    std::chrono::duration<double, std::milli>(
        tr1 - tr0).count();

        /*
         * Gli indici prodotti dal kernel sono locali
         * al chunk. merge_topk() aggiunge chunk_offset
         * per ottenere gli indici globali.
         */
        merge_topk(
            indices,
            values,
            chunk_offset,
            result);
    }

    auto t1 =
        std::chrono::steady_clock::now();

    clReleaseMemObject(value_buf);
    clReleaseMemObject(index_buf);
    clReleaseMemObject(query_buf);
    clReleaseMemObject(db_buf);
    clReleaseKernel(kernel);
    clReleaseProgram(program);
    clReleaseCommandQueue(queue);
    clReleaseContext(ctx);

std::cout
    << "GPU WRITE TOTAL : "
    << write_ms
    << " ms\n";

std::cout
    << "GPU KERNEL TOTAL: "
    << kernel_ms
    << " ms\n";

std::cout
    << "GPU READ TOTAL  : "
    << read_ms
    << " ms\n";

    return std::chrono::duration<double, std::milli>(
        t1 - t0).count();
}

int main(int argc, char **argv) {

    size_t n = 100000;

    if (argc > 1)
        n = std::strtoull(
            argv[1],
            nullptr,
            10);

    std::cout
        << "=== BATCH KNN BENCHMARK ===\n";

    std::cout
        << "DATABASE : "
        << n
        << "\n";

    std::cout
        << "QUERIES  : "
        << QUERIES
        << "\n";

    std::cout
        << "DIM      : "
        << DIM
        << "\n";

    std::cout
        << "TOPK     : "
        << TOPK
        << "\n";

    std::cout
        << "SEED     : 123456789\n\n";

    std::cout
        << "Generating database...\n";

    auto db =
        make_embeddings(n);

    std::cout
        << "Generating queries...\n";

    auto queries =
        make_embeddings(QUERIES);

    std::vector<Neighbor> cpu;
    std::vector<Neighbor> gpu;

    std::cout
        << "\n=== CPU ===\n";

    double cpu_ms =
        benchmark_cpu(
            db,
            queries,
            cpu,
            n);

    std::cout
        << std::fixed
        << std::setprecision(3);

    std::cout
        << "time_ms = "
        << cpu_ms
        << "\n";

    std::cout
        << "q0_top1 = "
        << cpu[0].index
        << "\n";

    std::cout
        << "\n=== OPENCL MALI ===\n";

    double gpu_ms =
        benchmark_opencl(
            db,
            queries,
            gpu,
            n);

    std::cout
        << "time_ms = "
        << gpu_ms
        << "\n";

    std::cout
        << "q0_top1 = "
        << gpu[0].index
        << "\n";

    size_t total =
        QUERIES * TOPK;

    size_t same = 0;

    for (size_t i = 0;
         i < total;
         ++i) {

        if (cpu[i].index == gpu[i].index)
            ++same;
    }

    std::cout
        << "exact_topk_positions = "
        << same
        << "/"
        << total
        << "\n";

    std::cout
        << "exact_ratio = "
        << (double(same) /
            double(total)) * 100.0
        << "%\n";

for (size_t i = 0; i < total; ++i) {
    if (cpu[i].index != gpu[i].index) {
        size_t q = i / TOPK;
        size_t k = i % TOPK;

        std::cout
            << "MISMATCH q=" << q
            << " k=" << k
            << " CPU index=" << cpu[i].index
            << " score=" << cpu[i].score
            << " GPU index=" << gpu[i].index
            << " score=" << gpu[i].score
            << "\n";
    }
}



    std::cout
        << "\n=== SPEEDUP ===\n";

    std::cout
        << "CPU/OpenCL = "
        << cpu_ms / gpu_ms
        << "x\n";

    double pairs =
        double(n) *
        double(QUERIES);

    double billion =
        pairs *
        double(DIM) /
        1e9;

    std::cout
        << "\n=== WORK ===\n";

    std::cout
        << "similarity_ops = "
        << billion
        << " billion dot-product elements\n";

    std::cout
        << "\n=== THROUGHPUT ===\n";

    std::cout
        << "database_vectors_per_second = "
        << (double(n) *
            double(QUERIES))
           / (gpu_ms / 1000.0)
        << "\n";

    return 0;
}