#ifdef MEMBENCH_LIBUSE
#include "Membench.h"
#endif // MEMBENCH_LIBUSE

#include "csaccess.h"
#include "pmu-cache.h"
#include "prem_counter.h"

#ifdef BAREMETAL
//------Threads-overrides---------------------------------------------------

typedef size_t pthread_t;
typedef size_t pthread_barrier_t;
typedef size_t pthread_attr_t;

static inline void CPU_ZERO(cpu_set_t *set) { set->bits = 0; set->count = 0; }
static inline void CPU_SET(int cpu, cpu_set_t *set) { set->bits |= (1 << cpu); set->count++; }
static inline void CPU_CLR(int cpu, cpu_set_t *set) { set->bits ^= (1 << cpu); set->count--; }
static inline int  CPU_ISSET(int cpu, cpu_set_t *set) { return (set->bits & (1 << cpu)); }
static inline int CPU_COUNT(cpu_set_t *set) { return (int) set->count; }

static inline int pthread_setaffinity_np(pthread_t id, unsigned long affty, cpu_set_t * set) {return 0;}

static inline void pthread_barrier_init(pthread_barrier_t * bar, pthread_attr_t * attr, unsigned count) {}
static inline void pthread_barrier_wait(pthread_barrier_t * bar) {}
static inline void pthread_barrier_destroy(pthread_barrier_t * bar) {}
static inline void pthread_join(pthread_t id, void * smth){}

static int pthread_create(pthread_t *thread, const pthread_attr_t *attr,
                          void *(*start_routine) (void *), void *arg)
{
	start_routine(arg);
	return 0;
}
//-END:-Threads-overrides---------------------------------------------------

#ifdef JAILHOUSE
#include <inmate.h>

#define CMDLINE_BUFFER_SIZE	256
CMDLINE_BUFFER(CMDLINE_BUFFER_SIZE);

//uintstd in jailhouse way
#define uint32_t u32
#define uint64_t u64


#define POLLUTE_CACHE_SIZE	(512 * 1024)

#ifdef CONFIG_UART_OXPCIE952
#define UART_BASE		0xe010
#else
#define UART_BASE		0x3f8
#endif
#define UART_LSR		0x5
#define UART_LSR_THRE		0x20
#define UART_IDLE_LOOPS		100


#define perror(FUNC) (printk(#FUNC))

#define printf printk

#define exit(SIG) { printk("exit with %d\n", SIG); asm volatile("hlt");}

//-----Time-and-and-randomization-overrides-------------------------------
static inline unsigned long time(unsigned long * seconds)
{
	 u64 time_sec = tsc_read();
	 if (seconds != NULL){
	 	(*seconds) = time_sec;
	 }
	 return time_sec;
}

/*
*	Tables of Maximally-Equidistributed Combined Lfsr Generators (1998)
*	by Pierre L'Ecuyer
*	taken from: http://stackoverflow.com/questions/1167253/implementation-of-rand
*/
static unsigned int z1 = 12345, z2 = 12345, z3 = 12345, z4 = 12345;

static unsigned int lfsr113_Bits (void)
{
   unsigned int b;
   b  = ((z1 << 6) ^ z1) >> 13;
   z1 = ((z1 & 4294967294U) << 18) ^ b;
   b  = ((z2 << 2) ^ z2) >> 27;
   z2 = ((z2 & 4294967288U) << 2) ^ b;
   b  = ((z3 << 13) ^ z3) >> 21;
   z3 = ((z3 & 4294967280U) << 7) ^ b;
   b  = ((z4 << 3) ^ z4) >> 12;
   z4 = ((z4 & 4294967168U) << 13) ^ b;
   return (z1 ^ z2 ^ z3 ^ z4);
}

static unsigned int rand(void)
{
	return lfsr113_Bits();
}

static void srand(unsigned int seed)
{
   //seed into z1 z2 z3 z4
   z1 = z2 = z3 = z4 = seed;
}
//-END:-Time-and-and-randomization-overrides-------------------------------

//-----Assertion-overrides-------------------------------------------------
#ifdef NDEBUG
# define assert(EX)
#else
# define assert(EX) (void)((EX) || (__assert (#EX, __FILE__, __LINE__),0))
#endif

static inline void __assert(const char *msg, const char *file, int line)
{
	printk("Assertion %s in %s:%s failed.\n", msg, file, line);
	exit(-1);
}
//-END:-Assertion-overrides-------------------------------------------------

//just remove fflush
#define fflush (void)sizeof
#define stdout 0
#define stdin 0
#define stderr 0
#endif // JAILHOUSE

#else

#endif // BAREMETAL

#define STRINGIFY(val) #val
#define TOSTRING(val) STRINGIFY(val)
#define LOC __FILE__ ":" TOSTRING(__LINE__) ": "

#define CHECK(cmd) ({ int ret = (cmd); if (ret == -1) { perror(LOC #cmd); exit(1); }; ret; })
#define CHECKPTR(cmd) ({ void *ptr = (cmd); if (ptr == (void*)-1) { perror(LOC #cmd); exit(1); }; ptr; })
#define CHECKNULL(cmd) ({ typeof(cmd) ptr = (cmd); if (ptr == NULL) { perror(LOC #cmd); exit(1); }; ptr; })
//#define CHECKFGETS(s, size, stream) ({ void *ptr = fgets(s, size, stream); if (ptr == NULL) { if (feof(stream)) fprintf(stderr, LOC "fgets(" #s "): Unexpected end of stream\n"); else perror(LOC "fgets(" #s ")"); exit(1); }; ptr; })
#define CHECKTRUE(bool, msg) ({ if (!(bool)) { printf("Error: " msg "\n"); exit(1); }; })

struct s {
        struct s *ptr;
        uint32_t dummy[(64 - sizeof(struct s*))/sizeof(uint32_t)];
};

static_assert(sizeof(struct s) == 64, "Struct size differs from cacheline size");

#ifdef BAREMETAL
#define MAX_CPUS 1
#else
#define MAX_CPUS 4
#endif

#ifdef __aarch64__
#define MRS32(reg) ({ uint32_t v; asm volatile ("mrs %0," # reg : "=r" (v)); v; })
#define MRS64(reg) ({ uint64_t v; asm volatile ("mrs %0," # reg : "=r" (v)); v; })

#define MSR(reg, v) ({ asm volatile ("msr " # reg ",%0" :: "r" (v)); })

static void ccntr_init(void)
{
	MSR(PMCNTENSET_EL0, 0x80000000);
	MSR(PMCR_EL0, MRS32(PMCR_EL0) | 1);
}

static uint64_t ccntr_get(void)
{
	return MRS64(PMCCNTR_EL0);
}

#else //!aarch64

#ifdef JAILHOUSE

unsigned long tsc_freq;
static void ccntr_init(void)
{
	tsc_freq = tsc_init();
}

#else //!JAILHOUSE

static void ccntr_init(void){}

#endif

static uint64_t ccntr_get(void)
{
//taken from lib/timing.c
#ifdef __x86_64__
	uint64_t lo, hi;

	asm volatile("rdtsc" : "=a" (lo), "=d" (hi));
	return (uint64_t)lo | (((uint64_t)hi) << 32);
#elseif __x86__ //x86
	uint64_t v;

	asm volatile("rdtsc" : "=A" (v));
	return v;
#endif

}
#endif

static uint64_t get_time(struct cfg *cfg)
{
	if (cfg->use_cycles) {
#ifdef __ARM_ARCH
		return ccntr_get();
#elif JAILHOUSE
	 	return tsc_read();
#elif __x86_64__ //!BAREMETAL
        struct timespec t;

        clock_gettime(CLOCK_MONOTONIC, &t);
        return (uint64_t)t.tv_sec * 1000000000 + t.tv_nsec;
#endif

	} else {
		//struct timeval tval;
		//gettimeofday(&tval, NULL);
		//return (long int)tval.tv_sec * 1000000 + (long int)tval.tv_usec;
		return ccntr_get() / 1200;
	}
}

#define ARRAY_SIZE(x) (sizeof(x) / sizeof(x[0]))

static void prepare(struct s *array, unsigned size, bool sequential)
{
	int i, j;
	int count = size / sizeof(struct s);

	if (sequential) {
		for (i = 0; i < count - 1; i++){
	 		array[i].ptr = &array[i+1];
		}
	 	array[count - 1].ptr = &array[0];
	} else {
		memset(array, 0, size);
		struct s *p = &array[0];
		for (i = 0; i < count - 1; i++) {
			p->ptr = (struct s*)1; /* Mark as occupied to avoid self-loop */
			for (j = rand() % count;
			     array[j].ptr != NULL;
			     j = (j >= count) ? 0 : j+1);
			p = p->ptr = &array[j];
		}
		p->ptr = &array[0];
	}
}

static void do_read(struct s *array, unsigned reads)
{
	unsigned i = reads / 32;
	volatile struct s *p = &array[0];
	while (--i) {
		p = p->ptr;	/* 0 */
		p = p->ptr;	/* 1 */
		p = p->ptr;	/* 2 */
		p = p->ptr;	/* 3 */
		p = p->ptr;	/* 4 */
		p = p->ptr;	/* 5 */
		p = p->ptr;	/* 6 */
		p = p->ptr;	/* 7 */
		p = p->ptr;	/* 8 */
		p = p->ptr;	/* 9 */
		p = p->ptr;	/* 10 */
		p = p->ptr;	/* 11 */
		p = p->ptr;	/* 12 */
		p = p->ptr;	/* 13 */
		p = p->ptr;	/* 14 */
		p = p->ptr;	/* 15 */
		p = p->ptr;	/* 16 */
		p = p->ptr;	/* 17 */
		p = p->ptr;	/* 18 */
		p = p->ptr;	/* 19 */
		p = p->ptr;	/* 20 */
		p = p->ptr;	/* 21 */
		p = p->ptr;	/* 22 */
		p = p->ptr;	/* 23 */
		p = p->ptr;	/* 24 */
		p = p->ptr;	/* 25 */
		p = p->ptr;	/* 26 */
		p = p->ptr;	/* 27 */
		p = p->ptr;	/* 28 */
		p = p->ptr;	/* 29 */
		p = p->ptr;	/* 30 */
		p = p->ptr;	/* 31 */
	}
}

static void do_write(struct s *array, unsigned accesses, unsigned ofs)
{
	unsigned i = accesses / 32;
	volatile struct s *p = &array[0];
	while (--i) {
		p->dummy[ofs]++; p = p->ptr; /* 0 */
		p->dummy[ofs]++; p = p->ptr; /* 1 */
		p->dummy[ofs]++; p = p->ptr; /* 2 */
		p->dummy[ofs]++; p = p->ptr; /* 3 */
		p->dummy[ofs]++; p = p->ptr; /* 4 */
		p->dummy[ofs]++; p = p->ptr; /* 5 */
		p->dummy[ofs]++; p = p->ptr; /* 6 */
		p->dummy[ofs]++; p = p->ptr; /* 7 */
		p->dummy[ofs]++; p = p->ptr; /* 8 */
		p->dummy[ofs]++; p = p->ptr; /* 9 */
		p->dummy[ofs]++; p = p->ptr; /* 10 */
		p->dummy[ofs]++; p = p->ptr; /* 11 */
		p->dummy[ofs]++; p = p->ptr; /* 12 */
		p->dummy[ofs]++; p = p->ptr; /* 13 */
		p->dummy[ofs]++; p = p->ptr; /* 14 */
		p->dummy[ofs]++; p = p->ptr; /* 15 */
		p->dummy[ofs]++; p = p->ptr; /* 16 */
		p->dummy[ofs]++; p = p->ptr; /* 17 */
		p->dummy[ofs]++; p = p->ptr; /* 18 */
		p->dummy[ofs]++; p = p->ptr; /* 19 */
		p->dummy[ofs]++; p = p->ptr; /* 20 */
		p->dummy[ofs]++; p = p->ptr; /* 21 */
		p->dummy[ofs]++; p = p->ptr; /* 22 */
		p->dummy[ofs]++; p = p->ptr; /* 23 */
		p->dummy[ofs]++; p = p->ptr; /* 24 */
		p->dummy[ofs]++; p = p->ptr; /* 25 */
		p->dummy[ofs]++; p = p->ptr; /* 26 */
		p->dummy[ofs]++; p = p->ptr; /* 27 */
		p->dummy[ofs]++; p = p->ptr; /* 28 */
		p->dummy[ofs]++; p = p->ptr; /* 29 */
		p->dummy[ofs]++; p = p->ptr; /* 30 */
		p->dummy[ofs]++; p = p->ptr; /* 31 */
	}
}

struct benchmark_thread {
	pthread_t id;
	unsigned cpu;
	uint64_t result_integral;
	uint64_t result_fractional;
	uint32_t cache_misses_pmu;
	uint32_t cache_misses_pl;
	struct cfg *cfg;
};

pthread_barrier_t barrier;
#ifdef BAREMETAL
struct s array[MAX_CPUS][MAX_WSS/sizeof(struct s)] __attribute__ ((aligned (2*1024*1024))) __attribute__ ((section (".bench-array")));
#else //!BAREMETAL
struct s array[MAX_CPUS][64*0x100000/sizeof(struct s)] __attribute__ ((aligned (2*1024*1024)));
#endif
bool print = true;

static void *benchmark_thread(void *arg)
{
	struct benchmark_thread *me = arg;
	cpu_set_t set;

	CPU_ZERO(&set);
 	CPU_SET(me->cpu, &set);

	if (pthread_setaffinity_np(me->id, sizeof(set), &set) != 0) {
		perror("pthread_setaffinity_np");
		exit(1);
	}

 	prepare(array[me->cpu], me->cfg->size, me->cfg->sequential);

 	pthread_barrier_wait(&barrier);

 	if (print)
 		printf("CPU %d starts measurement\n", me->cpu);

 	uint64_t tic, tac;
 	tic = get_time(me->cfg);

 	if (me->cfg->cache_miss) {
 		pmu_cachemiss_reset(me->cpu);
 	}

 	if (me->cfg->cache_miss_pl) {
 		_deprecated_pl_count_reset(me->cpu);
 	}

 	if (me->cfg->write == false)
 		do_read(array[me->cpu], me->cfg->read_count);
 	else
 		do_write(array[me->cpu], me->cfg->read_count, me->cfg->ofs);

 	if (me->cfg->cache_miss) {
 		me->cache_misses_pmu = pmu_cachemiss_read(me->cpu);
 	}

 	if (me->cfg->cache_miss_pl) {
 		me->cache_misses_pl = _deprecated_pl_count_read(me->cpu) * ETM_EVENTS_BUFF_NUM;
 	}

 	tac = get_time(me->cfg);
 	me->result_integral = (tac - tic) / me->cfg->read_count;
 	me->result_fractional = (((tac - tic) * 1000) / me->cfg->read_count) % 1000;
 //	printf("(tic=%lu,tac=%lu) ", tic, tac);
 	return NULL;
}

void run_benchmark(struct cfg *cfg)
{
	struct benchmark_thread thread[MAX_CPUS];
	unsigned i;
	cpu_set_t set = cfg->cpu_set;
	pthread_barrier_init(&barrier, NULL, cfg->num_threads);
	for (i = 0; i < cfg->num_threads; i++) {
		thread[i].cfg = cfg;
		if (CPU_COUNT(&set) == 0) {
			thread[i].cpu = i;
		} else {
			int j;
			for (j = 0; j < MAX_CPUS; j++) {
				if (CPU_ISSET(j, &set)) {
					thread[i].cpu = j;
					CPU_CLR(j, &set);
					break;
				}
			}
		}
		if (print)
			printf( "Running thread %d on CPU %d\n", i, thread[i].cpu);
		pthread_create(&thread[i].id, NULL, benchmark_thread, &thread[i]);
	}

	for (i = 0; i < cfg->num_threads; i++) {
		pthread_join(thread[i].id, NULL);
	}
	pthread_barrier_destroy(&barrier);

	printf("%d", cfg->size);
	for (i = 0; i < cfg->num_threads; i++) {
		//NOTE: Jailhouse is not able to print doubles.
		//printf("\t%#.3g", thread[i].result);
		printf("\t%lu.%03lu", thread[i].result_integral, thread[i].result_fractional);
		if (thread[i].cfg->cache_miss) {
			printf("\t%u", thread[i].cache_misses_pmu);
		}
		if (thread[i].cfg->cache_miss_pl) {
			printf("\t%u", thread[i].cache_misses_pl);
		}
	}
	printf("\n");
	fflush(stdout);
	print = false;
}

void do_benchmark(struct cfg * cfg)
{

	if (cfg->write) {
		struct s s;
		assert(cfg->ofs < ARRAY_SIZE(s.dummy));
	}

	if (cfg->cache_miss) {
			for (int i = 0; i < MAX_CPUS; i++) {
				pmu_init(i);
			}
		}
	if (cfg->cache_miss_pl) {
		cs_prem_count_init();
		for (uint32_t i = 0; i < MAX_CPUS; i++) {
			cs_prem_count_percpu_init(i);
		}
	}
#ifdef __ARM_ARCH
	ccntr_init();
	srand(time(NULL));
#endif

	for (;cfg->repeats > 0; cfg->repeats--) {
		if (cfg->size != 0) {
			run_benchmark(cfg);
		} else {
			 unsigned order, size, step;
			 for (order = 10; order <= 29; order++) {
				for (step = 0; step < 2; step++) {
					size = 1 << order;
					if (step == 1)
						size += size / 2;

					cfg->size = size;
			 		run_benchmark(cfg);
				}
			}
			cfg->size = 0;
		}
	}

}

#ifndef MEMBENCH_LIBUSE

#ifdef JAILHOUSE
const void * image_end;
void inmate_main(void)
#else //Linux
int main(int argc, char *argv[])
#endif
{
	struct cfg cfg = {
		.sequential = true,
		.num_threads = 1,
		.size = 0,
		.read_count = 0x2000000,
		.write = false,
		.ofs = 0,
		.use_cycles = false, // i.e. use nanoseconds /
		.repeats = 1
	};
#ifdef JAILHOUSE
	//initialize UART
	unsigned int n;
	printk_uart_base = UART_BASE;
	do {
		for (n = 0; n < UART_IDLE_LOOPS; n++)
			if (!(inb(UART_BASE + UART_LSR) & UART_LSR_THRE))
				break;
	} while (n < UART_IDLE_LOOPS);

	printk("cmdline opts: '%s'\n", cmdline);
	comm_region->pm_timer_address = 0x408;

	//parse cmdline
	cfg.read_count = cmdline_parse_int("-c", cfg.read_count);
	cfg.repeats = cmdline_parse_int("-R", cfg.repeats);
	cfg.ofs = cmdline_parse_int("-o", cfg.ofs);
	cfg.sequential = !cmdline_parse_bool("-r");
	cfg.size = cmdline_parse_int("-s", cfg.size);
	if (cmdline_parse_bool("-t")) {
		printk("Threads are not supported. '-t' was ignored.\n");
	}
	if (cmdline_parse_bool("-C")) {
		printk("CPU selection is not supported. '-C' was ignored.\n");
	}
	cfg.write = cmdline_parse_bool("-w");
	cfg.use_cycles = cmdline_parse_bool("-y");
	//initialize timing
	ccntr_init();
	printk("Calibrated TSC frequency: %lu.%03u kHz\n", tsc_freq / 1000,
	       tsc_freq % 1000);

	u8 * start_memreg = (u8 *) array;
	u64 end_memreg_addr = (u64) &image_end;
	while ( ((u64) start_memreg )< end_memreg_addr) {
		//printk("%p\n",start_memreg);
		map_range(start_memreg, HUGE_PAGE_SIZE, MAP_CACHED);
		start_memreg += HUGE_PAGE_SIZE;

	}

#else //Linux param's parsing
	CPU_ZERO(&cfg.cpu_set);

	int opt;
	while ((opt = getopt(argc, argv, "c:C:o:rs:t:wyR:")) != -1) {
		switch (opt) {
		case 'c':
			cfg.read_count = atol(optarg);
			break;
		case 'o':
			cfg.ofs = atol(optarg);
			break;
		case 'r':	// random //
			cfg.sequential = false;
			break;
		case 's':
			cfg.size = atol(optarg);
			assert(cfg.size <= sizeof(array[0]));
			break;
		case 't':
			cfg.num_threads = atol(optarg);
			break;
		case 'C':
			CPU_SET(atol(optarg), &cfg.cpu_set);
			break;
		case 'w':
			cfg.write = true;
			break;
		case 'y':
			cfg.use_cycles = true;
			break;
		case 'R':
			cfg.repeats = atol(optarg);
			break;
		default: // '?' //
			fprintf(stderr, "Usage: %s ... TODO\n", argv[0]);
			exit(1);
		}
	}
#endif

	srand(time(NULL));

	if (cfg.write) {
		struct s s;
		assert(cfg.ofs < ARRAY_SIZE(s.dummy));
	}

	if (cfg.use_cycles)
		ccntr_init();

	for (;cfg.repeats > 0; cfg.repeats--) {
		if (cfg.size != 0) {
			run_benchmark(&cfg);
		} else {
			 unsigned order, size, step;
			 for (order = 10; order <= 24; order++) {
				for (step = 0; step < 2; step++) {
					size = 1 << order;
					if (step == 1)
						size += size / 2;

					cfg.size = size;
			 		run_benchmark(&cfg);
				}
			}
			cfg.size = 0;
		}
	}

// unsigned long tic, tac, i;
// i = 1;
// while (comm_region->msg_to_cell != JAILHOUSE_MSG_SHUTDOWN_REQUEST)
// {
// 	tic = ccntr_get();
// 	delay_us(1000000*i);
// 	i *= 2;
// 	tac = ccntr_get();
// 	printf("tic=%lu, tac=%lu, diff=%lu, time=%lu.%03lu\n", tic, tac, tac-tic, (tac-tic) / tsc_freq, ((tac-tic) * 1000 / tsc_freq) % 1000);
// }
#ifdef JAILHOUSE
	comm_region->cell_state = JAILHOUSE_CELL_SHUT_DOWN;
	printk("done!\n");
#else //!JAILHOUSE
	return 0;
#endif
}
#endif // MEMBENCH_LIBUSE
