// Source: https://github.com/ivanbgd/Matrix-Multiplication-MatMul-C/blob/master/matmul_1d_seq.c
#define MATMUL_1D
#ifdef MATMUL_1D

/* Matrices are represented as 1-D arrays in memory.
 * That means they are contiguous in memory.
 * Minimum dimension is 1, not 0, and internal dimensions must match. */

// #include <math.h>
// #include <omp.h>
// #include <stdio.h>
// #include <stdlib.h>
// #include <string.h>
// #include <time.h>
#include <cheriintrin.h>
#include <stddef.h>
#include <stdint.h>

void free(void * __capability ptr);
void * __capability malloc(size_t size);

/* Initializes vector or matrix, sequentially, with indices. */
void init_seq(int * __capability a, const unsigned n_rows_a, const unsigned n_cols_a) {
    for (size_t i = 0; i < n_rows_a; i++) {
        for (size_t j = 0; j < n_cols_a; j++) {
            a[i*n_cols_a + j] = 2;
        }
    }
}

// static unsigned int seed = 123456789;  // must be initialized somehow

// unsigned int rand_test(void)
// {
//     seed = (1103515245 * seed + 12345);
//     return (seed >> 16) & 0x7FFF; // 15-bit result
// }

// /* Initializes vector or matrix, randomly. */
// void init_rand(double *a, const unsigned n_rows_a, const unsigned n_cols_a) {
//     for (size_t i = 0; i < n_rows_a; i++) {
//         for (size_t j = 0; j < n_cols_a; j++) {
//             a[i*n_cols_a + j] = rand_test() / 32767.0;
//             // rand_test() / (double)RAND_MAX;
//         }
//     }
// }

/*  Takes and returns a new matrix, t, which is a transpose of the original one, m.
    It's also flat in memory, i.e., 1-D, but it should be looked at as a transpose
    of m, meaning, n_rows_t == n_cols_m, and n_cols_t == n_rows_m.
    The original matrix m stays intact. */
int __capability * transpose(const int __capability *m, const unsigned n_rows_m, const unsigned n_cols_m, int __capability *t) {
    for (size_t i = 0; i < n_rows_m; i++) {
        for (size_t j = 0; j < n_cols_m; j++) {
            t[j*n_rows_m + i] = m[i*n_cols_m + j];
        }
    }

    return t;
}

/* Dot product of two arrays, or matrix product
 * Allocates and returns an array.
 * This variant doesn't transpose matrix b, and it's a lot slower. */
int * __capability dot_simple(const int * __capability a, const unsigned n_rows_a, const unsigned n_cols_a,\
                   const int * __capability b, const unsigned n_rows_b, const unsigned n_cols_b) {

    if (n_cols_a != n_rows_b) {
        // printf("#columns A must be equal to #rows B!\n");
        // system("pause");
        // exit(-2);
        while(1);
    }

    int * __capability c = malloc(n_rows_a * n_cols_b * sizeof(*c));
    if (c == NULL) {
        // printf("Couldn't allocate memory!\n");
        // system("pause");
        // exit(-1);
        while(1);
    }

    for (size_t i = 0; i < n_rows_a; i++) {
        for (size_t k = 0; k < n_cols_b; k++) {
            int sum = 0;
            for (size_t j = 0; j < n_cols_a; j++) {
                sum += a[i*n_cols_a + j] * b[j*n_cols_b + k];
            }
            c[i*n_cols_b + k] = sum;
        }
    }

    return c;
}

/* Dot product of two arrays, or matrix product
 * Allocates and returns an array.
 * This variant transposes matrix b, and it's a lot faster. */
int * __capability dot(const int  * __capability a, const unsigned n_rows_a, const unsigned n_cols_a, \
            const int * __capability b, const unsigned n_rows_b, const unsigned n_cols_b) {

    if (n_cols_a != n_rows_b) {
        // printf("#columns A must be equal to #rows B!\n");
        // system("pause");
        // exit(-2);
    }

    int * __capability bt = malloc(n_rows_b * n_cols_b * sizeof(*b));

    int * __capability c = malloc(n_rows_a * n_cols_b * sizeof(*c));    

    if ((c == NULL) || (bt == NULL)) {
        // printf("Couldn't allocate memory!\n");
        // system("pause");
        // exit(-1);
    }

    bt = transpose(b, n_rows_b, n_cols_b, bt);

    for (size_t i = 0; i < n_rows_a; i++) {
        for (size_t k = 0; k < n_cols_b; k++) {
            int sum = 0;
            for (size_t j = 0; j < n_cols_a; j++) {
                sum += a[i*n_cols_a + j] * bt[k*n_rows_b + j];
            }
            c[i*n_cols_b + k] = sum;
        }
    }

    free(bt);

    return c;
}

/* Prints vector, or matrix. */
// void print(const double *a, const unsigned n_rows_a, const unsigned n_cols_a) {
//     for (size_t i = 0; i < n_rows_a; i++) {
//         for (size_t j = 0; j < n_cols_a; j++) {
//             printf("%8.3f ", a[i*n_cols_a + j]);
//         }
//         printf("\n");
//     }
//     printf("\n");
// }

int main(void) {
    /* Intializes random number generator */
    // time_t t;
    // srand((unsigned)time(&t));
    // srand(0);

    /* For measuring time */
    int t0, t1;

    const unsigned scale = 20;
    const unsigned n_rows_a = 4 * scale;
    const unsigned n_cols_a = 3 * scale;
    const unsigned n_rows_b = 3 * scale;
    const unsigned n_cols_b = 2 * scale;

    int __capability *a = malloc(n_rows_a * n_cols_a * sizeof(*a));
    int __capability *b = malloc(n_rows_b * n_cols_b * sizeof(*b));
    int __capability *c = NULL;
    int __capability *d = NULL;

    if (!a || !b) {
        // printf("Couldn't allocate memory!\n");
        while(1);
        // system("pause");
        // exit(-1);
    }

    // init_rand(a, n_rows_a, n_cols_a);
    // init_rand(b, n_rows_b, n_cols_b);

    init_seq(a, n_rows_a, n_cols_a);
    init_seq(b, n_rows_b, n_cols_b);

    // for (size_t i = 0; i < n_rows_a; i++) {
    //     for (size_t j = 0; j < n_cols_a; j++) {
    //         a[i*n_cols_a + j] = i*n_cols_a + j;
    //     }
    // }

    // t0 = omp_get_wtime();
    c = dot_simple(a, n_rows_a, n_cols_a, b, n_rows_b, n_cols_b);
    // t1 = omp_get_wtime();
    // printf("Dot Simple: Elapsed time %.3f s\n", t1 - t0);

    // // t0 = omp_get_wtime();
    d = dot(a, n_rows_a, n_cols_a, b, n_rows_b, n_cols_b);
    // t1 = omp_get_wtime();
    // printf("Dot: Elapsed time %.3f s\n", t1 - t0);

    // if (scale == 1) {
    //     // printf("Matrix A:\n");
    //     // print(a, n_rows_a, n_cols_a);
    //     // // printf("Matrix B:\n");
    //     // print(b, n_rows_b, n_cols_b);
    //     // // printf("Matrix C:\n");
    //     // print(c, n_rows_a, n_cols_b);
    //     // // printf("Matrix D:\n");
    //     // print(d, n_rows_a, n_cols_b);
    // }

    free(a);
    free(b);
    free(c);
    free(d);

    // system("pause");

    return(0);
}

#endif // MATMUL_1D