#include "test_TA_openmp_example.h"
#include <omp.h>
#include <pthread.h>
#include <stdio.h>
#include <string.h>
#include <tee_internal_api.h>

enum { PRIMES_OPENMP };
enum { MULTIPLIER = 2};
enum { HIGH_LIMIT = 32768UL};

#define COUNT_PRIMES_BTWN_1_AND_N(i,j,n,prime,total)\
    for (i = 2; i <= n; i++) {                      \
        prime = 1;                                  \
        for (j = 2; j < i; j++) {                   \
            if (i % j == 0) {                       \
                prime = 0;                          \
                break;                              \
            }                                       \
        }                                           \
        total = total + prime;                      \
    }

#define NOTIFY(substr)                                          \
    printf("   Number of primes between 1 and %d ("substr"):\n",\
           HIGH_LIMIT);

static TEE_Result openmp_primes(void);
static uint32_t prime_count_mp(uint32_t n);
static uint32_t prime_count_no_mp(uint32_t n);

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;

    switch (commandID) {
    case PRIMES_OPENMP:
        return openmp_primes();

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

}

static TEE_Result count_primes_and_notify(uint32_t(*counter) (uint32_t))
{
    uint32_t n;
    uint32_t primes;
    double wtime;

    printf("         N        Pi          Time\n\n");
    n = 1;
    while (n <= HIGH_LIMIT) {
        wtime = omp_get_wtime();
        primes = counter(n);
        wtime = omp_get_wtime() - wtime;
        printf("  %8d  %8d  %14f\n", n, primes, wtime);
        n = n * MULTIPLIER;
    }

    return TEE_SUCCESS;
}

static TEE_Result openmp_primes(void)
{
    TEE_Result res = TEE_SUCCESS;

    NOTIFY("MP on");
    count_primes_and_notify(prime_count_mp);

    NOTIFY("MP off");
    count_primes_and_notify(prime_count_no_mp);

    return res;
}

static uint32_t prime_count_mp(uint32_t n)
{
    uint32_t i;
    uint32_t j;
    uint32_t prime;
    uint32_t total = 0;

#pragma omp parallel \
    shared ( n )     \
    private ( i, j, prime )

#pragma omp for reduction ( + : total )
    COUNT_PRIMES_BTWN_1_AND_N(i, j, n, prime, total);

    return total;
}

static uint32_t prime_count_no_mp(uint32_t n)
{
    uint32_t i;
    uint32_t j;
    uint32_t prime;
    uint32_t total = 0;

    COUNT_PRIMES_BTWN_1_AND_N(i, j, n, prime, total);

    return total;
}
