#include <iostream>
#include <chrono>
#include <vector>
#include <random>

#include <arrayfire.h>

void ind(std::vector<int>& rows, std::vector<int>& columns, std::vector<float>& values, int maksimi, int Nx, int Ny) {
	rows.emplace_back(0);
	std::default_random_engine generator;
	std::normal_distribution<float> distribution(0.f, 1.0f);

	for (int kk = 0; kk < Nx; kk++) {
		int N_row = std::rand() % maksimi + 1;
		int temp = 0;
		for (int ll = 0; ll < N_row; ll++) {
			int apu = std::rand() % (Ny - N_row + ll - 1);
			if (apu <= temp)
				temp++;
			else
				temp = apu;
			columns.emplace_back(temp);
			values.emplace_back(distribution(generator));
		}
		rows.emplace_back(rows[kk] + N_row);
	}
}

int main() {
    // std::cout << "---------- OPENCL ------------\n";
    // af::setBackend(AF_BACKEND_OPENCL);
    // af::setDevice(1);

    std::cout << "---------- CUDA ------------\n";
    af::setBackend(AF_BACKEND_CUDA);
    af::setDevice(0);

    af::info();

	int Nx = 128 * 128 * 109 * 2;
	int maksimi = 128;
	int Ny = maksimi * maksimi * 109;
    
	std::vector<int> columns;
	std::vector<int> rows;
	std::vector<float> values;
    
	for (int ii = 0; ii < 3; ii++) {
		ind(rows, columns, values, maksimi, Nx, Ny);

		af::array element_ar(values.size(), &values.front());
		af::array indices_ar(columns.size(), &columns.front());
		af::array lor_ar(rows.size(), &rows.front());
		af::array H = af::sparse(lor_ar.dims(0) - 1, Ny, element_ar, lor_ar, indices_ar);

        // Calculate result once to profile runtime instead of JIT compilation
		af::array Summ = af::matmul(H, af::constant(1, H.dims(0), 1), AF_MAT_TRANS);
        Summ.eval();
        Summ = af::array{};

        // Prepare inputs
        H.eval();
        af::sync();

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

		Summ = af::matmul(H, af::constant(1, H.dims(0), 1), AF_MAT_TRANS);

        // Force evaluation
        Summ.eval();   
        af::sync();

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

        std::cout << "iteration " << ii << ", " << std::chrono::duration_cast<std::chrono::microseconds>(end - begin).count() << " us \n";
		rows.clear();
		columns.clear();
		values.clear();
	}

    std::cout << std::endl;
	return 0;
}