#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <float.h>
#include <math.h>

#include <assert.h>

#include <mpi.h>

#include "smut.h"
#include "stripesmut.h"
#include "timer.h"
#include "auction.h"
#include "stripe_auction.h"

#if !defined(HAVE_RESTRICT)
#define restrict
#endif

MPI_Comm comm;
int rank, csize, am_root;

#if !defined(HAVE_MPI_ALLOC_MEM)
static int
MPI_Alloc_mem (MPI_Aint size, MPI_Info info, void *baseptr)
{
  void **out = (void**)baseptr;
  *out = malloc (size);
  if (*out) return MPI_SUCCESS;
  return MPI_ERR_NO_MEM;
}
static int
MPI_Free_mem (void *baseptr)
{
  free(baseptr);
  return MPI_SUCCESS;
}
#if !defined(MPI_INFO_NULL)
#define MPI_INFO_NULL 0
#endif
#endif


struct cbuf_entry {
  double price;
  int i, j;
};

struct pricebuf_entry {
  double price;
  int i;
};


struct comm_state_t {
  int csize, rank;
  int phase;

  int max_lnc;

  int send_neighbor;
  int n_to_send;
  struct cbuf_entry *sendbuf;

  int merged;

  int recv_neighbor;
  int n_recvd;
  struct cbuf_entry *recvbuf;

  int bcast_mode, bcast_root;
  int *n_unmatched;

  MPI_Request req[2];
  MPI_Status status[2];

  int *g2l;

  MPI_Datatype cbuf_entry_type;
};

int
comm_get_changes (struct comm_state_t *state, MPI_Comm comm,
		  int *n_changes_out, const struct cbuf_entry **clist)
{
  const int phase = ++state->phase;
  struct cbuf_entry *tmp;

  if (phase >= state->csize) {
    *n_changes_out = 0;
    *clist = NULL;
    return 0;
  }

  if (state->bcast_mode) {
    /*
    fprintf (stderr, "%d: bcasting %d from %d\n",
	     state->rank, state->n_to_send, state->bcast_root);
    */
    assert(state->rank == state->bcast_root || state->n_to_send == 0);
    MPI_Bcast (&state->n_to_send, 1, MPI_INT, state->bcast_root, comm);
    MPI_Bcast (state->sendbuf, state->n_to_send, state->cbuf_entry_type,
	       state->bcast_root, comm);
    *n_changes_out = state->n_to_send;
    *clist = state->sendbuf;
    state->phase = state->csize;
    return 1;
  }

  MPI_Waitall (2, &state->req[0], &state->status[0]);
  MPI_Get_count (&state->status[0], state->cbuf_entry_type,
		 &state->n_recvd);
  tmp = state->recvbuf;
  state->recvbuf = state->sendbuf;
  state->sendbuf = tmp;
  state->n_to_send = state->n_recvd;

  if (phase < state->csize) {
    MPI_Irecv (state->recvbuf, state->max_lnc,
	       state->cbuf_entry_type,
	       state->recv_neighbor, phase,
	       comm, &state->req[0]);
    MPI_Isend (state->sendbuf, state->n_to_send,
	       state->cbuf_entry_type,
	       state->send_neighbor, phase, comm, &state->req[1]);
  }

  /* not kosher; send should be over first */
  *n_changes_out = state->n_to_send;
  *clist = state->sendbuf;

  return 1;
}

void
free_comm_state (struct comm_state_t *state, MPI_Comm comm)
{
  if (state->g2l) free (state->g2l);
  state->g2l = NULL;
  MPI_Free_mem (state->sendbuf); state->sendbuf = NULL;
  MPI_Free_mem (state->recvbuf); state->recvbuf = NULL;
  MPI_Type_free (&state->cbuf_entry_type);
}

static void
make_pricebuf_entry_type (MPI_Datatype *pricebuf_entry_type)
{
  struct pricebuf_entry buf[2];
  MPI_Datatype bufent[3] = { MPI_DOUBLE, MPI_INT, MPI_UB };
  int blocklen[3] = { 1, 1, 1 };
  MPI_Aint disp[3];
  MPI_Aint base;

  MPI_Address (&buf[0], disp);
  MPI_Address (&buf[0].i, disp+1);
  MPI_Address (&buf[1], disp+2);
  base = disp[0];
  disp[0] -= base;
  disp[1] -= base;
  disp[2] -= base;
  MPI_Type_struct (3, blocklen, disp, bufent, pricebuf_entry_type);
  MPI_Type_commit (pricebuf_entry_type);

}

static void
make_cbuf_entry_type (MPI_Datatype *cbuf_entry_type)
{
  struct cbuf_entry buf[2];
  MPI_Datatype bufent[4] = { MPI_DOUBLE, MPI_INT, MPI_INT, MPI_UB };
  int blocklen[4] = { 1, 1, 1, 1 };
  MPI_Aint disp[4];
  MPI_Aint base;

  MPI_Address (&buf[0], disp);
  MPI_Address (&buf[0].i, disp+1);
  MPI_Address (&buf[0].j, disp+2);
  MPI_Address (&buf[1], disp+3);
  base = disp[0];
  disp[0] -= base;
  disp[1] -= base;
  disp[2] -= base;
  disp[3] -= base;
  MPI_Type_struct (4, blocklen, disp, bufent, cbuf_entry_type);
  MPI_Type_commit (cbuf_entry_type);

}

void
init_comm_state (struct comm_state_t *state, const struct gcsr_t *gA,
		 MPI_Comm comm)
{
  int rank, csize;
  const int gnr = gA->nr, gnc = gA->nc;
  int lnc = gA->local.nc, max_lnc;
  int j;
  int *g2l = NULL;

  MPI_Comm_size (comm, &csize); state->csize = csize;
  MPI_Comm_rank (comm, &rank); state->rank = rank;
  MPI_Allreduce (&lnc, &max_lnc, 1, MPI_INT, MPI_MAX, comm);
  state->max_lnc = max_lnc;

  state->send_neighbor = rank - 1;
  if (state->send_neighbor < 0) state->send_neighbor = csize - 1;
  state->recv_neighbor = (rank + 1);
  if (state->recv_neighbor >= csize) state->recv_neighbor = 0;

  state->n_to_send = 0;
  MPI_Alloc_mem (max_lnc * sizeof(struct cbuf_entry),
		 MPI_INFO_NULL, &state->sendbuf);
  state->n_recvd = 0;
  MPI_Alloc_mem (max_lnc * sizeof(struct cbuf_entry),
		 MPI_INFO_NULL, &state->recvbuf);

  state->bcast_mode = 0;
  state->bcast_root = -1;
  MPI_Alloc_mem (csize * sizeof(int), MPI_INFO_NULL,
		 &state->n_unmatched);

  make_cbuf_entry_type (&state->cbuf_entry_type);

  if (gA->colmap.l2g) {
    const int lnc = gA->local.nc;
    const int * restrict l2g = gA->colmap.l2g;
    g2l = malloc (gnc * sizeof(int));
    memset (g2l, -1, gnc * sizeof(int));
    for (j = 0; j < lnc; ++j) g2l[l2g[j]] = j;
  }
  state->g2l = g2l;
}

int
comm_g_n_unmatched (struct comm_state_t *state, int n_unmatched, int n_changes,
		    MPI_Comm comm)
{
  int out = 0;
  int nactive = 0;
  int active_p = -1;
  int p, csize;
  csize = state->csize;
  out = n_unmatched + n_changes;
  state->bcast_mode = 0;
/*
  fprintf (stderr, "%d: allgathering %d %p %p\n", state->rank, out,
	   &out, state->n_unmatched);
*/
  MPI_Allgather (&out, 1, MPI_INT,
		 state->n_unmatched, 1, MPI_INT,
		 comm);
//  fprintf (stderr, "%d: allgathered\n");
  for (p = 0; p < csize; ++p) {
    /*
    fprintf (stderr, "%d:  seeing %d active on %d\n", state->rank,
	     state->n_unmatched[p], p);
    */
    if (state->n_unmatched[p]) {
      ++nactive;
      active_p = p;
    }
    out += state->n_unmatched[p];
  }
  if (nactive == 1) {
    state->bcast_mode = 1;
    state->bcast_root = active_p;
    /*
    fprintf (stderr, "%d: setting bcast mode, root %d\n", state->rank,
	     active_p);
    */
  }
  return out;
}

void
comm_give_changes (struct comm_state_t *state, const struct gcsr_t *gA,
		   const int infeasible,
		   const int n_changes,
		   const int *changed_col_list_in,
		   const int *col_changed_flag_in,
		   const int *lmatch_in, const double *lprice_in,
		   MPI_Comm comm)
{
  struct cbuf_entry * restrict sendbuf = state->sendbuf;

  const int * restrict changed_col_list = changed_col_list_in;
  const int * restrict col_changed_flag = col_changed_flag_in;
  const int * restrict lmatch = lmatch_in;
  const double * restrict lprice = lprice_in;

  const int gnr = gA->nr;
  const int gnc = gA->nc;
  const int * restrict l2g = gA->colmap.l2g;
  const int lindstart = gA->rowmap.lindstart;
  /*const int lnind = gA->rowmap.lnind;*/

  int k;

  if (infeasible) { /* infeasible */
    state->n_to_send = 1;
    sendbuf[0].i = -1;
    sendbuf[0].j = -1;
    sendbuf[0].price = -1.0;
  }
  else {
    for (k = 0; k < n_changes; ++k) {
      int j, gj;
      j = changed_col_list[k];
      gj = (l2g? l2g[j] : j);

      sendbuf[k].j = gj;
      sendbuf[k].i = lmatch[j];
      if (lmatch[j] >= 0)
	sendbuf[k].i += lindstart;
      sendbuf[k].price = lprice[j];
    }

    state->n_to_send = n_changes;
  }

  state->phase = 0;
  state->merged = 1;
  if (csize > 0 && !state->bcast_mode) {
    MPI_Irecv (state->recvbuf, state->max_lnc,
	       state->cbuf_entry_type,
	       state->recv_neighbor, 0,
	       comm, &state->req[0]);
    MPI_Isend (state->sendbuf, state->n_to_send,
	       state->cbuf_entry_type,
	       state->send_neighbor, 0, comm, &state->req[1]);
  }
}


void
merge_cbuf_into_vars (const int n_changes,
		      const struct cbuf_entry* cbuf_in,
		      const int *g2l_in,
		      const struct gcsr_t *gA,
		      int *lmatch_in, double *lprice_in,
		      int *n_unmatched_out, int *unmatched_list)
{
  const struct cbuf_entry * restrict cbuf = cbuf_in;

  int n_unmatched = *n_unmatched_out;
  int * restrict unmatched = unmatched_list;
  int * restrict lmatch = lmatch_in;
  double * restrict lprice = lprice_in;

  const int gnr = gA->nr;
  const int gnc = gA->nc;
  const int lnc = gA->local.nc;
  const int * restrict l2g = gA->colmap.l2g;
  const int * restrict g2l = g2l_in;
  const int lindstart = gA->rowmap.lindstart;
  const int lnind = gA->rowmap.lnind;

  int k;

  if (n_unmatched < 0) return; /* infeas */

  /*fprintf (stderr, "%d: merging\n", rank);*/

  for (k = 0; k < n_changes; ++k) {
    const double newp = cbuf[k].price;
    const int gi = cbuf[k].i;
    const int gj = cbuf[k].j;

    int i, old_i, j;
    double lp;

    j = (g2l? g2l[gj] : gj);
    i = gi - lindstart;

    if (j < 0) {
      *n_unmatched_out = j; /* infeas */
      return;
    }

    /*
    if (4 == j)
      fprintf (stderr, "%d: col %d\n", rank, j);
    */

    /*assert(i < 0 || i >= lnind);*/
    /* make idempotent instead */
    if (i < lnind && i >= 0) continue;

    lp = lprice[j];

    if (lp > newp) continue;
    if (lp == newp && i >= lnind) {
      /*
      fprintf (stderr, "%d: col %d: a tie: %g == %g, %d >= %d\n", rank,
	       j, lp, newp, i, lnind);
      */
      continue;
    }

    lprice[j] = newp;
    old_i = lmatch[j];
    lmatch[j] = -1;
    if (old_i >= 0) {
      unmatched[n_unmatched] = old_i;
      ++n_unmatched;
    }
  }
  *n_unmatched_out = n_unmatched;
}

void
comm_merge_pricebuf_into_vars (const struct pricebuf_entry *pricebuf_in,
			       const int *g2l_in,
			       const struct gcsr_t* gA,
			       int *lmatch_in, double *lprice_in,
			       int *g_n_unmatched,
			       int *n_unmatched_out, int *unmatched_list)
{
  const struct pricebuf_entry * restrict pricebuf = pricebuf_in;

  int n_unmatched = *n_unmatched_out;
  int * restrict unmatched = unmatched_list;
  int * restrict lmatch = lmatch_in;
  double * restrict lprice = lprice_in;

  const int gnr = gA->nr;
  const int gnc = gA->nc;
  const int lnc = gA->local.nc;
  const int * restrict l2g = gA->colmap.l2g;
  const int * restrict g2l = g2l_in;
  const int lindstart = gA->rowmap.lindstart;
  const int lnind = gA->rowmap.lnind;

  int j;

  for (j = 0; j < gnc; ++j) {
    int lj;
    int matchi;

    lj = (g2l? g2l[j] : j);
    if (!lnc) lj = -1;

    if (lj < 0) continue;
    matchi = pricebuf[j].i;

    if (matchi < 0) {
      *n_unmatched_out = -2*gnr;
      *g_n_unmatched = -2*gnr;
      return;
    }

    matchi -= lindstart; /* translate to local */

    if (matchi >= lnind || matchi < 0) /* map all external rows to -1 */
      matchi = -1;

    if (pricebuf[j].price > lprice[lj]) {
      assert(matchi < 0 || matchi >= lnind);
      lprice[lj] = pricebuf[j].price;
      if (lmatch[lj] >= 0) {
	unmatched[n_unmatched] = lmatch[lj];
	lmatch[lj] = -1;
	++n_unmatched;
      }
    }
    else if (pricebuf[j].price == lprice[lj]) {
      /*
	if (!(matchi == lmatch[lj] || matchi < 0)) {
	  fprintf (stderr, "%d:  j = %d, lj = %d, matchi = %d, lmatch = %d, lnind = %d\n",
	           rank, j, lj, matchi, lmatch[lj], lnind);
	  fprintf (stderr, "%d:  sent  matchi = %d (trans %d)\n",
	           rank, comm_state.pricebuf_send[j].i, comm_state.pricebuf_send[j].i - lindstart);
	}
      */
      assert (matchi == lmatch[lj] || matchi < 0);

      if (lprice[lj] == HUGE_VAL && matchi != lmatch[lj]) {
	n_unmatched = -2*gnr;
	  /*
	  fprintf (stderr,
		   "%d: tried to break inf price, col %d between %d(l %d) and %d (%d)\n",
		   rank,
		   j, lmatch[lj] + lindstart, lmatch[lj], comm_state.pricebuf[j].i, matchi);
	  */
	break;
      }

      if (matchi < 0 && lmatch[lj] != -1) {
	unmatched[n_unmatched] = lmatch[lj];
	lmatch[lj] = -1;
	++n_unmatched;
      }
    }
#if !defined(NDEBUG)
    else {
      assert(lmatch[lj] < 0 || matchi != lmatch[lj]);
      assert(matchi < 0);
    }
#endif
  }

  MPI_Allreduce (&n_unmatched, g_n_unmatched, 1, MPI_INT, MPI_SUM,
		 comm);
  *n_unmatched_out = n_unmatched;
}

int*
comm_state_g2l (struct comm_state_t* state)
{
  return state->g2l;
}

int
main (int argc, char **argv)
{
  FILE *f;
  int err;

  double mu_min;
  int relgap = 0;

  struct spcsr_t A, Acopy;
  struct gcsr_t gA;

  int lnr, lnc, lnent;

  double *R = NULL, *C = NULL;
  double *expint = NULL, *lexpint = NULL, *expint_check = NULL;
  int *rdisps = NULL, *rcounts = NULL;

  int *match = NULL;
  double *price = NULL;

  int *lmatch = NULL;
  double *lprice = NULL;

  int errflg, i, j;
  double primal, dual;

  struct Timer timer;

  struct comm_state_t comm_state;

  MPI_Init (&argc, &argv);
  comm = MPI_COMM_WORLD;
  MPI_Comm_size (comm, &csize);
  MPI_Comm_rank (comm, &rank);
  am_root = (0 == rank);

  assert(sizeof(int) == 4);

  spcsr_init_clear (&A);

  if (am_root) {
    if (argc <= 1) {
      printf("No files\n");
      return -1;
    }

    f = fopen(argv[1], "rb");
    if (!f) {
      perror("Error opening file: ");
      return -1;
    }

    err = spcsr_load_binfile (f, &A);
    if (err) {
      printf ("Error reading: %d\n", err);
      return -1;
    }
    fclose(f);

    R = malloc (A.nr * sizeof(double));
    C = malloc (A.nc * sizeof(double));

    printf ("nr = %d nc = %d nent = %d\n", A.nr, A.nc, A.nent);

#if 0
    spcsr_lascaling (&A, R, C);
    spcsr_apply_scaling (&A, R, C);
#elif 1
    spcsr_lascale (&A, R, C);
#endif

    spcsr_copy (&A, &Acopy);

    expint_check = malloc (Acopy.nent * sizeof(double));

    auction_toexpint (&Acopy, expint_check);
    auction_shift (&Acopy, expint_check);
  }

  gcsr_take_csr_root (&A, &gA, MPI_COMM_WORLD);
  gcsr_redist_ents (&gA, 0, MPI_COMM_WORLD);

  {
    int p;
    for (p = 0; p < csize; ++p) {
      MPI_Barrier(MPI_COMM_WORLD);
      if (p == rank) {
	printf("++++ %d: gA(%d, %d), lA(%d, %d; %d) rows [%d, %d)\n",
	       rank, gA.nr, gA.nc, gA.local.nr, gA.local.nc, gA.local.nent,
	       gA.rowmap.lindstart, gA.rowmap.lindstart + gA.rowmap.lnind);

	printf("++++ %d: local rowoff[1] = %d\n", rank,
	       (gA.local.rowoff? gA.local.rowoff[1] : -1));
      }
      MPI_Barrier(MPI_COMM_WORLD);
    }
  }

  lnr = gA.local.nr;
  lnc = gA.local.nc;
  lnent = gA.local.nent;

  /*
  fprintf (stderr, "%d: (lnr, lnc; lnent) = (%d, %d; %d)\n",
	   rank, lnr, lnc, lnent);
  fprintf (stderr, "%d: (gnr, gnc) = (%d, %d)\n",
	   rank, gA.nr, gA.nc);
  */

  lexpint = malloc (lnent * sizeof(double));
  auction_toexpint (&gA.local, lexpint);
  stripe_auction_shift (&gA, lexpint, MPI_COMM_WORLD);

  if (am_root) {
    int k;
    int die = 0;
    for (k = 0; k < gA.local.nent; ++k) {
      if (gA.local.entry[k] != Acopy.entry[k]) {
	fprintf (stderr, "gA.local.entry[%d] = %g    Acopy.entry[%d] = %g\n",
		 k, gA.local.entry[k], k, Acopy.entry[k]);
	die = 1;
      }
      if (gA.local.colind[k] != Acopy.colind[k]) {
	fprintf (stderr, "gA.local.colind[%d] = %d    Acopy.colind[%d] = %d\n",
		 k, gA.local.colind[k], k, Acopy.colind[k]);
	die = 1;
      }
      if (lexpint[k] != expint_check[k]) {
	fprintf (stderr, "lexpint[%d] = %g    expint_check[%d] = %g\n",
		 k, lexpint[k], k, expint_check[k]);
	die = 1;
      }
      if (die) MPI_Abort (comm, -9999);
    }
  }

  mu_min = 1.0 / gA.nr;
  if (argc > 2)
    mu_min = strtod(argv[2], NULL);
  if (argc > 3)
    relgap = 1;

  lmatch = malloc (lnc * sizeof (int));
  lprice = malloc (lnc * sizeof (double));

  MPI_Barrier(MPI_COMM_WORLD);
  initialize_timer (&timer);
  start_timer (&timer);
#if 0
  errflg = stripe_auction_simple (&gA, lexpint, lmatch, lprice, mu_min,
				  &comm_state, MPI_COMM_WORLD);
#else
  errflg = stripe_auction_scaling (&gA, lexpint, lmatch, lprice,
				   mu_min, relgap,
				   &comm_state, MPI_COMM_WORLD);
#endif
  stop_timer(&timer);
  MPI_Barrier(MPI_COMM_WORLD);

  if (am_root) {
    price = malloc (gA.nc * sizeof(double));
    match = malloc (gA.nc * sizeof(int));
    expint = malloc (Acopy.nent * sizeof(double));
    rcounts = malloc (csize * sizeof(int));
    rdisps = malloc (csize * sizeof(int));
  }

  stripe_auction_vars_l2g (&gA, lmatch, lprice, match, price, comm);

  if (1 == csize) {
    int j;
    for (j = 0; j < gA.nc; ++j) {
      assert(lmatch[j] == match[j]);
      if (lprice[j] != price[j]) {
	fprintf(stderr, "prices at %d don't match: %g %g\n",
		j, lprice[j], price[j]);
      }
      assert(lprice[j] == price[j]);
    }
  }

  /* now get expint */
  MPI_Gather (&gA.local.nent, 1, MPI_INT,
	      rcounts, 1, MPI_INT,
	      0, comm);
  if (am_root) {
    int p;
    rdisps[0] = 0;
    for (p = 1; p < csize; ++p)
      rdisps[p] = rdisps[p-1] + rcounts[p-1];
  }
  MPI_Gatherv (lexpint, gA.local.nent, MPI_DOUBLE,
	       expint, rcounts, rdisps, MPI_DOUBLE,
	       0, comm);

  auction_eval_primal_mdual (gA.local, lexpint, lmatch, lprice,
			     &primal, &dual);
  {
    double pd_in[2];
    double pd[2];
    pd_in[0] = primal; pd_in[1] = dual;
    MPI_Allreduce (&pd_in[0], &pd[0], 2, MPI_DOUBLE, MPI_SUM, comm);
    if (am_root) {
      printf ("Reduced primal: %g    dual %g\n", pd[0], pd[1]);
    }
  }

  gcsr_redist_root (&gA, comm);

#if 1
  if (am_root) {
    int k;
    for (k = 0; k < Acopy.nent; ++k) {
      if (expint[k] != expint_check[k]) {
	fprintf (stderr, "expint[%d] = %g    expint_check[%d] = %g  le = %g\n",
		 k, expint[k], k, expint_check[k], lexpint[0]);
      }
      assert(expint[k] == expint_check[k]);
      assert(gA.local.entry[k] == Acopy.entry[k]);
      assert(gA.local.colind[k] == Acopy.colind[k]);
    }
  }
#endif

  if (am_root) {
    auction_eval_primal_mdual (Acopy, expint, match, price,
			      &primal, &dual);
    printf ("Primal: %20g\nDual: %20g\nTime: %20g\nmu_min: %20g\n",
	    primal, dual, timer_duration(timer), mu_min);
    auction_eval_primal_mdual (gA.local, expint_check, match, price,
			      &primal, &dual);
    printf ("Primal: %20g\nDual: %20g\n", primal, dual);

#if 1
    {
      int i, j;
      int *invmatch;
      invmatch = malloc(Acopy.nr * sizeof(int));
      memset (invmatch, -1, Acopy.nr * sizeof(int));
      for (j = 0; j < Acopy.nc; ++j) {
	if (match[j] < 0) printf("column %d is unmatched\n", j);
	if (match[j] >= Acopy.nr) printf("column %d's match out of range (%d/%d)\n",
				     j, match[j], Acopy.nr);
	invmatch[ match[j] ] = j;
      }
      for (i = 0; i < Acopy.nr; ++i) {
	int k;
	int found_match = 0;
	if (invmatch[i] < 0)
	  printf("row %d isn't matched\n", i);
	for (k = Acopy.rowoff[i]; k < Acopy.rowoff[i+1]; ++k) {
	  const int j = Acopy.colind[k];
	  if (i == match[j]) ++found_match;
	}
	if (!found_match)	printf ("row %d is unmatched\n", i);
	if (found_match > 1) printf ("row %d is multiply matched\n", i);
      }
      free(invmatch);
    }
#endif
  }


  /* ------------ */
  gcsr_free (&gA, comm);
  spcsr_free (&A);

  if (am_root) {
    free (rdisps);
    free (rcounts);
    free (expint_check);
    free (expint);
    free (price);
    free (match);
    free (R);
    free (C);
  }

  MPI_Finalize ();
  return 0;
}