#include <stdio.h>
#include <stdlib.h>

#include "../common.h"

__global__ void matmul(int **a, int **b, int **c, int n) {
    printf("%d\t%d\n", blockIdx.x, threadIdx.x);
    int i = blockIdx.x, j = threadIdx.x;
    c[i][j] = 0;
    for (int k = 0; k < n; k++) {
        c[i][j] += a[i][k] * b[k][j];
    }
}

void init(int ***pa, int ***pb, int ***pc, int n) {
    CHECK(cudaMallocManaged(pa, n * sizeof(int*)));
    CHECK(cudaMallocManaged(pb, n * sizeof(int*)));
    CHECK(cudaMallocManaged(pc, n * sizeof(int*)));
    int **a = *pa, **b = *pb, **c = *pc;

    char file[16];
    sprintf(file, "%d.in", n);

    FILE *fp = fopen(file, "rb");
    for (int i = 0; i < n; i++) {
        CHECK(cudaMallocManaged(a + i, n * sizeof(int)));
        CHECK(cudaMallocManaged(c + i, n * sizeof(int)));
        fread(a[i], sizeof(int), n, fp);
    }
    for (int i = 0; i < n; i++) {
        CHECK(cudaMallocManaged(b + i, n * sizeof(int)));
        fread(b[i], sizeof(int), n, fp);
    }
    fclose(fp);
}

int main(int argc, char **argv) {
    int n = atoi(argv[1]);
    int **a, **b, **c;
    init(&a, &b, &c, n);

    matmul<<<n, n>>>(a, b, c, n);
    CHECK(cudaDeviceSynchronize());

    printf("a: %d\nb: %d\nc: %d\n", a[100][100], b[100][100], c[100][100]);

    char file[16];
    sprintf(file, "%d.out", n);
    FILE *fp = fopen(file, "wb");
    for (int i = 0; i < n; i++) {
        fwrite(c[i], sizeof(int), n, fp);
    }

    fclose(fp);
    for (int i = 0; i < n; i++) {
        CHECK(cudaFree(a[i]));
        CHECK(cudaFree(b[i]));
        CHECK(cudaFree(c[i]));
    }
    CHECK(cudaFree(a));
    CHECK(cudaFree(b));
    CHECK(cudaFree(c));

    return 0;
}