#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <malloc.h>
#include <pthread.h>
#include <semaphore.h>
#include <string.h>
#include <sys/time.h>
#include <sched.h>
#include <sys/resource.h>
#include <sys/mman.h>

#define MB	(1024*1024)
#define QUAD_CORE	0

# define handle_error_en(en, msg) \
	do { errno = en; perror(msg); exit(EXIT_FAILURE); } while (0);

unsigned long l_megabytes = 1;
unsigned long l_repeat = 1000;

void SingleFct(void)
{
	unsigned long l_size = l_megabytes*MB;
	unsigned char * l_memBuffer = memalign(32, l_size);
	int i,j;
	unsigned long l_time;
	double l_datarate;
	double t_datarate = 0.0;
	struct timeval  start, end;

	printf("Single Thread - data size %ld MB, runs = %ld\n", l_megabytes, l_repeat);
	for(j=0; j<10; j++)
	{
                gettimeofday(&start, NULL);
		for(i=0; i<l_repeat; i++)
		{
			memset(l_memBuffer, 0x5a, l_size);
		}
                gettimeofday(&end, NULL);

		l_time = (end.tv_sec*1000000 + end.tv_usec) - (start.tv_sec*1000000 + start.tv_usec);
		// printf("end time = %u.%06u, start time = %u.%06u\n", end.tv_sec, end.tv_usec, start.tv_sec, start.tv_usec);
		// printf("l_time = %u\n", l_time);
		l_datarate = (double)l_size*(double)l_repeat/(double)l_time;
		printf("Single Thread Datarate: %f MB/s\n", l_datarate);
		t_datarate += l_datarate;
	}
	printf("Single Thread Average Datarate: %f MB/s\n", t_datarate/10.0f);
	free(l_memBuffer);
}

void Fct(void)
{
	unsigned long l_size = l_megabytes*MB;
	unsigned char * l_memBuffer = memalign(32, l_size);
	unsigned long l_time;
	double l_datarate;
	double t_datarate = 0.0;
	int i,j;
	struct timeval  start, end;

	printf("Thread %u - data size %ld MB, runs = %ld\n", (unsigned int)pthread_self(), l_megabytes, l_repeat);
	for(j=0; j<10; j++)
	{
                gettimeofday(&start, NULL);
		for(i=0; i<l_repeat; i++)
		{
			memset(l_memBuffer, 0x5a, l_size);
		}
                gettimeofday(&end, NULL);
		l_time = (end.tv_sec*1000000 + end.tv_usec) - (start.tv_sec*1000000 + start.tv_usec);
		l_datarate = (double)l_size*(double)l_repeat/(double)l_time;
		printf("Thread :%u: Datarate: %lf MB/s\n", (unsigned int)pthread_self(), l_datarate);
		t_datarate += l_datarate;

	}

	printf("Thread :%u: Average Datarate: %lf MB/s\n", (unsigned int)pthread_self(), t_datarate/10.0f);
	free(l_memBuffer);

	pthread_exit((void *) 0);
}

int main(int argc, char** argv) {

	pthread_t t1, t2, t3, t4;

	pthread_attr_t t1_attr, t2_attr, t3_attr, t4_attr;
	cpu_set_t t1_cpuset, t2_cpuset, t3_cpuset, t4_cpuset;

	int rv;
	int affinity = 0;
	int iarg = 1;

	while (iarg < argc) {
	    if (!strcmp(argv[iarg],"-a1")) {
                affinity = 1;
            }
            else if (!strcmp(argv[iarg],"-a2")) {
                affinity = 2;
            } 
            else if (!strcmp(argv[iarg],"-s1")) {
                l_megabytes = 1;
            }
            else if (!strcmp(argv[iarg],"-s2")) {
                l_megabytes = 10;
            }
            else if (!strcmp(argv[iarg],"-s3")) {
                l_megabytes = 100;
            }
            else {
                printf("Usage: %s [-a1|-a2 ] [-s1|-s2|-s3]\n",argv[0]);
                exit(0);
            }
            iarg++;
	}
	printf("========= Single Thread =========\n\n");

	SingleFct();

	printf("\n\n========= Multi Thread =========\n\n");

	rv = pthread_attr_init(&t1_attr);
	if (rv != 0) handle_error_en(rv,"pthread_attr_init thread #1");

	rv = pthread_attr_init(&t2_attr);
	if (rv != 0) handle_error_en(rv,"pthread_attr_init thread #2");

	if (affinity == 1) 
	{
	    CPU_ZERO(&t1_cpuset);
	    CPU_SET(0,&t1_cpuset);
	    rv = pthread_attr_setaffinity_np(&t1_attr,sizeof(t1_cpuset),&t1_cpuset);
	    if (rv != 0) handle_error_en(rv,"pthread_attr_setaffinity_np thread #1");

	    CPU_ZERO(&t2_cpuset);
	    CPU_SET(1,&t2_cpuset);
	    rv = pthread_attr_setaffinity_np(&t2_attr,sizeof(t2_cpuset),&t2_cpuset);
	    if (rv != 0) handle_error_en(rv,"pthread_attr_setaffinity_np thread #2");
	}
	else if (affinity == 2) 
	{
	    CPU_ZERO(&t1_cpuset);
	    CPU_SET(0,&t1_cpuset);
	    rv = pthread_attr_setaffinity_np(&t1_attr,sizeof(t1_cpuset),&t1_cpuset);
	    if (rv != 0) handle_error_en(rv,"pthread_attr_setaffinity_np thread #1");

	    CPU_ZERO(&t2_cpuset);
	    CPU_SET(0,&t2_cpuset);
	    rv = pthread_attr_setaffinity_np(&t2_attr,sizeof(t2_cpuset),&t2_cpuset);
	    if (rv != 0) handle_error_en(rv,"pthread_attr_setaffinity_np thread #2");
	}

	rv = pthread_create(&t1, &t1_attr, (void *)&Fct, (void *)NULL);
	if (rv != 0) handle_error_en(rv,"pthread_create thread #1");

	rv = pthread_create(&t2, &t2_attr, (void *)&Fct, (void *)NULL);
	if (rv != 0) handle_error_en(rv,"pthread_create thread #2");

#if QUAD_CORE
	rv = pthread_attr_init(&t3_attr);
	if (rv != 0) handle_error_en(rv,"pthread_attr_init thread #3");

	rv = pthread_attr_init(&t4_attr);
	if (rv != 0) handle_error_en(rv,"pthread_attr_init thread #4");

	if (affinity == 1) 
	{
	    CPU_ZERO(&t3_cpuset);
	    CPU_SET(2,&t3_cpuset);
	    rv = pthread_attr_setaffinity_np(&t3_attr,sizeof(t3_cpuset),&t3_cpuset);
	    if (rv != 0) handle_error_en(rv,"pthread_attr_setaffinity_np thread #3");

	    CPU_ZERO(&t4_cpuset);
	    CPU_SET(3,&t4_cpuset);
	    rv = pthread_attr_setaffinity_np(&t4_attr,sizeof(t4_cpuset),&t4_cpuset);
	    if (rv != 0) handle_error_en(rv,"pthread_attr_setaffinity_np thread #4");
	}
	else if (affinity == 2) 
	{
	    CPU_ZERO(&t3_cpuset);
	    CPU_SET(0,&t3_cpuset);
	    rv = pthread_attr_setaffinity_np(&t3_attr,sizeof(t3_cpuset),&t3_cpuset);
	    if (rv != 0) handle_error_en(rv,"pthread_attr_setaffinity_np thread #3");

	    CPU_ZERO(&t4_cpuset);
	    CPU_SET(0,&t4_cpuset);
	    rv = pthread_attr_setaffinity_np(&t4_attr,sizeof(t4_cpuset),&t4_cpuset);
	    if (rv != 0) handle_error_en(rv,"pthread_attr_setaffinity_np thread #4");
	}

	rv = pthread_create(&t3, &t3_attr, (void *)&Fct, (void *)NULL);
	if (rv != 0) handle_error_en(rv,"pthread_create thread #3");

	rv = pthread_create(&t4, &t4_attr, (void *)&Fct, (void *)NULL);
	if (rv != 0) handle_error_en(rv,"pthread_create thread #4");

#endif

    pthread_join(t1, NULL);
    pthread_join(t2, NULL);
#if QUAD_CORE
    pthread_join(t3, NULL);
    pthread_join(t4, NULL);
#endif

	printf("Finished!\n");
	return 0;
}