#define CL_TARGET_OPENCL_VERSION 300
#include <CL/cl.h>

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

static constexpr int DIM = 128;
static constexpr int DIM4 = 32;
static constexpr int TOPK = 32;
static constexpr int QUERIES = 32;
static constexpr int CHUNK_SIZE = 5000;
static constexpr int LANES = 8;
static constexpr int TOTAL_DB = 12000;

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

static const char* kernel_source = R"CLC(

#define DIM4 32
#define TOPK 32
#define LANES 8
#define QUERIES 32

inline float dot_f4(
    __global const float4* db,
    __global const float4* queries,
    uint q,
    uint db_index)
{
    uint db_base = db_index * DIM4;
    uint q_base  = q * DIM4;

    float4 acc0 = (float4)(0.0f);
    float4 acc1 = (float4)(0.0f);
    float4 acc2 = (float4)(0.0f);
    float4 acc3 = (float4)(0.0f);

    for (uint d = 0; d < DIM4; d += 4) {
        acc0 += db[db_base + d + 0] * queries[q_base + d + 0];
        acc1 += db[db_base + d + 1] * queries[q_base + d + 1];
        acc2 += db[db_base + d + 2] * queries[q_base + d + 2];
        acc3 += db[db_base + d + 3] * queries[q_base + d + 3];
    }

    float4 s = acc0 + acc1 + acc2 + acc3;
    return s.x + s.y + s.z + s.w;
}

__kernel void knn_minreplace(
    __global const float4* db,
    __global const float4* queries,
    uint db_count,
    __global float* local_scores_out,
    __global uint* local_indices_out)
{
    uint gid  = get_global_id(0);
    uint lane = gid % LANES;
    uint q    = gid / LANES;

    if (q >= QUERIES)
        return;

    float best[TOPK];
    uint  best_idx[TOPK];

    for (uint k = 0; k < TOPK; ++k) {
        best[k] = -INFINITY;
        best_idx[k] = 0;
    }

    float min_score = -INFINITY;
    uint min_pos = 0;

    for (uint i = lane; i < db_count; i += LANES) {

        float score = dot_f4(db, queries, q, i);

        if (score > min_score) {

            best[min_pos] = score;
            best_idx[min_pos] = i;

            min_score = best[0];
            min_pos = 0;

            for (uint k = 1; k < TOPK; ++k) {
                if (best[k] < min_score) {
                    min_score = best[k];
                    min_pos = k;
                }
            }
        }
    }

    uint base = (q * LANES + lane) * TOPK;

    for (uint k = 0; k < TOPK; ++k) {
        local_scores_out[base + k] = best[k];
        local_indices_out[base + k] = best_idx[k];
    }
}

__kernel void merge_topk(
    __global const float* local_scores,
    __global const uint* local_indices,
    __global float* final_scores,
    __global uint* final_indices)
{
    uint q = get_global_id(0);

    if (q >= QUERIES)
        return;

    float best[TOPK];
    uint  best_idx[TOPK];

    for (uint k = 0; k < TOPK; ++k) {
        best[k] = -INFINITY;
        best_idx[k] = 0;
    }

    float min_score = -INFINITY;
    uint min_pos = 0;

    for (uint lane = 0; lane < LANES; ++lane) {

        uint base = (q * LANES + lane) * TOPK;

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

            float score = local_scores[base + k];

            if (score > min_score) {

                best[min_pos] = score;
                best_idx[min_pos] = local_indices[base + k];

                min_score = best[0];
                min_pos = 0;

                for (uint j = 1; j < TOPK; ++j) {
                    if (best[j] < min_score) {
                        min_score = best[j];
                        min_pos = j;
                    }
                }
            }
        }
    }

    // Sort final TOPK descending, to match the semantic output
    // of the original benchmark.
    for (uint i = 0; i < TOPK; ++i) {
        uint max_pos = i;

        for (uint j = i + 1; j < TOPK; ++j) {
            if (best[j] > best[max_pos])
                max_pos = j;
        }

        if (max_pos != i) {
            float ts = best[i];
            best[i] = best[max_pos];
            best[max_pos] = ts;

            uint ti = best_idx[i];
            best_idx[i] = best_idx[max_pos];
            best_idx[max_pos] = ti;
        }
    }

    uint out_base = q * TOPK;

    for (uint k = 0; k < TOPK; ++k) {
        final_scores[out_base + k] = best[k];
        final_indices[out_base + k] = best_idx[k];
    }
}

)CLC";

int main() {
    cl_int err;

    cl_uint num_platforms = 0;
    check(
        clGetPlatformIDs(0, nullptr, &num_platforms),
        "clGetPlatformIDs");

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

    check(
        clGetPlatformIDs(
            num_platforms,
            platforms.data(),
            nullptr),
        "clGetPlatformIDs(list)");

    cl_platform_id platform = platforms[0];

    cl_uint num_devices = 0;

    check(
        clGetDeviceIDs(
            platform,
            CL_DEVICE_TYPE_GPU,
            0,
            nullptr,
            &num_devices),
        "clGetDeviceIDs(count)");

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

    check(
        clGetDeviceIDs(
            platform,
            CL_DEVICE_TYPE_GPU,
            num_devices,
            devices.data(),
            nullptr),
        "clGetDeviceIDs(list)");

    cl_device_id device = devices[0];

    char device_name[256] = {};

    clGetDeviceInfo(
        device,
        CL_DEVICE_NAME,
        sizeof(device_name),
        device_name,
        nullptr);

    std::cout << "=== FULL KNN FLOAT4 + MIN-REPLACE ===\n";
    std::cout << "Device       : " << device_name << "\n";
    std::cout << "DIM          : " << DIM << "\n";
    std::cout << "TOPK         : " << TOPK << "\n";
    std::cout << "QUERIES      : " << QUERIES << "\n";
    std::cout << "TOTAL_DB     : " << TOTAL_DB << "\n";
    std::cout << "CHUNK_SIZE   : " << CHUNK_SIZE << "\n";
    std::cout << "LANES        : " << LANES << "\n\n";

    cl_context context =
        clCreateContext(
            nullptr,
            1,
            &device,
            nullptr,
            nullptr,
            &err);

    check(err, "clCreateContext");

    cl_command_queue queue =
        clCreateCommandQueue(
            context,
            device,
            CL_QUEUE_PROFILING_ENABLE,
            &err);

    check(err, "clCreateCommandQueue");

    const size_t db_bytes =
        static_cast<size_t>(CHUNK_SIZE) *
        DIM4 *
        4 *
        sizeof(float);

    const size_t query_bytes =
        static_cast<size_t>(QUERIES) *
        DIM4 *
        4 *
        sizeof(float);

    const size_t lane_candidates =
        static_cast<size_t>(QUERIES) *
        LANES *
        TOPK;

    const size_t candidate_bytes =
        lane_candidates * sizeof(float);

    const size_t index_bytes =
        lane_candidates * sizeof(cl_uint);

    const size_t final_count =
        static_cast<size_t>(QUERIES) * TOPK;

    const size_t final_score_bytes =
        final_count * sizeof(float);

    const size_t final_index_bytes =
        final_count * sizeof(cl_uint);

    cl_mem db_buf =
        clCreateBuffer(
            context,
            CL_MEM_READ_ONLY,
            db_bytes,
            nullptr,
            &err);

    check(err, "clCreateBuffer(db)");

    cl_mem query_buf =
        clCreateBuffer(
            context,
            CL_MEM_READ_ONLY,
            query_bytes,
            nullptr,
            &err);

    check(err, "clCreateBuffer(query)");

    cl_mem local_scores_buf =
        clCreateBuffer(
            context,
            CL_MEM_READ_WRITE,
            candidate_bytes,
            nullptr,
            &err);

    check(err, "clCreateBuffer(local_scores)");

    cl_mem local_indices_buf =
        clCreateBuffer(
            context,
            CL_MEM_READ_WRITE,
            index_bytes,
            nullptr,
            &err);

    check(err, "clCreateBuffer(local_indices)");

    cl_mem final_scores_buf =
        clCreateBuffer(
            context,
            CL_MEM_READ_WRITE,
            final_score_bytes,
            nullptr,
            &err);

    check(err, "clCreateBuffer(final_scores)");

    cl_mem final_indices_buf =
        clCreateBuffer(
            context,
            CL_MEM_READ_WRITE,
            final_index_bytes,
            nullptr,
            &err);

    check(err, "clCreateBuffer(final_indices)");

    cl_program program =
        clCreateProgramWithSource(
            context,
            1,
            &kernel_source,
            nullptr,
            &err);

    check(err, "clCreateProgramWithSource");

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

    if (err != CL_SUCCESS) {

        size_t log_size = 0;

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

        std::vector<char> log(log_size + 1);

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

        std::cerr << log.data() << "\n";

        check(err, "clBuildProgram");
    }

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

    check(err, "clCreateKernel(knn_minreplace)");

    cl_kernel merge =
        clCreateKernel(
            program,
            "merge_topk",
            &err);

    check(err, "clCreateKernel(merge_topk)");

    std::vector<float> db(
        static_cast<size_t>(TOTAL_DB) * DIM);

    std::vector<float> queries(
        static_cast<size_t>(QUERIES) * DIM);

    std::mt19937 rng_db(123456789);
    std::mt19937 rng_q(987654321);

    std::normal_distribution<float> normal(0.0f, 1.0f);

    for (int i = 0; i < TOTAL_DB; ++i) {

        float norm = 0.0f;

        for (int d = 0; d < DIM; ++d) {

            float x = normal(rng_db);

            db[
                static_cast<size_t>(i) * DIM + d
            ] = x;

            norm += x * x;
        }

        norm = 1.0f / std::sqrt(norm);

        for (int d = 0; d < DIM; ++d)
            db[
                static_cast<size_t>(i) * DIM + d
            ] *= norm;
    }

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

        float norm = 0.0f;

        for (int d = 0; d < DIM; ++d) {

            float x = normal(rng_q);

            queries[
                static_cast<size_t>(q) * DIM + d
            ] = x;

            norm += x * x;
        }

        norm = 1.0f / std::sqrt(norm);

        for (int d = 0; d < DIM; ++d)
            queries[
                static_cast<size_t>(q) * DIM + d
            ] *= norm;
    }

    std::vector<float> final_scores(
        final_count);

    std::vector<cl_uint> final_indices(
        final_count);

    /*
     * TOP-K globale persistente fra tutti i chunk.
     *
     * final_scores/final_indices contengono il risultato
     * del chunk corrente.
     *
     * global_scores/global_indices contengono invece il
     * TOP-K complessivo elaborato fino a quel momento.
     */
    std::vector<float> global_scores(
        final_count,
        -INFINITY);

    std::vector<cl_uint> global_indices(
        final_count,
        0);

    double total_knn_ms = 0.0;
    double total_merge_ms = 0.0;
    double total_global_merge_ms = 0.0;

    const int chunks =
        (TOTAL_DB + CHUNK_SIZE - 1) /
        CHUNK_SIZE;

    for (int chunk = 0; chunk < chunks; ++chunk) {

        const int begin =
            chunk * CHUNK_SIZE;

        const int count =
            std::min(
                CHUNK_SIZE,
                TOTAL_DB - begin);

        const size_t db_upload_bytes =
            static_cast<size_t>(count) *
            DIM4 *
            4 *
            sizeof(float);

        check(
            clEnqueueWriteBuffer(
                queue,
                db_buf,
                CL_TRUE,
                0,
                db_upload_bytes,
                db.data() +
                    static_cast<size_t>(begin) * DIM,
                0,
                nullptr,
                nullptr),
            "write db");

        check(
            clEnqueueWriteBuffer(
                queue,
                query_buf,
                CL_TRUE,
                0,
                query_bytes,
                queries.data(),
                0,
                nullptr,
                nullptr),
            "write queries");

        check(
            clSetKernelArg(
                knn,
                0,
                sizeof(cl_mem),
                &db_buf),
            "arg 0");

        check(
            clSetKernelArg(
                knn,
                1,
                sizeof(cl_mem),
                &query_buf),
            "arg 1");

        cl_uint db_count = count;

        check(
            clSetKernelArg(
                knn,
                2,
                sizeof(cl_uint),
                &db_count),
            "arg 2");

        check(
            clSetKernelArg(
                knn,
                3,
                sizeof(cl_mem),
                &local_scores_buf),
            "arg 3");

        check(
            clSetKernelArg(
                knn,
                4,
                sizeof(cl_mem),
                &local_indices_buf),
            "arg 4");

        size_t global =
            static_cast<size_t>(QUERIES) *
            LANES;

        size_t local = LANES;

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

        check(
            clEnqueueNDRangeKernel(
                queue,
                knn,
                1,
                nullptr,
                &global,
                &local,
                0,
                nullptr,
                nullptr),
            "enqueue knn");

        check(clFinish(queue), "finish knn");

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

        total_knn_ms +=
            std::chrono::duration<double, std::milli>(
                t1 - t0).count();

        check(
            clSetKernelArg(
                merge,
                0,
                sizeof(cl_mem),
                &local_scores_buf),
            "merge arg 0");

        check(
            clSetKernelArg(
                merge,
                1,
                sizeof(cl_mem),
                &local_indices_buf),
            "merge arg 1");

        check(
            clSetKernelArg(
                merge,
                2,
                sizeof(cl_mem),
                &final_scores_buf),
            "merge arg 2");

        check(
            clSetKernelArg(
                merge,
                3,
                sizeof(cl_mem),
                &final_indices_buf),
            "merge arg 3");

        size_t merge_global = QUERIES;

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

        check(
            clEnqueueNDRangeKernel(
                queue,
                merge,
                1,
                nullptr,
                &merge_global,
                nullptr,
                0,
                nullptr,
                nullptr),
            "enqueue merge");

        check(clFinish(queue), "finish merge");

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

        total_merge_ms +=
            std::chrono::duration<double, std::milli>(
                tm1 - tm0).count();

        check(
            clEnqueueReadBuffer(
                queue,
                final_scores_buf,
                CL_TRUE,
                0,
                final_score_bytes,
                final_scores.data(),
                0,
                nullptr,
                nullptr),
            "read final scores");

        check(
            clEnqueueReadBuffer(
                queue,
                final_indices_buf,
                CL_TRUE,
                0,
                final_index_bytes,
                final_indices.data(),
                0,
                nullptr,
                nullptr),
            "read final indices");

        /*
         * Merge CPU fra:
         *
         *   TOP-K globale accumulato finora
         *   TOP-K prodotto dal chunk corrente
         *
         * Gli indici restituiti dalla GPU sono locali al
         * chunk. Li convertiamo quindi in indici globali
         * aggiungendo begin.
         */
        auto global_merge_start =
            std::chrono::steady_clock::now();

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

            const size_t base =
                static_cast<size_t>(q) * TOPK;

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

                const float candidate_score =
                    final_scores[base + k];

                const cl_uint candidate_index =
                    static_cast<cl_uint>(begin) +
                    final_indices[base + k];

                /*
                 * Trova il risultato peggiore nel TOP-K
                 * globale corrente.
                 */
                size_t worst_pos = base;
                float worst_score =
                    global_scores[base];

                for (int j = 1; j < TOPK; ++j) {

                    const size_t pos =
                        base +
                        static_cast<size_t>(j);

                    if (global_scores[pos] <
                        worst_score)
                    {
                        worst_score =
                            global_scores[pos];

                        worst_pos = pos;
                    }
                }

                /*
                 * Sostituzione del minimo.
                 */
                if (candidate_score > worst_score) {

                    global_scores[worst_pos] =
                        candidate_score;

                    global_indices[worst_pos] =
                        candidate_index;
                }
            }
        }

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

        total_global_merge_ms +=
            std::chrono::duration<
                double,
                std::milli>(
                    global_merge_end -
                    global_merge_start).count();
    }

    /*
     * Ordina il TOP-K globale di ogni query in ordine
     * decrescente di similarita.
     */
    for (int q = 0; q < QUERIES; ++q) {

        const size_t base =
            static_cast<size_t>(q) * TOPK;

        for (int i = 0; i < TOPK; ++i) {

            int max_pos = i;

            for (int j = i + 1;
                 j < TOPK;
                 ++j)
            {
                if (global_scores[
                        base +
                        static_cast<size_t>(j)
                    ] >
                    global_scores[
                        base +
                        static_cast<size_t>(max_pos)
                    ])
                {
                    max_pos = j;
                }
            }

            if (max_pos != i) {

                std::swap(
                    global_scores[
                        base +
                        static_cast<size_t>(i)
                    ],
                    global_scores[
                        base +
                        static_cast<size_t>(max_pos)
                    ]);

                std::swap(
                    global_indices[
                        base +
                        static_cast<size_t>(i)
                    ],
                    global_indices[
                        base +
                        static_cast<size_t>(max_pos)
                    ]);
            }
        }
    }

    /*
     * Verifica intervallo indici globali.
     */
    size_t invalid_indices = 0;

    for (cl_uint index : global_indices) {
        if (index >=
            static_cast<cl_uint>(TOTAL_DB))
        {
            ++invalid_indices;
        }
    }

    double checksum_scores = 0.0;
    double checksum_indices = 0.0;

    for (size_t i = 0; i < final_count; ++i) {
        checksum_scores += global_scores[i];
        checksum_indices += global_indices[i];
    }

    double vectors =
        static_cast<double>(TOTAL_DB) *
        QUERIES;

    double vector_rate =
        vectors /
        total_knn_ms /
        1000.0;

    double gflops =
        vector_rate *
        DIM *
        2.0 /
        1000.0;

    double db_traffic =
        vector_rate *
        DIM *
        sizeof(float) /
        1000.0;

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

    std::cout << "KNN kernel time : "
              << std::fixed
              << std::setprecision(3)
              << total_knn_ms
              << " ms\n";

    std::cout << "Merge time      : "
              << std::fixed
              << std::setprecision(3)
              << total_merge_ms
              << " ms\n";

    std::cout << "Global CPU merge: "
              << std::fixed
              << std::setprecision(3)
              << total_global_merge_ms
              << " ms\n";

    std::cout << "Total GPU time  : "
              << std::fixed
              << std::setprecision(3)
              << total_knn_ms + total_merge_ms
              << " ms\n";

    std::cout << "Search core time: "
              << std::fixed
              << std::setprecision(3)
              << total_knn_ms +
                     total_merge_ms +
                     total_global_merge_ms
              << " ms\n";

    std::cout << "Invalid indices : "
              << invalid_indices
              << "\n";

    std::cout << "Vector rate     : "
              << std::fixed
              << std::setprecision(3)
              << vector_rate
              << " M vector/s\n";

    std::cout << "Dot rate        : "
              << std::fixed
              << std::setprecision(3)
              << gflops
              << " GFLOP/s\n";

    std::cout << "DB traffic      : "
              << std::fixed
              << std::setprecision(3)
              << db_traffic
              << " GB/s effective\n";

    std::cout << "Score checksum  : "
              << std::fixed
              << std::setprecision(3)
              << checksum_scores
              << "\n";

    std::cout << "Index checksum  : "
              << std::fixed
              << std::setprecision(3)
              << checksum_indices
              << "\n";

    clReleaseKernel(knn);
    clReleaseKernel(merge);
    clReleaseProgram(program);

    clReleaseMemObject(db_buf);
    clReleaseMemObject(query_buf);
    clReleaseMemObject(local_scores_buf);
    clReleaseMemObject(local_indices_buf);
    clReleaseMemObject(final_scores_buf);
    clReleaseMemObject(final_indices_buf);

    clReleaseCommandQueue(queue);
    clReleaseContext(context);

    return 0;
}