#include <algorithm>
#include <chrono>
#include <cstdio>
#include <execution>
#include <numeric>
#include <vector>
//#include "matmul_cpp_functor.h"

int main() {
    constexpr int N     = 16384;   // = 16k
    constexpr int total = N * N;   // = 256M, can be represented in int32

    // std::vector is backed by CUDA unified memory when compiled with
    // nvc++ -stdpar=gpu, so host and device can both access these arrays
    // without explicit transfers.
    std::vector<float> A(total), B(total), C(total);

    // ------------------------------------------------------------------
    // Index vector: the standard way to express 2D parallel work over a
    // flat integer range with C++ execution policies.  We fill it on the
    // GPU with par_unseq so the data never has to move through the host.
    //
    // A custom counting iterator would avoid this allocation, but keeping
    // the index vector makes the pattern maximally portable across stdpar
    // implementations (nvc++, Intel oneAPI, GCC with TBB).
    // ------------------------------------------------------------------
    std::vector<int> idx(total);
    std::iota(idx.begin(), idx.end(), 0);

    // Capture raw pointers: lambdas that run on the GPU must not capture
    // std::vector objects by reference - only trivially copyable values.
    float* pA = A.data();
    float* pB = B.data();
    float* pC = C.data();

    // ------------------------------------------------------------------
    // Initialization - par_unseq on GPU
    // Each thread owns one (i,j) element of the flat index space.
    // ------------------------------------------------------------------
    std::for_each(std::execution::par_unseq, idx.begin(), idx.end(),
        [=](int n) {
            int i    = n / N;
            int j    = n % N;
            pA[n] = static_cast<float>((i + j) %  97 + 1) * 0.01f;
            pB[n] = static_cast<float>((i * j) % 101 + 1) * 0.01f;
            pC[n] = 0.0f;
        });

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

    // ------------------------------------------------------------------
    // Dense matrix multiply  C = A * B - par_unseq on GPU
    //
    // par_unseq allows the implementation to parallelize AND vectorize
    // iterations in any order.  nvc++ -stdpar=gpu turns this into a CUDA
    // kernel: the flat index n maps to a GPU thread, and the inner k-loop
    // runs sequentially inside each thread, exactly like the Fortran
    // DO CONCURRENT version.
    //
    // stdpar calls block until the GPU finishes, so std::chrono timing
    // around this call accurately measures wall-clock GPU time.
    // ------------------------------------------------------------------
    std::for_each(std::execution::par_unseq, idx.begin(), idx.end(),
        [=](int n) {
            int   i   = n / N;
            int   j   = n % N;
            float acc = 0.0f;
            for (int k = 0; k < N; ++k)
                acc += pA[i * N + k] * pB[k * N + j];
            pC[n] = acc;
        });
    // To use the functor in matmul_cpp_functor.h instead of a lambda:
    // comment out the lambda above, then uncomment the #include at the
    // top of this file, as well as the next line -
    //    MatMul{ pA, pB, pC, N });

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

    double elapsed = std::chrono::duration<double>(t1 - t0).count();
    double gflops  = 2.0 * N * N * N / (elapsed * 1.0e9);

    // Spot-check: sequential loop on the host reads back a small corner.
    float checksum = 0.0f;
    for (int i = 0; i < 16; ++i)
        for (int j = 0; j < 16; ++j)
            checksum += C[i * N + j];

    std::printf("Matrix size : %d x %d\n", N, N);
    std::printf("Elapsed     : %.3f s\n",   elapsed);
    std::printf("Performance : %.3f GFLOP/s\n", gflops);
    std::printf("Checksum    : %e\n",        checksum);

    return 0;
}