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

#define N 2048
#define NUM_THREADS 16  // Set desired number of threads

void load_matrix(const char *filename, double *matrix) {
    FILE *file = fopen(filename, "rb");
    if (!file) {
        perror("Cannot read file");
        exit(EXIT_FAILURE);
    }
    fread(matrix, sizeof(double), N * N, file);
    fclose(file);
}

int main() {
    char matrix_a_path[256], matrix_b_path[256];
    sprintf(matrix_a_path, "../dataset/%d/matrix_A.bin", N);
    sprintf(matrix_b_path, "../dataset/%d/matrix_B.bin", N);

    double *A = (double *)malloc(N * N * sizeof(double));
    double *B = (double *)malloc(N * N * sizeof(double));
    double *C = (double *)calloc(N * N, sizeof(double));

    if (!A || !B || !C) {
        perror("Memory allocation failed");
        exit(EXIT_FAILURE);
    }

    load_matrix(matrix_a_path, A);
    load_matrix(matrix_b_path, B);

    omp_set_num_threads(NUM_THREADS);  // Set thread limit
    #pragma omp parallel for
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < N; j++) {
            for (int k = 0; k < N; k++) {
                C[i * N + j] += A[i * N + k] * B[k * N + j];
            }
        }
    }

    free(A);
    free(B);
    free(C);
    return 0;
}