#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <pvm3.h>
#include <math.h>

#include "main.h"
#include "image.h"

/* variable instantiation */
char *configFilename;

int maxiter;
float xmin, xmax, 
      ymin, ymax;

int *tids;
SlaveStat *slaveStats;
Compix *slaveComp;
Utilization *slaveUtil;
int tids_len;

int xres, yres;
unsigned char *image;
unsigned char *cmap;


Config *
newConfig()
{
  Config *conf;

  conf = (Config *) malloc(sizeof(Config));
  conf->slaves = 0;
  conf->spawn_option = 0;
  conf->where = (char *) malloc(sizeof(char)*WHERE_LEN);
  conf->where[0] = 0;
  conf->where_len = WHERE_LEN;

  return conf;
}

void
deleteConfig(Config *conf)
{
  if (conf != NULL) {
    if (conf->where != NULL) {
      free(conf->where);
    }

    free(conf);
  }
}

Config *
nextConfig(FILE *confFile, int *rc)
{
  Config *conf;
  int scanResult;
  int slaves;
  char spawn_option[50];
  char where[WHERE_LEN];

  conf = newConfig(); 
  *rc = 0;
  scanResult = fscanf(confFile, "%d%s%s", &slaves, spawn_option, where);
  switch (scanResult) {
  case 3:
    conf->slaves = slaves;
    strcpy(conf->where, where);

    if (strcmp(spawn_option, "HOST") == 0) {
      conf->spawn_option = PvmTaskHost;
    }
    else if (strcmp(spawn_option, "ARCH") == 0) {
      conf->spawn_option = PvmTaskArch;
    }
    else {
      fprintf(stderr, "Unknown spawn option: %s\n", spawn_option);
      *rc = 1;
    }
    break;

  case 2:
    conf->slaves = slaves;

    if (strcmp(spawn_option, "DEFAULT") == 0) {
      conf->spawn_option = PvmTaskDefault;
    }
    else {
      fprintf(stderr, "Unknown spawn option: %s\n", spawn_option);
      *rc = 1;
    }
    break;
    
  default:
    if (!feof(confFile)) {
      fprintf(stderr, "Error parsing config file\n");
      *rc = 1;
    }
    deleteConfig(conf);
    conf=NULL;
  }

  if (*rc != 0) {
    deleteConfig(conf);
    conf=NULL;
  }

  return conf;
}

int
spawnSlaves(FILE *confFile)
{
  Config *conf;
  int spawned = 0;
  int shouldSpawn = 0;
  int rc = 0;
  char maxit_a[20];
  char *args[2];

  sprintf(maxit_a, "%d", maxiter);
  args[0] = maxit_a;
  args[1] = 0;

  do {
    conf = nextConfig(confFile, &rc);
    if (conf != NULL) {
      printf("Spawning %d slave(s) of type %d on %s\n", conf->slaves, 
	     conf->spawn_option, conf->where);
      
      shouldSpawn += conf->slaves;
      
      spawned += pvm_spawn(SLAVENAME, args, conf->spawn_option, conf->where,
			   conf->slaves, tids+spawned);
      tids_len = spawned;

      if (spawned != shouldSpawn) {
	rc = 1;
      }
    }
    deleteConfig(conf);
  } while(conf != NULL && rc == 0);

  return rc;
}

void
killSlaves()
{
  int i;
  int rc;

  for(i=0; i<tids_len; i++) {
    if ((rc = pvm_kill(tids[i])) != 0) {
      if (rc == PvmBadParam) {
	fprintf(stderr, "Tried to kill invalid tid %d\n", tids[i]);
      }
      else {
	fprintf(stderr, "Error while killing task %d\n", tids[i]);
      }
    }
  }
}

int
nextIdleSlave(int pos)
{
  int i = pos;

  do {
    i++;

    if (i >= tids_len) {
      i = 0;
    }
  } while (slaveStats[i] != SSIdle && i != pos);

  /* no free slave found */
  if (slaveStats[i] != SSIdle) {
    i = -1;
  }

  return i;
}

int
findSlave(int tid)
{
  int i;

  for (i=0; i<tids_len; i++) {
    if (tids[i] == tid) {
      return i;
    }
  }

  return -1; /* not found */
}

int
sendPix(int slave, Compix c)
{
  int cmd;
  int rc;

  cmd = CmdPixel;
  
  pvm_initsend(PvmDataDefault);
  pvm_pkint(&cmd, 1, 1);
  pvm_pkfloat((float *)&c.c.real, 1, 1);
  pvm_pkfloat((float *)&c.c.imag, 1, 1);
  rc = pvm_send(tids[slave], MSGTAG);
  
  if (rc == 0) {
    slaveStats[slave] = SSBusy;
    slaveComp[slave] = c;
  }

  return rc;
}

int
receivePix()
{
  int bufid, iters, tid, slave;

  bufid = pvm_recv(-1, MSGTAG);
  pvm_bufinfo(bufid, NULL, NULL, &tid);
  pvm_upkint(&iters, 1, 1);
  /*fprintf(stderr, "Received iterations %d\n", iters);*/
  
  slave = findSlave(tid);
  slaveUtil[slave].pixels++;

  /* WRITE PIXEL(iters) IN IMAGE(slaveComp[slave].p) */
  /* image[slaveComp[slave].p.x][slaveComp[slave].p.y] = iters*255/maxiter; */

  setPixWithCmap(slaveComp[slave].p, iters);

  /* fprintf(stderr, "Wrote pixel (%d,%d) (%d,%d,%d)\n", slaveComp[slave].p.x,
slaveComp[slave].p.y, r, g, b); */

  slaveStats[slave] = SSIdle;

  return slave;
}

int
sendQuit(int slave)
{
  int cmd;
  int rc;
  
  cmd = CmdQuit;
  pvm_initsend(PvmDataDefault);
  pvm_pkint(&cmd, 1, 1);
  rc = pvm_send(tids[slave], MSGTAG);

  return rc;
}

int
receiveStat()
{
  int bufid, tid, slave;
  long idle_sec, busy_sec;

  bufid = pvm_recv(-1, MSGTAG);
  pvm_bufinfo(bufid, NULL, NULL, &tid);
  pvm_upklong(&idle_sec, 1, 1);
  pvm_upklong(&busy_sec, 1, 1);
  
  slave = findSlave(tid);
  slaveUtil[slave].idle_sec = idle_sec;
  slaveUtil[slave].busy_sec = busy_sec;
 
  slaveStats[slave] = SSTerminated;

  return slave;
}

void
printutil()
{
  int j,k,nhost,narch,info;
  struct pvmhostinfo *hostp;
  info = pvm_config(&nhost,&narch,&hostp);

  if (info < 0) {
    fprintf(stderr, "Error while reading PVM configuration.\n");
    bailOut(1);
  }
  
  printf("\n\tMANDELBROT\tSLAVE\tUTILIZATION\n\n");
  printf("SLAVE\tHOST\t\tTIME IDLE\tTIME BUSY\tNUMBER OF PIXEL CALCULATIONS\n");
  for (j=0; j<tids_len ;j++) {
    k = 0;
    while(hostp[k].hi_tid !=pvm_tidtohost(tids[j])) {
      k++;
    }
    
    printf("%d%20s%12ld%15ld%12d\n",j,
	   hostp[k].hi_name,
	   slaveUtil[j].idle_sec,
	   slaveUtil[j].busy_sec,
	   slaveUtil[j].pixels);
  }
}

void
createCmap()
{
  double value;
  unsigned char r, g, b;
  int i, j, maxcols;

  cmap = (unsigned char *) malloc(sizeof(unsigned char)*maxiter*3);

  j = 0;
  maxcols = 20;
  for (i=0; i < maxiter; i++) {
    j++;

    if (j > maxcols) {
      j = 0;
    }

    value = (double)j/maxcols;

  /*
     This colormap is taken from the applet at 
       http://www.geocities.com/CapeCanaveral/Hall/2368/mandelbrot/home.html
     color[0][i] = new Color((float)Math.abs(Math.sin(2*Math.PI*value)), 1.0f - value, (float)Math.abs(Math.cos(2*Math.PI*value)));
     color[1][i] = new Color(value * 1.0f, (float)Math.abs(Math.cos(2*Math.PI*value)), 1.0f - value);
     color[2][i] = new Color(1.0f - value, (float)Math.abs(Math.cos(2*Math.PI*value)), (float)Math.abs(Math.sin(2*Math.PI*value)));
     color[3][i] = new Color((float)Math.abs(Math.cos(2*Math.PI*value)), (float)Math.abs(Math.sin(2*Math.PI*value)), 1.0f - value);
  */

    r = (unsigned char) (fabs(sin(M_PI*2*value))*255);
    g = (unsigned char) (((double)1 - value)*255);
    b = (unsigned char) (fabs(cos(M_PI*2*value))*255);

    *(cmap + i*3 + 0) = r;
    *(cmap + i*3 + 1) = g;
    *(cmap + i*3 + 2) = b;
  }

  *(cmap + maxiter*3 + 0) = 0;
  *(cmap + maxiter*3 + 1) = 0;
  *(cmap + maxiter*3 + 2) = 0;
}

void 
setPix(Coord p, unsigned char r, unsigned char g, unsigned char b)
{
  unsigned char *pix;

  if (p.x >=0 && p.x < xres &&
      p.y >=0 && p.y < yres) {
    pix = image + (p.x + xres*p.y)*3;
    *(pix+0) = r;
    *(pix+1) = g;
    *(pix+2) = b;
  }
  else {
    fprintf(stderr, "Trying to set invalid pixel\n");
  }
}

void 
setPixWithCmap(Coord p, int cmapIndex)
{
  unsigned char r, g, b;

  r = *(cmap + cmapIndex*3 + 0);
  g = *(cmap + cmapIndex*3 + 1);
  b = *(cmap + cmapIndex*3 + 2);

  setPix(p, r, g, b);
}

void
bailOut(int rc) {
  killSlaves();

  if (tids != NULL) {
    free(tids);
  }

  if (slaveStats != NULL) {
    free(slaveStats);
  }

  if (slaveComp != NULL) {
    free(slaveComp);
  }

  if (slaveUtil != NULL) {
    free(slaveUtil);
  }

  if (cmap != NULL) {
    free(cmap);
  }

  if (image != NULL) {
    free(image);
  }

  exit(rc);
}

int
main(int argc, char *argv[]) 
{
  char *configFilename;
  FILE *confFile;
  int rc;
  int x, y;
  Compix c;
  int slave;
  int i;

  tids = NULL;
  slaveStats = NULL;
  slaveComp = NULL;
  slaveUtil = NULL;
  tids_len = 0;
  image = NULL;
  cmap = NULL;

  if (argc < 7) {
    printf("usage: <config_filename> <xmin> <xmax> <ymin> <ymax> <maxiter>");
    maxiter = 480;
    image = (unsigned char *) malloc(sizeof(unsigned char)*640*480*3);
    createCmap();
    for (y=0; y < 480; y++) {
      if (y < maxiter) {
	for (x=0; x < 640; x++) {
	  *(image+(x+640*y)*3+0) = *(cmap +y*3 +0);
	  *(image+(x+640*y)*3+1) = *(cmap +y*3 +1);
	  *(image+(x+640*y)*3+2) = *(cmap +y*3 +2);
	}
      }
    }
    /* memset(image, 0xff, 640*240 *sizeof(unsigned char)*3); */
    writeTGA("test.tga", image, 640, 480);

    bailOut(1);
  }

  configFilename=argv[1];
  xmin = atof(argv[2]);
  xmax = atof(argv[3]);
  ymin = atof(argv[4]);
  ymax = atof(argv[5]);
  maxiter = atoi(argv[6]);

  xres = X_RES;
  yres = Y_RES;

  image = (unsigned char *) malloc(sizeof(unsigned char)*xres*yres*3);
  createCmap();

  printf("Creating Mandelbrot Set image of size (%d,%d)\n", xres, yres);

  /* run sequential code */
  printf("Sequential Code run\n");
  {
    int x,y, iters, temp, lengthsq;
    Complex c, z;
    Coord p;
    struct timeval tv1, tv2;
    int dt1;

    gettimeofday(&tv1, (struct timezone*)0);
    for (y=0; y < yres; y++) {
      for (x=0; x < xres; x++) {
	p.x = x; p.y = y;
	z.real = z.imag = 0.0;
	c.real = xmin + (float) x * (xmax - xmin)/xres;
	c.imag = ymax - (float) y * (ymax - ymin)/yres;
	iters = 0;

	do {
	  temp = z.real*z.real - z.imag*z.imag + c.real;
	  z.imag = 2.0*z.real*z.imag + c.imag;
	  z.real = temp;
	  lengthsq = z.real*z.real+z.imag*z.imag;
	  iters++;
	  
	} while (lengthsq < 4.0 && iters < maxiter);

	setPixWithCmap(p, iters); 
      }
    }
    gettimeofday(&tv2, (struct timezone*)0);

    dt1 = (tv2.tv_sec - tv1.tv_sec) * 1000000
      + tv2.tv_usec - tv1.tv_usec;

    printf("Creating image took %d usecs, that are %ld pix/sec\n", dt1, (long)xres*yres*1000000/dt1);
  }

  printf("\nParallel Code run\n");
  {
    struct timeval tv1, tv2;
    int dt1;

    gettimeofday(&tv1, (struct timezone*)0);

    /* Enroll In PVM */
    pvm_mytid();
    
    tids = (int *) malloc(sizeof(int)*MAX_NUMBER_OF_SLAVES);
    tids_len = 0;
    
    
    confFile = fopen(configFilename, "r");
    
    if (confFile == NULL) {
      fprintf(stderr, "Unable to open file: %s\n", configFilename);
      bailOut(1);
    }
    
    if (spawnSlaves(confFile) != 0) {
      fprintf(stderr, "Error while spawning slaves\n");
      fclose(confFile);
      bailOut(1);
    }
    
    if(fclose(confFile) != 0) {
      fprintf(stderr, "Unable to close file; %s\n", configFilename);
      bailOut(1);
    }
    
    slaveStats = (SlaveStat *) malloc(sizeof(SlaveStat)*tids_len);
    slaveComp = (Compix *) malloc(sizeof(Compix)*tids_len);
    slaveUtil = (Utilization *) malloc(sizeof(Utilization)*tids_len);
    
    for (i=0; i<tids_len; i++) {
      slaveStats[i] = SSIdle;
      slaveUtil[i].pixels = 0;
    }
    
    slave = 0;
    y = 0;
    do {
      x = 0;
      do {
	c.p.x = x;
	c.p.y = y;
	c.c.real = xmin + (float) x * (xmax - xmin)/xres;
	c.c.imag = ymax - (float) y * (ymax - ymin)/yres;
	
	/* send pixel */
	slave = nextIdleSlave(slave);
	
	if (slave != -1) { /* found one */
	  rc = sendPix(slave, c);
	  
	  if (rc == 0) { 
	    x++;
	  }
	  else { /* error */
	    fprintf(stderr, "Error while sending pixel to tid %d\n", tids[slave]);
	  }
	}
	else { /* wait for next */
	  slave = receivePix();
	}
	
      } while (x<xres);
      y++;
    } while (y<yres);
    
    /* receiving the remaining data */
    for (i=0; i<tids_len; i++) {
      slave = receivePix();
    }
    
    /* and send quit */
    for (i=0; i<tids_len; i++) {
      sendQuit(i);
    }
    
    for (i=0; i<tids_len; i++) {
      slave = receiveStat();
    }
    
    /*TODO print util stat here */
    printutil();

    gettimeofday(&tv2, (struct timezone*)0);
    
    dt1 = (tv2.tv_sec - tv1.tv_sec) * 1000000
      + tv2.tv_usec - tv1.tv_usec;
    
    printf("Creating image took %d usecs, that are %ld pix/sec\n", dt1, (long)xres*yres*1000000/dt1);
  }
  writeTGA("mandel.tga", image, xres, yres);

  bailOut(0);
  return 0;
}