#include "tee_pthread_big_matrix_mult_TA.h"
#include <errno.h>
#include <pthread.h>
#include <stdio.h>
#include <tee_internal_api.h>
#include <string.h>

enum { BENCHMARK_MATRIX_MULT };

static TEE_Result test_pthread_benchmark_matrix_mult(unsigned int ncpus);

static inline uint32_t calculate_diff(const TEE_Time *time1,
                                      const TEE_Time *time2)
{
    return (time1->millis - time2->millis) +
        (time1->seconds - time2->seconds) * 1000;
}

TEE_Result TA_CreateEntryPoint(void)
{
    return TEE_SUCCESS;
}

void TA_DestroyEntryPoint(void)
{
}

TEE_Result TA_OpenSessionEntryPoint(uint32_t paramTypes,
                                    TEE_Param params[4],
                                    void **sessionContext)
{
    (void) paramTypes;
    (void) params;
    (void) sessionContext;

    return TEE_SUCCESS;
}

void TA_CloseSessionEntryPoint(void *sessionContext)
{
    (void) sessionContext;
}

TEE_Result TA_InvokeCommandEntryPoint(void *sessionContext,
                                      uint32_t commandID,
                                      uint32_t paramTypes,
                                      TEE_Param params[4])
{
    (void) sessionContext;
    (void) paramTypes;
    (void) params;

    if (TEE_PARAM_TYPE_GET(paramTypes, 0) != TEE_PARAM_TYPE_VALUE_INPUT) {
        return TEE_ERROR_BAD_PARAMETERS;
    }

    unsigned int ncpus = params[0].value.a;

    switch (commandID) {
        case BENCHMARK_MATRIX_MULT:
            return test_pthread_benchmark_matrix_mult(ncpus);

        default:
            printf("Wrong command id %u\n", commandID);
            return TEE_ERROR_BAD_PARAMETERS;
    }

}

#define M       512
#define N       512
#define P       512

int A[M][N];
int B[N][P];
/* Result-matrix */
size_t C[M][P];

static void mat_mult_by_row(void *arg)
{
    size_t row = (size_t) arg;
    size_t col;
    size_t i;

    size_t sum;

    for (col = 0; col < P; ++col) {
        sum = 0;
        for (i = 0; i < N; ++i) {
            sum += A[row][i] * B[i][col];
        }
        C[row][col] = sum;
    }
}

static void mmult_single_threaded(size_t row_start, size_t row_end)
{
    size_t i = row_start;
    for (; i < row_end; ++i) {
        mat_mult_by_row((void *) i);
    }
}

struct thread_arg {
    size_t row_s;
    size_t row_e;
};

static void *mmult_multi_threaded(void *arg)
{
    struct thread_arg *_arg = arg;

    mmult_single_threaded(_arg->row_s, _arg->row_e);
    return NULL;
}

static TEE_Result test_pthread_benchmark_matrix_mult(unsigned int ncpus)
{
    unsigned int i;
    unsigned int threads_num = 0;
    int err = 0;
    TEE_Result res = TEE_SUCCESS;
    const size_t rows_slice = M / ncpus;
    pthread_t th[ncpus];
    struct thread_arg arg;
    TEE_Time start_time;
    TEE_Time end_time;
    uint32_t single_threaded_time_ms;
    uint32_t multi_threaded_time_ms;

    /* Fill A & B matrixes by random values */
    TEE_GenerateRandom(&A[0][0], sizeof(A));
    TEE_GenerateRandom(&B[0][0], sizeof(A));

    TEE_GetSystemTime(&start_time);
    mmult_single_threaded(0, M);
    TEE_GetSystemTime(&end_time);
    single_threaded_time_ms = calculate_diff(&end_time, &start_time);

    TEE_MemFill(&start_time, 0, sizeof(TEE_Time));
    TEE_MemFill(&end_time, 0, sizeof(TEE_Time));

    TEE_GetSystemTime(&start_time);
    for (i = 0; i < ncpus; ++i, ++threads_num) {
        arg.row_s = i * rows_slice;
        arg.row_e = (i + 1) * rows_slice;
        if (i == ncpus - 1) {
            arg.row_e = M;
        }

        err = pthread_create(&th[i], NULL, mmult_multi_threaded, &arg);
        if (err) {
            printf("pthread_create() failed, errno %d, iter %d\n", err, i);
            res = TEE_ERROR_GENERIC;
            break;
        }
    }

    for (i = 0; i < threads_num; ++i) {
        err = pthread_join(th[i], NULL);
        if (err) {
            printf("pthread_join() failed, errno %d, iter %d\n", err, i);
            res |= TEE_ERROR_GENERIC;
        }
    }
    TEE_GetSystemTime(&end_time);

    if (res != TEE_SUCCESS) {
        printf("Test failed\n");
        return res;
    }

    multi_threaded_time_ms = calculate_diff(&end_time, &start_time);

    printf("Profiling results:\n");
    printf("Single thread time, ms: %u\n", single_threaded_time_ms);
    printf("Multiple threads time, ms: %u\n", multi_threaded_time_ms);
    if (multi_threaded_time_ms > single_threaded_time_ms) {
        printf("Test failed\n");
        res = TEE_ERROR_GENERIC;
    } else {
        printf("Test passed\n");
    }

    return res;
}
