/****************************************************************
 *                                                              *
 *    Hypercubic lattice layout routines                        *
 *    Based on MILC lattice QCD code, pretty much modified      * 
 *                                                              *
 *    These determine the distribution of sites on nodes,       *
 *    and do the necessary setting up.                          *
 *                                                              *
 ***************************************************************/

#include "comdefs.h"
#include "generic.h"

/* static variables for node calculations */
int squaresize[NDIM];	/* dimensions of hypercubes */
int nsquares[NDIM];	/* number of hypercubes in each direction */

/* GLOBALS for communications; needed by com_XXX.c and block_lattice.c */
node_struct *allnodes;          /* structure for all nodes on this run */
comlist_struct *comlist;        /* gather pointer for all gathers */

#define swap(a,b) {register int t; t=a; a=b; b=t; }

void setup_layout( int siz[NDIM] );
void test_gather( lattice_struct *lat );
void make_gathers( lattice_struct *lat );

/***************************************************************
 *  BASIC CALL FOR SETUP
 *  
 *  setup_lattice(int size[NDIM]);
 */

void setup_lattice(int siz[NDIM])
{

  /* first, do the basic lattice layout */
  setup_layout( siz );

  /* then, set up the comm arrays */
  make_lattice_arrays( &lattice );

#ifdef MPI
  /* Initialize wait_array structures */
  initialize_wait_arrays();
#endif
}

/***************************************************************/

/* number of primes to be used in factorization */
#define NPRIMES 4
static int prime[NPRIMES] = {2,3,5,7};

/* Set up now squaresize and nsquares - arrays 
 * Print info to stdout as we proceed 
 */

void setup_layout( int siz[NDIM] )
{
  int n,i,j,dir,nfactors[NPRIMES];

  if(mynode()==0){
    printf(" LAYOUT OF THE LATTICE:\n %d dimensions, layout options: ",NDIM);
#ifdef GRAYCODE
    printf("GRAYCODE ");
#endif
#ifdef EVENFIRST
    printf("EVENFIRST ");
#endif
    printf("\n");
    fflush(stdout);
  }

  /* reset the blocking level (just in case) */
  foralldir(dir) current_blocking_level[dir] = 0;

  /* static global */
  this_node = mynode();

  /* Figure out dimensions of rectangle */

  lattice.volume = 1;
  foralldir(dir) {
    nsquares[dir] = 1;
    squaresize[dir] = lattice.size[dir] = siz[dir];
    lattice.volume *= lattice.size[dir];
  }

  /* store the baseline */
  base_lattice = lattice;

  if (lattice.volume % numnodes()) {
    printf0(" No hope of laying out the lattice using %d nodes\n",numnodes());
    finishrun();
  }

  /* Factorize the node number in primes
   * These factors must be used in slicing the lattice!
   */
  i = numnodes(); 
  for (n=0; n<NPRIMES; n++) {
    nfactors[n] = 0;
    while (i%prime[n] == 0) { i/=prime[n]; nfactors[n]++; }
  }
  if (i != 1) {
    printf0(" Cannot factorize %d nodes with primes up to %d\n",numnodes(),prime[NPRIMES-1]);
    finishrun();
  }
  
  for (n=NPRIMES-1; n>=0; n--) for(i=0; i<nfactors[n]; i++) {
    /* figure out which direction to divide -- start from the largest prime, because
     * we don't want this to be last divisor! (would probably wind up with size 1) 
     */

    /* find largest divisible dimension of h-cubes 
     *   - start the division from direction 0, because
     * these are likely to be within the same node!  */
    for(j=1,dir=0; dir<NDIM;dir++)
      if( squaresize[dir]>j && squaresize[dir]%prime[n] == 0 ) j=squaresize[dir];
    
    /* if one direction with largest dimension has already been
       divided, divide it again.  Otherwise divide first direction
       with largest dimension. */

    for (dir=0; dir<NDIM;dir++)
      if( squaresize[dir]==j && nsquares[dir]>1 && 
	  squaresize[dir]%prime[n] == 0) break;

    /* not previously sliced, take one direction to slice */
    if (dir >= NDIM) for (dir=0; dir<NDIM; dir++)
      if( squaresize[dir]==j && squaresize[dir]%prime[n] == 0) break;

    if (dir >= NDIM) {
      /* This cannot happen! */
      printf("CANNOT HAPPEN! in layout.c\n");
      finishrun();
    }

    /* Now slice it */
    squaresize[dir] /= prime[n]; nsquares[dir] *= prime[n];

  }
  
  if (mynode() == 0) {
    printf(" Processor layout: ");
    foralldir(dir) {
      if (dir > 0) printf(" x ");
      printf("%d",nsquares[dir]);
    }
    printf("\n Sites on node: ");
    foralldir(dir) {
      if (dir > 0) printf(" x ");
      printf("%d",squaresize[dir]);
    }
    printf("\n");
  }
}

/**************** Get the node number for (BLOCKED) coordinates */
int node_number(int loc[NDIM])
{
  register int i,dir;

  i = (loc[NDIM-1] << current_blocking_level[NDIM-1]) / squaresize[NDIM-1];
  for (dir=NDIM-2; dir>=0; dir--) {
    i = i*nsquares[dir] + 
      ((loc[dir] << current_blocking_level[dir]) / squaresize[dir]);
  }

#ifdef GRAYCODE
  return( i ^ (i>>1) );	/* Gray code of i */
#else
  return( i );
#endif
}

/************** fast routine for clarifying if we're on THIS node */

int is_on_node(int loc[NDIM])
{
  register int d,dir;

  foralldir(dir) {
    d = loc[dir] - node.xmin[dir];
    if (d < 0 || d >= node.nodesize[dir] ) return(0);
  }
  return(1);
}

/************** give site index for ON NODE sites */

int node_index(int loc[NDIM], node_struct *node)
{
  int dir,l,i,s;
  
  i = l = loc[NDIM-1] - node->xmin[NDIM-1];
  s = loc[NDIM-1];
  for (dir=NDIM-2; dir>=0; dir--) {
    l = loc[dir] - node->xmin[dir];
    i = i*node->nodesize[dir] + l;
    s += loc[dir];
  }
  
  /* now i contains the `running index' for site */
#ifdef EVENFIRST
  if (s%2 == 0) return( i/2 );    /* even site index */
  else return( i/2 + node->evensites );  /* odd site */
#else
  return( i );
#endif
}

/******************************************************
 * routines for stepping through the lattice in 
 * coordinates, as in
 * #define forallcoordinates(x) \
 *  for(zero_arr(x); is_coord(x,&lattice); step_coord(x,&lattice) )
 */

void zero_arr(int x[NDIM]) { register int d; foralldir(d) x[d] = 0; }

int is_allowed_coord(int x[NDIM],lattice_struct *l) 
{ 
  int d,i;
  i = 1;
  foralldir(d) i = (i && (x[d] >= 0) && (x[d] < l->size[d]));
  return(i);
}

void step_coord(int x[NDIM],lattice_struct *l) 
{
  int d;

  for(d=0; d<NDIM && ++(x[d]) >= l->size[d]; x[d++] = 0) ;

  /* check if the lattice is 'full' */
  if (d >= NDIM) x[NDIM-1] = l->size[NDIM-1];
}

    
/************** 
 * set up the node structure for all of the nodes in
 * the run (for BLOCKED coordinates).
 */

void setup_node( int loc[NDIM], node_struct *n )
{
  register int offset,dir,blev,l,c0,c1,s;

  n->sites = 1;
  s = 0;
  foralldir(dir) {
    blev = 1 << current_blocking_level[dir];
    l = loc[dir] << current_blocking_level[dir];  /* normalized coord */
    offset = l % squaresize[dir];         /* normalized coord from node 'origin' */
    c0 = l - offset;                      /* coordinate of the origin */
    c1 = c0 + squaresize[dir] - 1;        /* coordinate of the last point */
    
    /* calculate the coordinate of the first blocked point on the
     * node.  If the origin is divisible by the blocking factor,
     * then it belongs to the blocked lattice and the coordinate is
     * just l0 / blev.  However, if not, then the first blocked point
     * is l0 / blev + 1.
     */
    if (c0 % blev == 0) n->xmin[dir] = c0 >> current_blocking_level[dir];
    else n->xmin[dir] = (c0 >> current_blocking_level[dir]) + 1;
    
    /* Now the coordinate of the last blocked point.  This is 
     * always c1/blev, regardless if it is divisible or not.
     */

    c1 = c1 >> current_blocking_level[dir];
    
    /* now the length of the blocked lattice */
    n->nodesize[dir] = c1 - n->xmin[dir] + 1;
    
    /* need to accumulate size */
    n->sites *= n->nodesize[dir];
    /* and parity of the origin */
    s += n->xmin[dir];
  }

  if ( n->sites % 2 ) {
    /* now odd sized node */
    if ( s % 2 == 0) n->evensites = n->sites/2 + 1;
    else n->evensites = n->sites/2;
    n->oddsites = n->sites - n->evensites;
  } else {
    n->evensites = n->oddsites = n->sites/2;
  }
}

/************************************************************
 * set up the node struct for all nodes
 */

node_struct * setup_nodes(lattice_struct *lat)
{
  int i,l,d,n,x[NDIM];
  node_struct *p;

  /* allocate the node array */
  p = (node_struct *)memalloc( l=numnodes(), sizeof(node_struct) );
  for (i=0; i<l; i++) p[i].sites = -1;   /* flag as not initialized */ 

  forallcoordinates(x) {
    /* now x contains the coordinates of the site */
    n = node_number(x);  /* which node? */
    /* set up the node structure, if not done before */
    if (p[n].sites < 0) setup_node(x,p+n);
  }

  for (i=0; i<l; i++) {
    /* not initialized, reset */
    if (p[i].sites == -1) {
      p[i].sites = p[i].evensites = p[i].oddsites = 0;
      foralldir(d) {
	p[i].xmin[d] = p[i].nodesize[0] = 0;
      }
    }

    /* find now the neighbour node indices */
    if (p[i].sites) {
      int xn[NDIM];
      /* bottom corner coordinate of node */
      foralldir(d) xn[d] = x[d] = p[i].xmin[d];
      foralldir(d) {
	/* directions down */
	xn[d] = (x[d] - 1 + lattice.size[d]) % lattice.size[d];
	p[i].down_node[d] = node_number(xn);
	xn[d] = x[d];
      }
      /* upper corner coordinate of node */
      foralldir(d) xn[d] = x[d] = p[i].xmin[d] + p[i].nodesize[d];
      foralldir(d) {
	/* directions up */
	xn[d] = (x[d] + 1) % lattice.size[d];
	p[i].up_node[d] = node_number(xn);
	xn[d] = x[d];
      }
    }
  }
  return (p);
}

/************************************************************/


void make_lattice_arrays(lattice_struct * l)
{
  int x[NDIM],i,d,j,p;

  /* First, set up the node structure */

  allnodes = setup_nodes(l);
  node = allnodes[ this_node ];

  /* Setup the SITE array */

  site = (site_struct *)memalloc(node.sites+3, sizeof(site_struct));

  i = 0;  /* index to sites */
  forallcoordinates(x) {
    /* now x contains the coordinates of the site */
    
    if (is_on_node(x)) {
      j = node_index(x,&node);    /* array index to site */

      /* set the site */
      p = 0;
      foralldir(d) {
	site[j].x[d] = x[d];
	p += x[d];
      }
      if (p % 2 == 0) site[j].parity = EVEN; 
      else site[j].parity = ODD;
      site[j].index = i;
    } /* else printf("Offsite!\n"); */
    i++;
  }

  /* and then, set up the neighbour arrays and bookkeeping */

  make_gathers(l);

  /* and test it */
  test_gather(l);

}

/***************************************************************************
 * A long routine to make the gather datastructs (comlists) 
 * and neighbour arrays
 */

void make_gathers( lattice_struct *lat )
{
  int d,i,j,k,c_offset,x[NDIM];
  int *nodes, *index, *parity, *here, *itmp, num;
  int par,n,off,od;
  receive_struct **r,*p;
  send_struct **s,*q;
  
  /* First, allocate neighbour arrays */
  foralldir(d) {
    neighb[d]          = (int *)memalloc(node.sites, sizeof(int));
    neighb[opp_dir(d)] = (int *)memalloc(node.sites, sizeof(int));
  }

  /* allocate comlist structure */

  comlist = (comlist_struct *)memalloc(NDIRS, sizeof(comlist_struct));

  /* work arrays, will be released later */
  nodes = (int *)memalloc(node.sites, sizeof(int));  /* node number */
  index = (int *)memalloc(node.sites, sizeof(int));  /* index on node */
  parity = (int *)memalloc(node.sites, sizeof(int)); /* and parity */
  here = (int *)memalloc(node.sites, sizeof(int));   /* index of original site */
  itmp = (int *)memalloc(node.sites, sizeof(int));   /* temporary index array */

  c_offset = node.sites;  /* current offset in arrays */

  /* now, construct nn-gather to direction d */

  for (d=0; d<NDIRS; d++) {

    /* first pass over sites */
    num = 0;  /* number of sites off node */
    forallsites(i) {
      foralldir(j) x[j] = site[i].x[j];
      if (is_up_dir(d)) {
	x[d] = (x[d] + 1) % lat->size[d];   /* neighbour of site */
      } else {
	k = opp_dir(d);
	x[k] = (x[k] - 1 + lat->size[k]) % lat->size[k];   /* neighbour of site */
      }
      if (is_on_node(x)) neighb[d][i] = node_index(x,&node); 
      else {
	nodes[num] = node_number(x);
	index[num] = node_index(x, allnodes + nodes[num] );
	parity[num] = site[i].parity;  /* parity of THIS */
	here[num] = i;
	num++;
      }
    }
    
    comlist[d].n_receive = 0;
    if (num > 0) {
      /* now, get the number of nodes to be gathered from */
      for (i=0; i<num; i++) {

	/* chase the list until the node found */
	for (j=0, r=&(comlist[d].from_node); j<comlist[d].n_receive && 
	       nodes[i] != (*r)->node; j++) r = &((*r)->next); 
	if (j == comlist[d].n_receive) { 
	  /* NEW NODE to receive from */
	  comlist[d].n_receive++;  
	  (*r) = p = (receive_struct *)memalloc(1,sizeof(receive_struct));
	  /* and fill in the node structure */
	  p->node = nodes[i];
	  p->n = 1; /* first site */
	  p->n_even = p->n_odd = 0;
	  if ( parity[i] == EVEN ) p->n_even = 1;  else p->n_odd = 1;
	  p->next = NULL;
	} else {
	  /* add to OLD NODE */
	  p = *r;
	  p->n ++;
	  if ( parity[i] == EVEN ) p->n_even ++;  else p->n_odd ++;
	}
      }
      
      /* Calculate the offsets for the gathers */
      for (j=0, p=comlist[d].from_node; j<comlist[d].n_receive; j++, p = p->next) {
	p->offset = c_offset;
	c_offset += p->n;  /* and increase the offset */
      }

      /* and NOW, finish the NEIGHBOR array */
      
      for (j=0, p=comlist[d].from_node; j<comlist[d].n_receive; j++, p = p->next) {
	/* Now, accumulate the locations to itmp-array, and sort the
	 * array according to the index of the sending node .
	 * First even neighbours 
	 */
	for (par=EVEN; par<=ODD; par++) {
	  for (n=i=0; i<num; i++) if (nodes[i] == p->node && parity[i] == par) {
	    itmp[n++] = i;
	    /* bubble sort the tmp-array */
	    for (k=n-1; k > 0 && index[itmp[k]] < index[itmp[k-1]]; k--) 
	      swap( itmp[k], itmp[k-1] );
	  }
	  off = p->offset;
	  if (par == ODD) off += p->n_even;
	  /* finally, root indices according to offset */
	  for (k=0; k<n; k++) neighb[d][here[itmp[k]]] = off + k;
	}
      }
    } /* num > 0 */
	
    /* receive done, now opposite send. This is just the gather
     * inverted
     */
      
    od = opp_dir(d);
    comlist[od].n_send = comlist[d].n_receive;

    if (num > 0) {
      p = comlist[d].from_node;
      for (j=0, s=&(comlist[od].to_node); j<comlist[od].n_send; 
	   j++, s = &((*s)->next), p = p->next) {
	(*s) = q = (send_struct *)memalloc(1,sizeof(send_struct));
	q->node   = p->node;
	q->n      = p->n;
	q->n_even = p->n_odd;     /* Note the swap !  even/odd refers to type of gather */
	q->n_odd  = p->n_even;
	q->next   = NULL;
	q->sitelist = (int *)memalloc(q->n, sizeof(int));
	
	/* now, initialize sitelist -- Now, we first want ODD parity, since
	 * this is what even gather asks for! 
	 */
	
	for (n=0,par=ODD; par>=EVEN; par--) {
	  for (i=0; i<num; i++) if (nodes[i] == q->node && parity[i] == par) {
	    (q->sitelist)[n++] = here[i];
	  }
	  if (par == ODD && n != q->n_even) halt("Parity odd error 3");
	  if (par == EVEN && n != q->n) halt("Parity even error 3");
	}
      }
    }
  } /* directions */

  free(nodes); 
  free(index);
  free(parity); 
  free(here); 
  free(itmp);

  /* Finally, set the site to the final offset (better be right!) */
  node.latfield_size = c_offset;

}


/************************************************************************
 * Do some test to validate the correctness of the gather
 */

typedef struct t {
  int x[NDIM],parity;
} tst_struct;


void gather_test_error( char *abuse, int dir, tst_struct *a, 
			tst_struct *n, int par )
{
  int l;
  
  printf(" *** %s, parity %d, from dir %d: ( ",abuse,par,dir);
  foralldir(l) printf("%d ",a->x[l]);
  printf(") -> ( ");
  foralldir(l) printf("%d ",n->x[l]);
  printf("), parity %d -> %d\n",a->parity,n->parity);
}


void test_gather( lattice_struct *lat )
{
  int i,d,k,j,n,off,dir,n_err,par,checkparity;
  tst_struct *a;
  msg_tag *tag[NDIM];

  a = new_latfield( tst_struct );

  /* ignore parity if blocked lattice - usually OK */
  checkparity = 1;
  foralldir(d) if (current_blocking_level[d]) checkparity = 0;

  n_err = 0;
  for (k=0; k<2; k++) {
    for (par=EVEN; par<=EVENODD; par++) {

      forallsites(i) {
	foralldir(d) a[i].x[d] = site[i].x[d];
	a[i].parity = site[i].parity;
      }

      foralldir(d) {
	if (k) dir = opp_dir(d); else dir = d;
	tag[d]  = start_get( a, dir, par );
      }

      foralldir(d) {	
	if (k) dir = opp_dir(d); else dir = d;

	wait_get(tag[d]);
	
	if (is_up_dir(dir)) off = 1; else off = lat->size[d] - 1;

	forparity(i,par) foralldir(j) {
	  n = nb(dir,i);
	  if (( j != d && a[n].x[j] != a[i].x[j]) ||
	      ( j == d && a[n].x[j] != ((a[i].x[j] + off) % lat->size[d]))
#ifndef IGNORE_PARITY
	      || (( a[i].parity != opp_parity(a[n].parity)) && checkparity ) 
#endif
	      ) {
	    if (n_err < 10)
	      gather_test_error("HALOO! Gather error",dir,a+i,a+n,par);
	    n_err ++;
	  }
	}
      }
    }
  }

  /* test scatter too - inverse.  Sensible only for EVEN or ODD */
  /* can be up or down */
  for (dir=0; dir<NDIRS; dir++) {
    int odir = opp_dir(dir);
    for (par=EVEN; par<=EVENODD; par++) {
      int opar = opp_parity(par);
      int error = 0;

      forallsites(i) foralldir(d) a[i].x[d] = 0;

      forparity(i,par) {
	/* put other parity (neighb) sites to data values */
	n = nb(dir,i);
	foralldir(d) {
	  if (d == dir) 
	    a[n].x[d] = ((site[i].x[d] + 1)%lat->size[d]);
	  else if (d == odir)
	    a[n].x[d] = ((site[i].x[d] - 1 + lat->size[d])%lat->size[d]);
	  else
	    a[n].x[d] = site[i].x[d];
	}
	a[n].parity = opp_parity(site[i].parity);
      }

      wait_put( start_put( a, dir, par ) );
		  
      forparity(i,opar) {
	error = 0;
#ifndef IGNORE_PARITY
	if (checkparity && a[i].parity != site[i].parity) error = 1;
#endif
	foralldir(d) if (a[i].x[d] != site[i].x[d]) error = 1;
	if (error) {
	  if (n_err < 10) gather_test_error("HALOO! Scatter error",
					      dir,a+i,a+nb(odir,i),par);
	  n_err ++; 
	}
	}
    } 
  }

  if (n_err > 0) halt(" Lattice layout error (BUG in com_mpi.c or layout.c)");
  else printf0(" Gather/Scatter tests passed\n");

  free_latfield(a);
}


/****************************************************************/

char *copy_latfield_func( char *f, int siz )
{
  char *t;

  t = new_latfield_size( siz );
  memcpy( t, f, (siz * node.latfield_size) );
  return( t );
}

/****************************************************************/

char *latfield_alloc(int size)
{
  char *t;
  char f[150];
  
  t = (char *)malloc(node.latfield_size * size 
		     + GATHER_STATUS_SIZE );
  if (t == NULL) {
    sprintf(f,"Could not allocate a latfield of %d chars",size);
    halt(f);
  }

#ifdef MPI
  gather_status_reset( t, size );
#endif
  return( t );
}


char *memalloc(int n, int size)
{
  char *t;
  char f[150];
  
  t = (char *)malloc(n * size);
  if (t == NULL) {
    sprintf(f,"Memalloc: could not allocate %d x %d bytes",n,size);
    halt(f);
  }
  return( t );
}


void *halt(char *s)
{
  printf("*** %s\n",s);
  terminate(0);
  return((void *)NULL);
}


void time_stamp(char *msg)
{
  time_t time_stamp;

  if (this_node == 0) {
    time(&time_stamp);
    if (msg != NULL) printf("%s",msg);
    printf("%s\n", ctime(&time_stamp));
    fflush(stdout);
  }
}

