/* vaguely derived from code vaguely derived by Pathscale from:
 * Derived from code supplied by Ohio State University:
 * Copyright (C) 2002-2003 the Network-Based Computing Laboratory
 * (NBCL), The Ohio State University.  
 *  http://nowlab.cis.ohio-state.edu/projects/mpi-iba/
 */

#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <unistd.h>
#include <sys/time.h>
#include <values.h>
#include <stdarg.h>
#include <string.h>
#include <errno.h>
#include <sched.h>

#define BUFSIZE (4*1024*1024)
char buf[BUFSIZE];
FILE *fp;

double gtod() {
    struct timeval tv;
    gettimeofday(&tv,0);
    return tv.tv_sec + 1e-6 * tv.tv_usec;
}
typedef unsigned long long u64;
#define serialize() \
__asm__ __volatile__("cpuid"::"a"(0):"eax","ebx","ecx","edx")

u64 rdtsc() {
    u64 tsc;
    serialize();
    __asm__ __volatile__("rdtsc" : "=A" (tsc));
    serialize();
    return tsc;
}

//#define DEBUG
#ifdef DEBUG
# define debug if (0) {} else debug_printf
#else
# define debug if (1) {} else debug_printf
#endif

void debug_printf(char *format, ...) {
    va_list ap;
    va_start(ap,format);

    vfprintf(stderr,format,ap);
    if (format[strlen(format)-1] != '\n') fputc('\n',stderr);

    va_end(ap);
}

void mysleep(double delay) {
    double sec = floor(delay);
    double usec = 1e6 * delay - sec;
    struct timeval tv = {(unsigned)sec,(unsigned)usec};
    select(0,0,0,0,&tv);
}

void yield() {
    struct timeval tv = {0,10000};
    select(0,0,0,0,&tv);
}

class Stats {
    unsigned n;
    double sum,min,max;
    float data[1000];

  public:
    Stats(): n(0), sum(0), min(FLT_MAX), max(FLT_MIN) {}
    void put(double v) {
	if (v > max) max = v;
	if (v < min) min = v;
	sum += v;
	if (n < (sizeof(data)/sizeof(data[0])))
	    data[n] = v;
	n++;
	
    }
    unsigned get_n() { return n; }
    double get_min() { if (!n) return 0; return min; }
    double get_max() { if (!n) return 0; return max; }
    double get_avg() { if (!n) return 0; return sum/n; }
    double dump(char *fname) { 
	FILE *fp = fopen(fname,"w");
	for (int i=0; i<n; i++)
	    fprintf(fp,"%f\n",data[i]);
	fclose(fp);
    }

    void spew() {
	fprintf(fp,"n=%u, min=%e, avg=%e, max=%e\n", n, min, get_avg(), max);
    }
};

const unsigned tag_name = 42;
const unsigned tag_size = 43;
const unsigned tag_data = 1;

const double warm_time = .001;
const double bench_time = .020;
const int min_bench_times = 60;
const int max_bench_times = 200;
double clock;
double ticks_to_secs;

int rank, nprocs;
char name[100];
int name_len;
char done[10000];

int name_to_node(char *name) {
    int peer_node;
    sscanf(name,"%*[a-z]%d",&peer_node);
    return peer_node;
}

double do_tx(int peer, int csizes, int sizes[]) {
    MPI_Status stat;
    int isize,bufsize;

    Stats *stats = new Stats[csizes];

    done[name_to_node(name)] = 1;

    char peer_name[1000];
    int peer_name_len = sizeof(peer_name);
    MPI_Recv(peer_name, peer_name_len, MPI_CHAR, peer, tag_name, MPI_COMM_WORLD, &stat);
    int peer_node = name_to_node(peer_name);

    debug("send(%s) -> %s (%d) done(%d)\n",name,peer_name,peer_node,done[peer_node]);

    if (!done[peer_node]) {
//	done[peer_node] = 1;
	for (isize=0; isize<csizes; isize++) {

	    Stats &st = stats[isize];
	    bufsize = sizes[isize];

	    debug("%s tx(%u) -> %s",name,bufsize,peer_name);

	    MPI_Send(&bufsize, 1, MPI_INT, peer, tag_size, MPI_COMM_WORLD);

	    buf[0] = 1;

	    // do a few warm-ups
	    for (int times = 0; times < 10; times++) {
		MPI_Send(buf, bufsize, MPI_INT, peer, tag_data, MPI_COMM_WORLD);
		MPI_Recv(buf, bufsize, MPI_INT, peer, tag_data, MPI_COMM_WORLD, &stat);
	    }

	    double start = gtod();
	    double end = start + bench_time;

	    while (1) {
		u64 a = rdtsc();
		MPI_Send(buf, bufsize, MPI_INT, peer, tag_data, MPI_COMM_WORLD);
		MPI_Recv(buf, bufsize, MPI_INT, peer, tag_data, MPI_COMM_WORLD, &stat);
		u64 b = rdtsc();

		if (b <= a) continue;	// just redo if the TSC wrapped
		st.put(ticks_to_secs * (b - a));

		// continue until either enough reps or enough time and reps
		if (st.get_n() > max_bench_times ||
		    (st.get_n() > min_bench_times && gtod() > end))
		    break;
	    }
	    debug("%s tx done (%u times, avg %.3f); sending end\n",
		  name,st.get_n(),1e6*st.get_avg());
	    buf[0] = 0;
	    MPI_Send(buf, bufsize, MPI_INT, peer, tag_data, MPI_COMM_WORLD);
	}
	fprintf(fp,"%d %d ",rank,peer);
	fprintf(fp,"%s %s ",name,peer_name);
	for (isize=0; isize<csizes; isize++) {
	    fprintf(fp,"%u %.3f/%.3f/%.3f ",
		   4*sizes[isize],
		   stats[isize].get_min() * 1e6,
		   stats[isize].get_avg() * 1e6,
		   stats[isize].get_max() * 1e6);
	}
	fprintf(fp,"\n");
    }
    bufsize = 0;
    MPI_Send(&bufsize, 1, MPI_INT, peer, tag_size, MPI_COMM_WORLD);
}

double do_rx(int peer) {
    MPI_Status stat;

    debug("%s rx announcing name\n",name);
    MPI_Send(name, 1+strlen(name), MPI_CHAR, peer, tag_name, MPI_COMM_WORLD);

    while (1) {
	int bufsize;
	MPI_Recv(&bufsize, 1, MPI_INT, peer, tag_size, MPI_COMM_WORLD, &stat);
	debug("%s rx got size %u\n",name,bufsize);
	if (!bufsize) break;

	while (1) {
	    MPI_Recv(buf, bufsize, MPI_INT, peer, tag_data, MPI_COMM_WORLD, &stat);
	    if (!buf[0]) break;
	    MPI_Send(buf, bufsize, MPI_INT, peer, tag_data, MPI_COMM_WORLD);
	}
    }
}

void calibrate_ticks() {
    int times = 0;
    double sum = 0;
    while (1) {
	double start = gtod();
	u64 a = rdtsc();
	mysleep(.050 + .3 * drand48());
	u64 b = rdtsc();
	double elapsed = gtod() - start;
	if (b <= a) continue;
	double tick_rate = (b-a) / elapsed;
	sum += tick_rate;
	times++;
//	printf("%8.6f GHz (%8.6f)\n",tick_rate*1e-9,(sum/times)*1e-9);
	if (times >= 10)
	    break;
    }
    clock = sum / times;
    ticks_to_secs = 1.0 / clock;
    debug("final %5.3f GHz",clock*1e-9);
}

void usage() {
    fprintf(stderr,"checker [-s size]\n");
    exit(0);
}

int main(int argc, char *argv[]) {
    int src,dst;
    char *oname = 0;

    // message sizes in 4B units.
    int sizes[10] = {1};
    int csizes = 1;

    int letter;
    while ((letter = getopt(argc,argv,"o:s:")) != -1) {
        switch(letter) {
	case 'o':
	    oname = optarg;
	    break;
        case 's':
	    sizes[csizes++] = atoi(optarg);
            break;
        default:
            usage();
        }
    }
    if (csizes == 1)
	sizes[csizes++] = 4096;

    for (int i=0; i<csizes; i++)
	debug("sizes(%d): %d\n",i,sizes[i]);

    fp = stdout;
    if (oname) {
	fp = fopen(oname,"w");
	if (!fp) {
	    debug_printf("failed to open(%s): %s",oname,strerror(errno));
	    exit(1);
	}
    }
    unsigned ntests = 0;

    double start = gtod();

    char *nn = getenv("SLURM_NNODES");
    char *np = getenv("SLURM_NPROCS");
//    printf("nn=%s, np=%s\n",nn,np);
//    printf("nn=%s, np=%s\n",nn,np);
    if (nn && np && !strcmp(nn,np)) {
	cpu_set_t set;
	int i;

//	printf("setting affinity to first cpu\n");

	if (1) {
	    sched_getaffinity(0,sizeof(set),&set);
//	    printf("original cpu affinity:");
	    for (i=0; i<sizeof(set)*8; i++)
		if (__CPU_ISSET(i,&set)) printf(" %u",i);
//	    printf("\n");
	}

	__CPU_ZERO(&set);
	__CPU_SET(0,&set);
	sched_setaffinity(0,sizeof(set),&set);
    }
    calibrate_ticks();

    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    if (rank == 0)
	printf("nodelist %s\n",getenv("SLURM_NODELIST"));

    name_len = sizeof(name);
    MPI_Get_processor_name(name, &name_len);

    if (rank == 0)
	fprintf(fp,"# sender receiver [bytes rtt-seconds]+\n");

    for (src = 0; src < nprocs; src++) {
	// to do only src->dst sends, s/nprocs/src/ here:
//	for (dst = 0; dst < nprocs; dst++) {
	for (dst = 0; dst < src; dst++) {

	    if (src == dst) continue;
	    MPI_Barrier(MPI_COMM_WORLD);

	    if (rank == src) do_tx(dst,csizes,sizes);
	    if (rank == dst) do_rx(src);
	    ntests++;
	}
    }
    MPI_Finalize();

    if (rank == 0) {
        double elapsed = gtod() - start;
	fprintf(fp,"# %u tests in %.3f seconds (%.1f)\n",
	       ntests,elapsed,ntests/elapsed);
    }

    return 0;
}