/* Implementation of parallel quicksort */

#include "csapp.h"
#include "taskq.h"
#include "pqsort.h"

/* Tuning parameters */
/* Verbosity level */
int verbose = 0;
/* Use library quicksort or homebrewed one? */
int use_qsort_lib = 0;
/* How small do partitions need to be relative to original size
   before using sequential sort & sequential partititoning */
size_t serial_sort_fraction = 32;
size_t serial_partition_fraction = 1;

/* Counting number of comparisons */
volatile size_t comp_cnt;
pthread_mutex_t count_mutex;

/* Debugging support */

#define SHOW_NELE_MAX 30
void show_data(data_t *base, size_t nele) {
  size_t i;
  for (i = 0; i < nele && i < SHOW_NELE_MAX; i++) {
    printf(" %llu", (printd_t) base[i]);
  }
  printf("\n");
}


/* For tracking debugging information */
static data_t *global_base = NULL;
static data_t *global_scratch_base = NULL;
static data_t *global_end = NULL;
static data_t *global_scratch_end = NULL;

size_t global_index(data_t *p) {
  if (p >= global_base && p <= global_end) 
    return p - global_base;
  if (p >= global_scratch_base && p <= global_scratch_end) 
    return p - global_scratch_base;
  return (size_t) -1;  
}

#if 0
ls #define CHK(p) check_range(p)
#else
#define CHK(p) p
#endif

#if LOGCOMPS
#define LTE(x,y) less_equal_than(x,y)
#else
#define LTE(x,y) ((x) <= (y))
#endif

/* Make sure pointer is within range of data */
static data_t *check_range(data_t *p) {
  if (p >= global_base && p <= global_end)
    return p;
  if (p >= global_scratch_base && p <= global_scratch_end)
    return p;
  printf("Error: Attempt to access element %ld or %ld\n",
	 (long int) (p - global_base),
	 (long int) (p - global_scratch_base));
  exit(1);
}

/* Compare two numbers and increment count */
static int less_equal_than(data_t x, data_t y) {
#ifdef LOGCOMPS
  pthread_mutex_lock(&count_mutex);
#endif  
  comp_cnt++;
#ifdef LOGCOMPS
  pthread_mutex_unlock(&count_mutex);
#endif
  return x <= y;
}

/*
  Partition size when start doing sequential sort
  (Computed at beginning)
*/
static size_t nele_max_sort_serial = (1UL << 24);
/*
  Partition size when start doing sequential sort
  (Computed at beginning)
*/
static size_t nele_max_partition_serial = (1UL << 24);

/* Comparision function for qsort */
int dcomp(const void *p1, const void *p2) {
  data_t x1 = *CHK((data_t *) p1);
  data_t x2 = *CHK((data_t *) p2);
  if (!LTE(x2,x1))
    return -1;
  if (x1 == x2)
    return 0;
  return 1;
}

/* Using library quicksort */
void qsort_lib(data_t *base, size_t nele, data_t *scratch_base) {
  qsort((void *) base, nele, sizeof(data_t), dcomp);
}

/* Pick pivot for quicksort */
static data_t *pivot(data_t *base, size_t nele) {
  return base + nele/2;
}

static void swap(data_t *p1, data_t *p2) {
  data_t x1 = *CHK(p1);  data_t x2 = *CHK(p2);
  *p1 = x2; *p2 = x1;
}

/* Partition array into elements <= and >= to pivot.
   Return index of pivot element.
*/
static size_t partition_serial(data_t *base, size_t nele);
static size_t partition_parallel(data_t *base, size_t nele,
				 data_t *scratch_base, size_t block_size);


static size_t partition_sweep(data_t *base, size_t nele, data_t pivot);

static size_t partition(data_t *base, size_t nele, data_t *scratch_base)
{
  if (nele <= nele_max_partition_serial)
    return partition_serial(base, nele);
  else
    return partition_parallel(base, nele, scratch_base, nele_max_partition_serial);
}

static size_t partition_serial(data_t *base, size_t nele) {
  /* Select pivot */
  data_t *p = pivot(base, nele);
  /* Move to right of array */
  data_t pval = *CHK(p);
  swap(p, base+nele-1);
  /* Do partititoning sweep */
  size_t pindex = partition_sweep(base, nele-1, pval);
  /* Move pivot to be first element of right partition */
  swap(base + pindex, base+nele-1);
  if (verbose >= 3) {
    size_t l = global_index(base);
    size_t m = global_index(base+pindex);
    size_t r = global_index(base + nele - 1);
    printf("Serial partition.  [%lu,%lu,%lu]\n",
	   (printi_t) l, (printi_t) m, (printi_t) r);
  }
  return pindex;
}

/* Partitioning sweep.
   Partitions into left (all elements <= p) and right (all elements >= p).
   Returns count of number of elements in left partition.
*/
size_t partition_sweep(data_t *base, size_t nele, data_t pval)
{
  data_t *left = base;
  data_t *right = base+nele-1;
   while (1) {
     while (left <= right && LTE(*CHK(left), pval)) {
      left++;
    }
     while (left <= right && LTE(pval, *CHK(right))) {
      right--;
    }
    if (left < right) {
      swap(left, right);
      left++;
      right--;
    } else
      break;
  }
  return left - base;
}

/* Parallel copy.  Needed as part of parallel partitioning */
/* Have structure that defines copying task */
typedef struct {
  data_t *src;
  data_t *dest;
  size_t nele;
} copy_task_t, *copy_task_ptr;

void *new_copy_task(data_t *src, data_t *dest, size_t nele) {
  copy_task_ptr p = (copy_task_ptr) Malloc(sizeof(copy_task_t));
  p->src = src;
  p->dest = dest;
  p->nele = nele;
  return (void *) p;
}

/* Copying thread function */
void *copy_thread(void *vargp) {
  copy_task_ptr p = (copy_task_ptr) vargp;
  data_t *src = p->src;
  data_t *dest = p->dest;
  size_t nele = p->nele;
  free(vargp);
  if (verbose >= 3) {
    printf("Copy thread.  %lu elements from %lu to %lu\n",
	   (printi_t) nele,
	   (printi_t) global_index(src),
	   (printi_t) global_index(dest));
  }
  bcopy((void *) src, (void *) dest, nele * sizeof(data_t));
  return NULL;
}

void dcopy_parallel(data_t *src, data_t *dest, size_t nele, size_t block_size) {
  int nblock = (nele+block_size-1) / block_size;
  int npb = nele / nblock;
  task_queue_ptr tq = new_task_queue();
  int b;
  if (verbose >= 2) {
    printf("Dcopy.  Copying %lu elements from %lu to %lu in %d blocks\n",
	   (printi_t) nele, (printi_t) global_index(src),
	   (printi_t) global_index(dest), nblock);
  }
  for (b = 0; b < nblock; b++) {
    data_t *bsrc = src + b*npb;
    data_t *bdest = dest + b*npb;
    size_t bnele = (b < nblock-1) ? npb : nele - b*npb;
    void *p = new_copy_task(bsrc, bdest, bnele);
    spawn_task(tq, copy_thread, p);
  }
  join_tasks(tq);
  free_task_queue(tq);
}

/* Parallel partitioning
   1. Copy (in parallel) to scratch space 
   2. Do partitioning
     a. Choose global pivot p.
     b. Break into k blocks and serially partition each block by p.
     c. Copy left subblocks into left hand side 
        & right subblocks into right hand side
*/

/* All information needed for parallel partitioning */
typedef struct {
  data_t *base;
  data_t *scratch_base;      /* Extra space for setting up partitions */
  size_t next_left_index;    /* Next position for accumulating left subblocks */
  size_t next_right_index;   /* Next position for accumulating subblocks */
  pthread_mutex_t mutex;     /* Mutex for shared info */
  data_t pval;               /* Pivot value */
} partition_t;

void init_partition(partition_t *par, data_t *base, size_t nele,
		    data_t *scratch_base, data_t pval)
{
  par->base = base;
  par->scratch_base = scratch_base;
  par->next_left_index = 0;
  par->next_right_index = nele;
  pthread_mutex_init(&par->mutex, NULL);
  par->pval = pval;
}

/* Information stored in each entry task queue for partitioning */
typedef struct {
  size_t nele;
  size_t src_index;
  partition_t *par;
} partition_task_t;

void *new_partition_task(size_t nele, size_t src_index, partition_t *par)
{
  partition_task_t *pt = Malloc(sizeof(partition_task_t));
  pt->nele = nele;
  pt->src_index = src_index;
  pt->par = par;
  return (void *) pt;
}

/* Thread routine for partitioning */
static void *partition_thread(void *vargp) {
  partition_task_t *pt = (partition_task_t *) vargp;
  size_t nele = pt->nele;
  size_t src_index = pt->src_index;
  partition_t *par = pt->par;
  free(vargp);
  /* Do sequential partitioning */
  size_t nleft = partition_sweep(par->scratch_base + src_index, nele,
				 par->pval);
  size_t nright = nele - nleft;
  /* Get destination positions */
  pthread_mutex_lock(&par->mutex);
  size_t left_index = par->next_left_index;
  par->next_left_index += nleft;
  par->next_right_index -= nright;
  size_t right_index = par->next_right_index;
  pthread_mutex_unlock(&par->mutex);
  if (verbose >= 2) {
    size_t glsrc = global_index(par->scratch_base + src_index);
    size_t gldest = global_index(par->base + left_index);
    size_t grsrc = global_index(par->scratch_base + src_index + nleft);
    size_t grdest = global_index(par->base + right_index);
    printf("Partition thread.  %lu elements\n", (printi_t) nele);
    printf("\tMoving left subblock (%lu elements) from sb[%lu] to b[%lu]\n",
	   (printi_t) nleft, (printi_t) glsrc, (printi_t) gldest);
    printf("\tMoving right subblock (%lu elements) from sb[%lu] to s[%lu]\n",
	   (printi_t) nright, (printi_t) grsrc, (printi_t) grdest);
  }
  bcopy((void *) (par->scratch_base + src_index),
	(void *) (par->base + left_index),
	nleft * sizeof(data_t)); 
  bcopy((void *) (par->scratch_base + src_index + nleft),
	(void *) (par->base + right_index),
	nright * sizeof(data_t));
  return NULL;
}

static size_t partition_parallel(data_t *base, size_t nele,
				 data_t *scratch_base, size_t block_size)
{
  /* Select pivot */
  data_t *p = pivot(base, nele);
  data_t pval = *CHK(p);
  /* Move to right of array */
  swap(p, base+nele-1);
  /* Copy everything else into scratch area */
  dcopy_parallel(base, scratch_base, nele-1, block_size);
  /* Now run multiple partition tasks */
  task_queue_ptr tq = new_task_queue();
  partition_t partition;
  init_partition(&partition, base, nele-1, scratch_base, pval);
  int nblock = (nele + block_size - 2) / block_size;
  size_t npb = (nele-1) / nblock;
  int b;
  for (b = 0; b < nblock; b++) {
    size_t src_index = b*npb;
    size_t bnele = (b < nblock-1) ? npb : nele-1 - b*npb;
    void *p = new_partition_task(bnele, src_index, &partition);
    spawn_task(tq, partition_thread, p);
  }
  join_tasks(tq);
  size_t left_size = partition.next_left_index;
  free_task_queue(tq);
  /* Swap pivot back */
  swap(base+left_size, base+nele-1);
  if (verbose >= 2) {
    printf("Parallel partitioning.  Pivot = %lu. [%lu,%lu,%lu] in %d blocks\n",
	   (printi_t) pval,
	   (printi_t) global_index(base),
	   (printi_t) global_index(base+left_size),
	   (printi_t) global_index(base+nele-1),
	   nblock);
  }
  if (verbose >= 3) {
    printf("Partitioned results: ");
    show_data(base, nele);
  }
  return left_size;
}

void qsort_serial(data_t *base, size_t nele,
		  data_t *scratch_base) {
  if (verbose >= 3) {
    size_t l = global_index(base);
    size_t r = global_index(base + nele - 1);
    printf("Qsort_serial.  [%lu,%lu]\n", (printi_t) l, (printi_t) r);
  }
  if (nele <= 1)
    return;
  if (nele == 2) {
    if (base[0] > base[1])
      swap(base, base+1);
    return;
  }
  size_t m = partition_serial(base, nele);
  if (m > 1)
    qsort_serial(base, m, scratch_base);
  if (nele-1 > m+1)
    qsort_serial(base+m+1, nele-m-1, scratch_base+m+1);
}

/* For threaded sorting, have structure that defines sorting task */
typedef struct {
  data_t *base;
  size_t nele;
  data_t *scratch_base;
  task_queue_ptr tq;
} sort_task_t, *sort_task_ptr;

static void init_task(size_t nele) {
  if (serial_sort_fraction > 1) {
    nele_max_sort_serial = nele / serial_sort_fraction;
    if (nele_max_sort_serial < 2)
      nele_max_sort_serial = 2;
  }
  else
    nele_max_sort_serial = nele;
  if (verbose >= 1) {
    printf("Serial sort cutover = %lu\n", (printi_t) nele_max_sort_serial);
  }
  if (serial_partition_fraction > 1) {
    nele_max_partition_serial = nele / serial_partition_fraction;
    if (nele_max_partition_serial < 2)
      nele_max_partition_serial = 2;
  }
  else
    nele_max_partition_serial = nele;
  if (verbose >= 1) {
    printf("Serial partition cutover = %lu\n", (printi_t) nele_max_partition_serial);
  }
#ifdef LOGCOMPS
  pthread_mutex_init(&count_mutex, NULL);
#endif  
}

/* Add a new task to the list.  Returns position in task list */
static sort_task_ptr new_task(data_t *base, size_t nele,
			      data_t *scratch_base, task_queue_ptr tq) {
  sort_task_ptr t = Malloc(sizeof(sort_task_t));
  t->base = base;
  t->nele = nele;
  t->scratch_base = scratch_base;
  t->tq = tq;
  if (verbose >= 3) {
    size_t l = global_index(base);
    size_t r = global_index(base + nele-1);
    printf("New_task. [%lu,%lu]\n", (printi_t) l, (printi_t) r);
  }
  return t;
}

static void tqsort_helper(data_t *base, size_t nele,
			  data_t *scratch_base, task_queue_ptr tq);

/* Thread routine for many-threaded quicksort */
static void *sort_thread(void *vargp) {
  sort_task_t *t = (sort_task_t *) vargp;
  data_t *base = t->base;
  size_t nele = t->nele;
  data_t *scratch_base = t->scratch_base;
  task_queue_ptr tq = t->tq;
  free(vargp);
  if (verbose >= 2) {
    size_t l = global_index(base);
    size_t r = global_index(base + nele-1);
    printf("sort_thread.  tid=0x%lx  [%lu,%lu]\n",
	   (printi_t) pthread_self(), (printi_t) l, (printi_t) r);
  }
  size_t m = partition(base, nele, scratch_base);
  if (m > 1)
    tqsort_helper(base, m, scratch_base, tq);
  if (nele-1 > m+1)
    tqsort_helper(base+m+1, nele-m-1, scratch_base+m+1, tq);
  return NULL;
}

/* Multi-threaded quicksort */
static void tqsort_helper(data_t *base, size_t nele,
			  data_t *scratch_base, task_queue_ptr tq) {
  if (verbose >= 2) {
    size_t l = global_index(base);
    size_t r = global_index(base+nele-1);
    printf("tqsort.  [%lu,%lu]\n", (printi_t) l, (printi_t) r);
  }
  if (nele <= nele_max_sort_serial) {
    /* Use sequential sort */
    if (use_qsort_lib)
      qsort_lib(base, nele, scratch_base);
    else
      qsort_serial(base, nele, scratch_base);
    return;
  }
  sort_task_t *t = new_task(base, nele, scratch_base, tq);
  spawn_task(tq, sort_thread, (void *) t);
}

void tqsort(data_t *base, size_t nele, data_t *scratch_base) {
  init_task(nele);
  global_base = base;
  global_end = global_base + nele - 1;
  global_scratch_base = scratch_base;
  global_scratch_end = global_scratch_base + nele - 1;
  task_queue_ptr tq = new_task_queue();
  tqsort_helper(base, nele, scratch_base, tq);
  int max_tasks = join_tasks(tq);
  free_task_queue(tq);
  if (verbose >= 1)
    printf("Sorted %lu elements with max %d concurrent tasks\n",
	   (printi_t) nele, max_tasks);
}