adding NamedMatrix.pprint() function

adding NamedMatrix.pprint() function

use LinearAlgebra,
    LinearAlgebra.Sparse,
    NumSuch,
    Random;

config param batchsize = 10000;

/*
 A matrix of data with named rows and columns.
 */
class NamedMatrix {
  var D: domain(2),
      SD = CSRDomain(D),
      X: [SD] real,  // the actual data
      rows: BiMap = new BiMap(),
      cols: BiMap = new BiMap();

   proc init() {
     this.complete(); // NAMEDMATRIX DOESNT EXPLICITLY INHERIT ANYTHING SO IS THIS NECESSARY??
   }

   proc init(X) {
     this.D = {X.domain.dim(1), X.domain.dim(2)};
     this.complete();
     this.loadX(X);
   }

   proc init(X, names: [] string) {
     this.init(X);
     try! this.setRowNames(names);
     try! this.setColNames(names);
   }

   proc init(X, rownames: [] string, colnames: [] string) {
     this.init(X);
     try! this.setRowNames(rownames);
     try! this.setColNames(colnames);
   }

   proc init(rownames: [] string, colnames: [] string) {
     this.D = {1..rownames.size, 1..colnames.size};
     this.complete();
     try! this.setRowNames(rownames);
     try! this.setColNames(colnames);
   }

   proc init(N: NamedMatrix) {
     this.init(N.X);
     this.rows = N.rows;
     this.cols = N.cols;
   }

   proc init(X, rows: BiMap, cols: BiMap) {
     this.init(X);
     this.rows = rows;
     this.cols = cols;
   }
}

/*
Returns the number of rows in the matrix frame
 */
proc NamedMatrix.nrows() {
  return X.domain.dim(1).size;
}

/*
Returns the number of columns in the matrix frame
 */
proc NamedMatrix.ncols() {
  return X.domain.dim(2).size;
}

/*
 Finds the max along the row indicated by row number
 */

proc NamedMatrix.rowMax(i: int) {
  return aMax(this.X, 1)[i];
}
proc NamedMatrix.rowMax(s: string) {
  return this.rowMax(this.rows.get(s));
}

proc NamedMatrix.rowMax() {
  return aMax(this.X, 1);
}

proc NamedMatrix.colMax(i: int) {
  return aMax(this.X, 2)[i];
}
proc NamedMatrix.colMax(s: string) {
  return this.colMax(this.cols.get(s));
}

proc NamedMatrix.colMax() {
  return aMax(this.X, 2);
}

proc NamedMatrix.rowArgMax(i: int) {
  return argMax(this.X, 1)[i];
}
proc NamedMatrix.rowArgMax(s: string) {
  return this.rowArgMax(this.rows.get(s));
}

proc NamedMatrix.rowArgMax() {
  return argMax(this.X, 1);
}

proc NamedMatrix.colArgMax(i: int) {
  return argMax(this.X, 2)[i];
}
proc NamedMatrix.colArgMax(s: string) {
  return this.colArgMax(this.cols.get(s));
}

proc NamedMatrix.colArgMax() {
  return argMax(this.X, 2);
}

proc NamedMatrix.rowMin(i: int) {
  return aMin(this.X, 1)[i];
}
proc NamedMatrix.rowMin(s: string) {
  return this.rowMin(this.rows.get(s));
}

proc NamedMatrix.rowMin() {
  return aMin(this.X, 1);
}

proc NamedMatrix.colMin(i: int) {
  return aMin(this.X, 2)[i];
}
proc NamedMatrix.colMin(s: string) {
  return this.colMin(this.cols.get(s));
}

proc NamedMatrix.colMin() {
  return aMin(this.X, 2);
}

proc NamedMatrix.rowArgMin(i:int) {
  return argMin(this.X, 1)[i];
}
proc NamedMatrix.rowArgMin(s: string) {
  return this.rowArgMin(this.rows.get(s));
}

proc NamedMatrix.rowArgMin() {
  return argMin(this.X, 1);
}

proc NamedMatrix.colArgMin(i) {
  return argMin(this.X, 2)[i];
}
proc NamedMatrix.colArgMin(s: string) {
  return this.colArgMin(this.cols.get(s));
}

proc NamedMatrix.colArgMin() {
  return argMin(this.X, 2);
}

/*
 Prints the matrix with row / column headers, not recommended for large matrices.
 Also width of entries should not exceed about 8 places.
 */
proc NamedMatrix.pprint() {
  write("     ");
  for c in this.cols.sorted() {
    write("       %5s  |".format(c[1]));
  }
  write("\n");
  for r in this.rows.sorted() {
    write(" %3s ".format(r[1]));
    for c in this.cols.sorted() {
      write(" %13.3dr ".format(this.get(r[2], c[2])));
    }
    write("\n");
  }

}

//proc NamedMatrix.argRowMax()

/*
Loads the data from X into the internal array, also called X.  We call them all X to keep it clear.

:arg real[]: Array representing the matrix
 */
proc NamedMatrix.loadX(X:[], shape: 2*int =(-1,-1)) {
  if shape(1) > 0 && shape(2) > 0 {
    this.D = {1..shape(1), 1..shape(2)};
  }
  for (i,j) in X.domain {
    this.SD += (i,j);
    this.X(i,j) = X(i,j);
  }
}

/*
Sets the row names for the matrix X
 */
proc NamedMatrix.setRowNames(rn:[]): string throws {
  //if (X.domain.dim(1).size > 0) && (rn.size != X.domain.dim(1).size) {
  //if (1 > 0) && (rn.size != X.domain.dim(1).size) {
  if rn.size != X.domain.dim(1).size {
    const err = new DimensionMatchError(expected = X.domain.dim(1).size, actual=rn.size);
    throw err;
    return "";
  } else {
    for i in 1..rn.size {
      this.rows.add(rn[i]);
    }
    return this.rows;
  }
}

/*
Sets the column names for the matrix X
*/
proc NamedMatrix.setColNames(cn:[]): string throws {
  if cn.size != X.domain.dim(2).size {
    const err = new DimensionMatchError(expected = X.domain.dim(2).size, actual=cn.size);
    throw err;
    return "";
  } else {
    for i in 1..cn.size {
      this.cols.add(cn[i]);
    }
  }
  return this.cols;
}

/*
 Gets the value of the (i,j) entry of the matrix X in the NamedMatrix
 */
proc NamedMatrix.get(i: int, j: int) {
   return this.X(i,j);
 }

 /*
 Get the values using the named rows and columns
  */
proc NamedMatrix.get(f: string, t: string) {
    return this.X(rows.get(f), cols.get(t));
 }

/*
 Set the value of this.X(i,j) by index row
 */
proc NamedMatrix.set(i: int, j: int, w: real) {
    if this.SD.member((i,j)) {
      this.X(i,j) = w;
    } else {
      this.SD += (i,j);
      this.X(i,j) = w;
    }
    return w;
}

/*
Set the values using the row and column names
*/
proc NamedMatrix.set(f: string, t: string, w: real) {
  return this.set(rows.get(f), cols.get(t), w);
}

/*
Update the value in X(i,j) with `w`
 */
proc NamedMatrix.update(i: int, j: int, w: real) {
  const x = this.get(i,j);
  return this.set(i,j, x + w);
}

/*
Update the value in X(i,j) with `w` using the row column names. This adds
`w` to the current value. To replace the current value, use `.set()`
 */
proc NamedMatrix.update(f: string, t: string, w: real) {
  const x = this.get(rows.get(f),cols.get(t));
  return this.set(rows.get(f),cols.get(t), x + w);
}

/*
 Removes the matrix entry `(i,j)` from the matrix and the underlying domain
 by row, column ids
 */
proc NamedMatrix.remove(i: int, j:int) {
  this.SD -= (i,j);
}

/*
 Removes the matrix entry `(i,j)` from the matrix and the underlying domain
 by row, column names
 */
proc NamedMatrix.remove(i: string, j:string) {
  this.remove(this.rows.get(i), this.cols.get(j));
}

/*
Find the number of non-zeroes in the matrix
 */
proc NamedMatrix.nnz() {
  var i: int = 0;
  for ij in this.X.domain {
    i += 1;
  }
  return i;
}

/*
Calculates the sparsity of the matrix: Number of entries / frame size
 */
proc NamedMatrix.sparsity() {
  const d = this.X.shape[1]:real * this.X.shape[2]: real;
  return this.nnz():real / d;
}

/*
 Takes a index (i,j) tuple and turns it into the sequential entry number on the
 given matrix.
 */
proc NamedMatrix.grid2seq(ij: 2*int) {
  return this.ncols() * ij[1] + ij[2] - this.ncols();
}
/*
 Takes a index values `i,j` and turns it into the sequential entry number on the
 given matrix.
 */
proc NamedMatrix.grid2seq(i: int, j: int) {
  return this.grid2seq((i,j));
}

/*
 Takes a sequential value and gives the row / column indices of that entry
 */
proc NamedMatrix.seq2grid(k: int) {
  return (divfloor(k, this.ncols())+1, k % this.ncols());
}

proc NamedMatrix.ntropic(N: NamedMatrix) {
  var T: NamedMatrix = new NamedMatrix(X = tropic(this.X, N.X), this.rows, N.cols);
  return T;
}

/*
 Multiplies the current NamedMatrix `X` against the argument `Y`, but first it aligns
 the names of `X.cols` with `Y.rows`.  Returns an appropriately named NamedMatrix
 :arg NamedMatrix Y:
 */
proc NamedMatrix.ndot(N: NamedMatrix) {
  var C: NamedMatrix = new NamedMatrix(X = dot(this.X,N.X), this.rows, N.cols);
  return C;
}

/*
 Routine to take two NamedMatrices and multiply them along just the row/column
 intersections.  If X has cols [red yellow blue green] and Y has rows [red blue green]
 this function produces a NamedMatrix with rows [red blue green]
 */
proc NamedMatrix.alignAndMultiply(Y: NamedMatrix) {
    var rcOverlap: domain(string) = this.cols.keys & Y.rows.keys;

    var xSD: sparse subdomain(this.D) dmapped CS(),
        ySD: sparse subdomain(Y.D) dmapped CS(),
        xX: [xSD] real,  // the actual data
        yX: [ySD] real,  // the actual data
        xrows: BiMap = new BiMap(),
        yrows: BiMap = new BiMap(),
        xcols: BiMap = new BiMap(),
        ycols: BiMap = new BiMap();

    for (left_row, right_col) in zip(this.D.dim(2), Y.D.dim(1)) {
      for rc in rcOverlap {
        const j = this.cols.get(rc);
        if this.SD.member(left_row, j) {
          xSD += (left_row, j);
          xX[left_row, j] = this.X[left_row, j];
        }
        if Y.SD.member(j, right_col) {
          ySD += (j, right_col);
          yX[j, right_col] = Y.X[j, right_col];
        }
      }
    }
    const z = xX.dot(yX);
    var n = new NamedMatrix(X=z);
    n.rows = this.rows;
    n.cols = Y.cols;
    return n;
}

/*
 Designed to be similar to SciPy's `amax() <https://docs.scipy.org/doc/numpy/reference/generated/numpy.amax.html>`_
 Useful on sparse matrices

 We'll get around to documenting here it eventually
 */
proc aMax(X:[], axis = 0) {
  var is: domain(1) = 1..max(X.domain.dim(1).size,X.domain.dim(2).size);
  var sis: sparse subdomain(is);
  var maxima: [sis] real;
  if axis == 0 {
    var wDom: domain(1) = 1..X.domain.size;
    var ws = reshape(X, wDom);
    sis += 1;
    maxima[1] = (max reduce ws);
  }
  if axis == 2 {
    is = X.domain.dim(1);
    for i in is {
      var wDom: domain(1);
      var ws: [wDom] real;
      for a in transpose(X).domain.dimIter(2, i) {
        sis += i;
        ws.push_back(X(a,i));
      }
      if sis.member(i) {
        maxima[i] = max reduce ws;
      }
    }
  } else if axis == 1 {
    is = X.domain.dim(2);
    for i in is {
      var wDom: domain(1);
      var ws: [wDom] real;
      for a in X.domain.dimIter(2, i) {
        sis += i;
        ws.push_back(X(i,a));
      }
      if sis.member(i) {
        maxima[i] = max reduce ws;
      }
    }
  }
  return maxima;
}

/*
 Designed to be similar to SciPy's `amin() <https://docs.scipy.org/doc/numpy/reference/generated/numpy.amin.html>`_
 Useful on sparse matrices.

 We'll get around to documenting here it eventually
 */
proc aMin(X:[], axis: int) {
  var is: domain(1) = 1..max(X.domain.dim(1).size,X.domain.dim(2).size);
  var sis: sparse subdomain(is);
  var minima: [sis] real;
  if axis == 0 {
    var wDom: domain(1) = 1..X.domain.size;
    var ws = reshape(X, wDom);
    sis += 1;
    minima[1] = (min reduce ws);
  }
  if axis == 2 {
    is = X.domain.dim(1);
    for i in is {
      var wDom: domain(1);
      var ws: [wDom] real;
      for a in transpose(X).domain.dimIter(2, i) {
        sis += i;
        ws.push_back(X(a,i));
      }
      if sis.member(i) {
        minima[i] = min reduce ws;
      }
    }
  } else if axis == 1 {
    is = X.domain.dim(2);
    for i in is {
      var wDom: domain(1);
      var ws: [wDom] real;
      for a in X.domain.dimIter(2, i) {
        sis += i;
        ws.push_back(X(i,a));
      }
      if sis.member(i) {
        minima[i] = min reduce ws;
      }
    }
  }
  return minima;
}

/*
Provides the "where" of the min along an axis in a sparse matrix
 */
proc argMin(X:[], axis: int) {
  var is: domain(1) = 1..max(X.domain.dim(1).size,X.domain.dim(2).size);
  var sis: sparse subdomain(is);
  var minima: [sis] 2*int;
  if axis == 0 {
    sis += 1;
    minima[1] = miniLoc_(axis == 0, id = 1, X);
  } else if axis == 2 {
    for i in is {
      var idm = miniLoc_(1,i,X);
      if idm != (0,0) {
        sis += i;
        minima[i] = idm;
      }
    }
  } else if axis == 1 {
    for i in is {
      var idm = miniLoc_(2,i,X);
      if idm != (0,0) {
        sis += i;
        minima[i] = idm;
      }
    }
  }
  return minima;
}

/*
 Provides the "where" of a max along a given axis in a sparse matrix
 */
proc argMax(X:[], axis: int) {
  var is: domain(1) = 1..max(X.domain.dim(1).size,X.domain.dim(2).size);
  var sis: sparse subdomain(is);
  var maxima: [sis] 2*int;
  if axis == 0 {
    sis += 1;
    maxima[1] = maxiLoc_(axis == 0, id = 1, X);
  } else if axis == 2 {
    for i in is {
      var idm = maxiLoc_(1,i,X);
      if idm != (0,0) {
        sis += i;
        maxima[i] = idm;
      }
    }
  } else if axis == 1 {
    for i in is {
      var idm = maxiLoc_(2,i,X);
      if idm != (0,0) {
        sis += i;
        maxima[i] = idm;
      }
    }
  }
  return maxima;
}

// HELPER FUNCTIONS SHOULD NOT BE EXPOSED TO USER
proc miniLoc_(axis:int, id, X:[]) {
  var idm: (int,int),
      minimum: real = 1000000;
  if axis == 0 {
    for (i,j) in X.domain {
      if X(i,j) < minimum {
        minimum = X(i,j);
        idm = (i,j);
      }
    }
  }
  if axis == 1 {
    for a in transpose(X).domain.dimIter(2, id) {
      if X(a,id) < minimum {
        minimum = X(a,id);
        idm = (a,id);
      }
    }
  }
  if axis == 2 {
    for a in X.domain.dimIter(2, id) {
      if X(id,a) < minimum {
        minimum = X(id,a);
        idm = (id,a);
      }
    }
  }
  return idm;
}

proc maxiLoc_(axis:int, id, X:[]) {
  var idm: (int,int),
      maximum: real;
  if axis == 0 {
    for (i,j) in X.domain {
      if X(i,j) > maximum {
        maximum = X(i,j);
        idm = (i,j);
      }
    }
  }
  if axis == 1 {
    for a in transpose(X).domain.dimIter(2, id) {
      if X(a,id) > maximum {
        maximum = X(a,id);
        idm = (a,id);
      }
    }
  }
  if axis == 2 {
    for a in X.domain.dimIter(2, id) {
      if X(id,a) > maximum {
        maximum = X(id,a);
        idm = (id,a);
      }
    }
  }
  return idm;
}

proc sparseEquals(A:[], B:[]) {}

proc tropic(A:[],B:[]) { // UNDER CONSTRUCTION
  var dom: domain(2) = {A.domain.dim(1),B.domain.dim(2)};
  var sps = CSRDomain(dom);
  var BT = transpose(B);
  var T: [sps] real;
  for (i,j) in dom {
    var mini = 100000000: real;
    var wdom: domain(1);
    var wids: [wdom] int;
    for w in A.domain.dimIter(2,i) {
      wids.push_back(w);
    }
    for w2 in BT.domain.dimIter(2,j) {
      wids.push_back(w2);
    }
    var wDom: domain(1);
    var witness: [wDom] real;
    for w in wids {
      if A(i,w) > 0 && BT(j,w) > 0 {
        witness.push_back(A(i,w) + BT(j,w));
      }
    }
    if A.domain.member((i,j)) {
      witness.push_back(A(i,j));
    }
    if B.domain.member((i,j)) {
      witness.push_back(B(i,j));
    }
    mini = min reduce witness;
    if mini < 1000000000 {
      sps += (i,j);
      T(i,j) = mini;
    }
  }
  return T;
}

proc tropicLimit(B:[] real) {
 var R = tropic(B,B);
 for n in B.domain.dim(2) {
   var S = tropic(R,B);
   if S.domain != R.domain {
    R = S; // Intended to just reassign the handle "R" to the contents of "S" i.o.w. destructive update of R
   } else {
     break;
   }
 }
 return R;
}

/*
 Build a random sparse matrix.  Good for testing;
 */
proc generateRandomSparseMatrix(size: int, sparsity: real) {
  const D: domain(2) = {1..size, 1..size};
  var SD: sparse subdomain(D) dmapped CS();
  var R: [SD] real;
  var da: domain(1) = {1..size};
  var array: [da] int = 1..size;
  var dom: domain(1) = {1..0};
  var indices: [dom] 2*int;
  var N = floor(size*(1-sparsity)): int;
  forall i in array {
    forall j in array {
      indices.push_back((i,j));
    }
  }
  shuffle(indices);
  var sparseids = indices[1..N];
  SD.bulkAdd(sparseids);
  forall (i,j) in sparseids {
    R(i,j) = 1;
  }
  //ndices = zip(array1, array2)
  return R;
}

/*
 General sparsity function for any `real` matrix.
 */
proc sparsity(X) {
  const d = X.shape[1]:real * X.shape[2]: real;
  var i: real = 0.0;
  for ij in X.domain {
    i += 1.0;
  }
  return i / d;
}

/*
 Error classes
 */
class NumSuchError : Error {
  proc init() {
    super.init();
    this.complete();
  }
  proc message() {
    return "Generic NumSuch Error";
  }
}

/*
 Used to indidcate dimension mismatches on NamedMatrices and vectors
 */
class DimensionMatchError : NumSuchError {
  var expected: int,
      actual: int;

  proc init(expected: int, actual:int) {
    super.init();
    this.complete();
    this.expected = expected;
    this.actual = actual;
  }

  proc message() {
    return "Error matching dimensions.  Expected: " + this.expected + " Actual: " + this.actual;
  }
}

use LinearAlgebra,
    LinearAlgebra.Sparse,
    NumSuch,
    Random;

config param batchsize = 10000;

/*
 A matrix of data with named rows and columns.
 */
class NamedMatrix {
  var D: domain(2),
      SD = CSRDomain(D),
      X: [SD] real,  // the actual data
      rows: BiMap = new BiMap(),
      cols: BiMap = new BiMap();

   proc init() {
     this.complete(); // NAMEDMATRIX DOESNT EXPLICITLY INHERIT ANYTHING SO IS THIS NECESSARY??
   }

   proc init(X) {
     this.D = {X.domain.dim(1), X.domain.dim(2)};
     this.complete();
     this.loadX(X);
   }

   proc init(X, names: [] string) {
     this.init(X);
     try! this.setRowNames(names);
     try! this.setColNames(names);
   }

   proc init(X, rownames: [] string, colnames: [] string) {
     this.init(X);
     try! this.setRowNames(rownames);
     try! this.setColNames(colnames);
   }

   proc init(rownames: [] string, colnames: [] string) {
     this.D = {1..rownames.size, 1..colnames.size};
     this.complete();
     try! this.setRowNames(rownames);
     try! this.setColNames(colnames);
   }

   proc init(N: NamedMatrix) {
     this.init(N.X);
     this.rows = N.rows;
     this.cols = N.cols;
   }

   proc init(X, rows: BiMap, cols: BiMap) {
     this.init(X);
     this.rows = rows;
     this.cols = cols;
   }
}

/*
Returns the number of rows in the matrix frame
 */
proc NamedMatrix.nrows() {
  return X.domain.dim(1).size;
}

/*
Returns the number of columns in the matrix frame
 */
proc NamedMatrix.ncols() {
  return X.domain.dim(2).size;
}

/*
 Finds the max along the row indicated by row number
 */

proc NamedMatrix.rowMax(i: int) {
  return aMax(this.X, 1)[i];
}
proc NamedMatrix.rowMax(s: string) {
  return this.rowMax(this.rows.get(s));
}

proc NamedMatrix.rowMax() {
  return aMax(this.X, 1);
}

proc NamedMatrix.colMax(i: int) {
  return aMax(this.X, 2)[i];
}
proc NamedMatrix.colMax(s: string) {
  return this.colMax(this.cols.get(s));
}

proc NamedMatrix.colMax() {
  return aMax(this.X, 2);
}

proc NamedMatrix.rowArgMax(i: int) {
  return argMax(this.X, 1)[i];
}
proc NamedMatrix.rowArgMax(s: string) {
  return this.rowArgMax(this.rows.get(s));
}

proc NamedMatrix.rowArgMax() {
  return argMax(this.X, 1);
}

proc NamedMatrix.colArgMax(i: int) {
  return argMax(this.X, 2)[i];
}
proc NamedMatrix.colArgMax(s: string) {
  return this.colArgMax(this.cols.get(s));
}

proc NamedMatrix.colArgMax() {
  return argMax(this.X, 2);
}

proc NamedMatrix.rowMin(i: int) {
  return aMin(this.X, 1)[i];
}
proc NamedMatrix.rowMin(s: string) {
  return this.rowMin(this.rows.get(s));
}

proc NamedMatrix.rowMin() {
  return aMin(this.X, 1);
}

proc NamedMatrix.colMin(i: int) {
  return aMin(this.X, 2)[i];
}
proc NamedMatrix.colMin(s: string) {
  return this.colMin(this.cols.get(s));
}

proc NamedMatrix.colMin() {
  return aMin(this.X, 2);
}

proc NamedMatrix.rowArgMin(i:int) {
  return argMin(this.X, 1)[i];
}
proc NamedMatrix.rowArgMin(s: string) {
  return this.rowArgMin(this.rows.get(s));
}

proc NamedMatrix.rowArgMin() {
  return argMin(this.X, 1);
}

proc NamedMatrix.colArgMin(i) {
  return argMin(this.X, 2)[i];
}
proc NamedMatrix.colArgMin(s: string) {
  return this.colArgMin(this.cols.get(s));
}

proc NamedMatrix.colArgMin() {
  return argMin(this.X, 2);
}

//proc NamedMatrix.argRowMax()

/*
Loads the data from X into the internal array, also called X.  We call them all X to keep it clear.

:arg real[]: Array representing the matrix
 */
proc NamedMatrix.loadX(X:[], shape: 2*int =(-1,-1)) {
  if shape(1) > 0 && shape(2) > 0 {
    this.D = {1..shape(1), 1..shape(2)};
  }
  for (i,j) in X.domain {
    this.SD += (i,j);
    this.X(i,j) = X(i,j);
  }
}

/*
Sets the row names for the matrix X
 */
proc NamedMatrix.setRowNames(rn:[]): string throws {
  //if (X.domain.dim(1).size > 0) && (rn.size != X.domain.dim(1).size) {
  //if (1 > 0) && (rn.size != X.domain.dim(1).size) {
  if rn.size != X.domain.dim(1).size {
    const err = new DimensionMatchError(expected = X.domain.dim(1).size, actual=rn.size);
    throw err;
    return "";
  } else {
    for i in 1..rn.size {
      this.rows.add(rn[i]);
    }
    return this.rows;
  }
}

/*
Sets the column names for the matrix X
*/
proc NamedMatrix.setColNames(cn:[]): string throws {
  if cn.size != X.domain.dim(2).size {
    const err = new DimensionMatchError(expected = X.domain.dim(2).size, actual=cn.size);
    throw err;
    return "";
  } else {
    for i in 1..cn.size {
      this.cols.add(cn[i]);
    }
  }
  return this.cols;
}

/*
 Gets the value of the (i,j) entry of the matrix X in the NamedMatrix
 */
proc NamedMatrix.get(i: int, j: int) {
   return this.X(i,j);
 }

 /*
 Get the values using the named rows and columns
  */
proc NamedMatrix.get(f: string, t: string) {
    return this.X(rows.get(f), cols.get(t));
 }

/*
 Set the value of this.X(i,j) by index row
 */
proc NamedMatrix.set(i: int, j: int, w: real) {
    if this.SD.member((i,j)) {
      this.X(i,j) = w;
    } else {
      this.SD += (i,j);
      this.X(i,j) = w;
    }
    return w;
}

/*
Set the values using the row and column names
*/
proc NamedMatrix.set(f: string, t: string, w: real) {
  return this.set(rows.get(f), cols.get(t), w);
}

/*
Update the value in X(i,j) with `w`
 */
proc NamedMatrix.update(i: int, j: int, w: real) {
  const x = this.get(i,j);
  return this.set(i,j, x + w);
}

/*
Update the value in X(i,j) with `w` using the row column names. This adds
`w` to the current value. To replace the current value, use `.set()`
 */
proc NamedMatrix.update(f: string, t: string, w: real) {
  const x = this.get(rows.get(f),cols.get(t));
  return this.set(rows.get(f),cols.get(t), x + w);
}

/*
 Removes the matrix entry `(i,j)` from the matrix and the underlying domain
 by row, column ids
 */
proc NamedMatrix.remove(i: int, j:int) {
  this.SD -= (i,j);
}

/*
 Removes the matrix entry `(i,j)` from the matrix and the underlying domain
 by row, column names
 */
proc NamedMatrix.remove(i: string, j:string) {
  this.remove(this.rows.get(i), this.cols.get(j));
}

/*
Find the number of non-zeroes in the matrix
 */
proc NamedMatrix.nnz() {
  var i: int = 0;
  for ij in this.X.domain {
    i += 1;
  }
  return i;
}

/*
Calculates the sparsity of the matrix: Number of entries / frame size
 */
proc NamedMatrix.sparsity() {
  const d = this.X.shape[1]:real * this.X.shape[2]: real;
  return this.nnz():real / d;
}

/*
 Takes a index (i,j) tuple and turns it into the sequential entry number on the
 given matrix.
 */
proc NamedMatrix.grid2seq(ij: 2*int) {
  return this.ncols() * ij[1] + ij[2] - this.ncols();
}
/*
 Takes a index values `i,j` and turns it into the sequential entry number on the
 given matrix.
 */
proc NamedMatrix.grid2seq(i: int, j: int) {
  return this.grid2seq((i,j));
}

/*
 Takes a sequential value and gives the row / column indices of that entry
 */
proc NamedMatrix.seq2grid(k: int) {
  return (divfloor(k, this.ncols())+1, k % this.ncols());
}

proc NamedMatrix.ntropic(N: NamedMatrix) {
  var T: NamedMatrix = new NamedMatrix(X = tropic(this.X, N.X), this.rows, N.cols);
  return T;
}

/*
 Multiplies the current NamedMatrix `X` against the argument `Y`, but first it aligns
 the names of `X.cols` with `Y.rows`.  Returns an appropriately named NamedMatrix
 :arg NamedMatrix Y:
 */
proc NamedMatrix.ndot(N: NamedMatrix) {
  var C: NamedMatrix = new NamedMatrix(X = dot(this.X,N.X), this.rows, N.cols);
  return C;
}

/*
 Routine to take two NamedMatrices and multiply them along just the row/column
 intersections.  If X has cols [red yellow blue green] and Y has rows [red blue green]
 this function produces a NamedMatrix with rows [red blue green]
 */
proc NamedMatrix.alignAndMultiply(Y: NamedMatrix) {
    var rcOverlap: domain(string) = this.cols.keys & Y.rows.keys;

    var xSD: sparse subdomain(this.D) dmapped CS(),
        ySD: sparse subdomain(Y.D) dmapped CS(),
        xX: [xSD] real,  // the actual data
        yX: [ySD] real,  // the actual data
        xrows: BiMap = new BiMap(),
        yrows: BiMap = new BiMap(),
        xcols: BiMap = new BiMap(),
        ycols: BiMap = new BiMap();

    for (left_row, right_col) in zip(this.D.dim(2), Y.D.dim(1)) {
      for rc in rcOverlap {
        const j = this.cols.get(rc);
        if this.SD.member(left_row, j) {
          xSD += (left_row, j);
          xX[left_row, j] = this.X[left_row, j];
        }
        if Y.SD.member(j, right_col) {
          ySD += (j, right_col);
          yX[j, right_col] = Y.X[j, right_col];
        }
      }
    }
    const z = xX.dot(yX);
    var n = new NamedMatrix(X=z);
    n.rows = this.rows;
    n.cols = Y.cols;
    return n;
}

/*
 Designed to be similar to SciPy's `amax() <https://docs.scipy.org/doc/numpy/reference/generated/numpy.amax.html>`_
 Useful on sparse matrices

 We'll get around to documenting here it eventually
 */
proc aMax(X:[], axis = 0) {
  var is: domain(1) = 1..max(X.domain.dim(1).size,X.domain.dim(2).size);
  var sis: sparse subdomain(is);
  var maxima: [sis] real;
  if axis == 0 {
    var wDom: domain(1) = 1..X.domain.size;
    var ws = reshape(X, wDom);
    sis += 1;
    maxima[1] = (max reduce ws);
  }
  if axis == 2 {
    is = X.domain.dim(1);
    for i in is {
      var wDom: domain(1);
      var ws: [wDom] real;
      for a in transpose(X).domain.dimIter(2, i) {
        sis += i;
        ws.push_back(X(a,i));
      }
      if sis.member(i) {
        maxima[i] = max reduce ws;
      }
    }
  } else if axis == 1 {
    is = X.domain.dim(2);
    for i in is {
      var wDom: domain(1);
      var ws: [wDom] real;
      for a in X.domain.dimIter(2, i) {
        sis += i;
        ws.push_back(X(i,a));
      }
      if sis.member(i) {
        maxima[i] = max reduce ws;
      }
    }
  }
  return maxima;
}

/*
 Designed to be similar to SciPy's `amin() <https://docs.scipy.org/doc/numpy/reference/generated/numpy.amin.html>`_
 Useful on sparse matrices.

 We'll get around to documenting here it eventually
 */
proc aMin(X:[], axis: int) {
  var is: domain(1) = 1..max(X.domain.dim(1).size,X.domain.dim(2).size);
  var sis: sparse subdomain(is);
  var minima: [sis] real;
  if axis == 0 {
    var wDom: domain(1) = 1..X.domain.size;
    var ws = reshape(X, wDom);
    sis += 1;
    minima[1] = (min reduce ws);
  }
  if axis == 2 {
    is = X.domain.dim(1);
    for i in is {
      var wDom: domain(1);
      var ws: [wDom] real;
      for a in transpose(X).domain.dimIter(2, i) {
        sis += i;
        ws.push_back(X(a,i));
      }
      if sis.member(i) {
        minima[i] = min reduce ws;
      }
    }
  } else if axis == 1 {
    is = X.domain.dim(2);
    for i in is {
      var wDom: domain(1);
      var ws: [wDom] real;
      for a in X.domain.dimIter(2, i) {
        sis += i;
        ws.push_back(X(i,a));
      }
      if sis.member(i) {
        minima[i] = min reduce ws;
      }
    }
  }
  return minima;
}

/*
Provides the "where" of the min along an axis in a sparse matrix
 */
proc argMin(X:[], axis: int) {
  var is: domain(1) = 1..max(X.domain.dim(1).size,X.domain.dim(2).size);
  var sis: sparse subdomain(is);
  var minima: [sis] 2*int;
  if axis == 0 {
    sis += 1;
    minima[1] = miniLoc_(axis == 0, id = 1, X);
  } else if axis == 2 {
    for i in is {
      var idm = miniLoc_(1,i,X);
      if idm != (0,0) {
        sis += i;
        minima[i] = idm;
      }
    }
  } else if axis == 1 {
    for i in is {
      var idm = miniLoc_(2,i,X);
      if idm != (0,0) {
        sis += i;
        minima[i] = idm;
      }
    }
  }
  return minima;
}

/*
 Provides the "where" of a max along a given axis in a sparse matrix
 */
proc argMax(X:[], axis: int) {
  var is: domain(1) = 1..max(X.domain.dim(1).size,X.domain.dim(2).size);
  var sis: sparse subdomain(is);
  var maxima: [sis] 2*int;
  if axis == 0 {
    sis += 1;
    maxima[1] = maxiLoc_(axis == 0, id = 1, X);
  } else if axis == 2 {
    for i in is {
      var idm = maxiLoc_(1,i,X);
      if idm != (0,0) {
        sis += i;
        maxima[i] = idm;
      }
    }
  } else if axis == 1 {
    for i in is {
      var idm = maxiLoc_(2,i,X);
      if idm != (0,0) {
        sis += i;
        maxima[i] = idm;
      }
    }
  }
  return maxima;
}

// HELPER FUNCTIONS SHOULD NOT BE EXPOSED TO USER
proc miniLoc_(axis:int, id, X:[]) {
  var idm: (int,int),
      minimum: real = 1000000;
  if axis == 0 {
    for (i,j) in X.domain {
      if X(i,j) < minimum {
        minimum = X(i,j);
        idm = (i,j);
      }
    }
  }
  if axis == 1 {
    for a in transpose(X).domain.dimIter(2, id) {
      if X(a,id) < minimum {
        minimum = X(a,id);
        idm = (a,id);
      }
    }
  }
  if axis == 2 {
    for a in X.domain.dimIter(2, id) {
      if X(id,a) < minimum {
        minimum = X(id,a);
        idm = (id,a);
      }
    }
  }
  return idm;
}

proc maxiLoc_(axis:int, id, X:[]) {
  var idm: (int,int),
      maximum: real;
  if axis == 0 {
    for (i,j) in X.domain {
      if X(i,j) > maximum {
        maximum = X(i,j);
        idm = (i,j);
      }
    }
  }
  if axis == 1 {
    for a in transpose(X).domain.dimIter(2, id) {
      if X(a,id) > maximum {
        maximum = X(a,id);
        idm = (a,id);
      }
    }
  }
  if axis == 2 {
    for a in X.domain.dimIter(2, id) {
      if X(id,a) > maximum {
        maximum = X(id,a);
        idm = (id,a);
      }
    }
  }
  return idm;
}

proc sparseEquals(A:[], B:[]) {}

proc tropic(A:[],B:[]) { // UNDER CONSTRUCTION
  var dom: domain(2) = {A.domain.dim(1),B.domain.dim(2)};
  var sps = CSRDomain(dom);
  var BT = transpose(B);
  var T: [sps] real;
  for (i,j) in dom {
    var mini = 100000000: real;
    var wdom: domain(1);
    var wids: [wdom] int;
    for w in A.domain.dimIter(2,i) {
      wids.push_back(w);
    }
    for w2 in BT.domain.dimIter(2,j) {
      wids.push_back(w2);
    }
    var wDom: domain(1);
    var witness: [wDom] real;
    for w in wids {
      if A(i,w) > 0 && BT(j,w) > 0 {
        witness.push_back(A(i,w) + BT(j,w));
      }
    }
    if A.domain.member((i,j)) {
      witness.push_back(A(i,j));
    }
    if B.domain.member((i,j)) {
      witness.push_back(B(i,j));
    }
    mini = min reduce witness;
    if mini < 1000000000 {
      sps += (i,j);
      T(i,j) = mini;
    }
  }
  return T;
}

proc tropicLimit(B:[] real) {
 var R = tropic(B,B);
 for n in B.domain.dim(2) {
   var S = tropic(R,B);
   if S.domain != R.domain {
    R = S; // Intended to just reassign the handle "R" to the contents of "S" i.o.w. destructive update of R
   } else {
     break;
   }
 }
 return R;
}

/*
 Build a random sparse matrix.  Good for testing;
 */
proc generateRandomSparseMatrix(size: int, sparsity: real) {
  const D: domain(2) = {1..size, 1..size};
  var SD: sparse subdomain(D) dmapped CS();
  var R: [SD] real;
  var da: domain(1) = {1..size};
  var array: [da] int = 1..size;
  var dom: domain(1) = {1..0};
  var indices: [dom] 2*int;
  var N = floor(size*(1-sparsity)): int;
  forall i in array {
    forall j in array {
      indices.push_back((i,j));
    }
  }
  shuffle(indices);
  var sparseids = indices[1..N];
  SD.bulkAdd(sparseids);
  forall (i,j) in sparseids {
    R(i,j) = 1;
  }
  //ndices = zip(array1, array2)
  return R;
}

/*
 General sparsity function for any `real` matrix.
 */
proc sparsity(X) {
  const d = X.shape[1]:real * X.shape[2]: real;
  var i: real = 0.0;
  for ij in X.domain {
    i += 1.0;
  }
  return i / d;
}

/*
 Error classes
 */
class NumSuchError : Error {
  proc init() {
    super.init();
    this.complete();
  }
  proc message() {
    return "Generic NumSuch Error";
  }
}

/*
 Used to indidcate dimension mismatches on NamedMatrices and vectors
 */
class DimensionMatchError : NumSuchError {
  var expected: int,
      actual: int;

  proc init(expected: int, actual:int) {
    super.init();
    this.complete();
    this.expected = expected;
    this.actual = actual;
  }

  proc message() {
    return "Error matching dimensions.  Expected: " + this.expected + " Actual: " + this.actual;
  }
}


Removed dead code from this test that I forgot to kill
before committing.

git-svn-id: 88467cb1fb04b8a755be7e1ee1026be4190196ef@10721 3a8e244f-b0f2-452b-bcba-4c88e055c3ca

config var n = 9;

def main() {
  const D = [1..n, 1..n];

  var A: [D] real;

  coforall t in 1..2 {
    for i in t..n by 2 {
      for j in 1..n {
        A(i,j) += i + j/10.0;
      }
    }
  }

  writeln("A is:\n", A);
}

config var n = 9;

def loopit(t) {
  for i in t..n by 2 {
    for j in 1..n {
      yield (i,j);
    }
  }
}

def main() {
  const D = [1..n, 1..n];

  var A: [D] real;

  coforall t in 1..2 {
    for i in t..n by 2 {
      for j in 1..n {
        A(i,j) += i + j/10.0;
      }
    }
  }

  writeln("A is:\n", A);
}


Fixed the CyclicDom.setup method when the LocDoms already exist.

It was creating new LocCyclicDom classes and replacing the ones that were
already created, but instead it should just change the LocCyclicDom.myBlock
field. With this change to setup() plus an empty CyclicArray.dsiReallocate()
method, one more distribution robustness test is passing for Cyclic.

git-svn-id: 88467cb1fb04b8a755be7e1ee1026be4190196ef@16812 3a8e244f-b0f2-452b-bcba-4c88e055c3ca

use DSIUtil;

def _tupleOfZero(param size: int, type t) {
  var tmp: size * t;
  return tmp;
}

config param debugCyclicDist = false;
config param verboseCyclicDistWriters = false;

class Cyclic: BaseDist {
  param rank: int = 1;
  type idxType = int(64);

  const lowIdx: rank*idxType;
  const targetLocDom: domain(rank);
  const targetLocs: [targetLocDom] locale;

  const locDist: [targetLocDom] LocCyclic(rank, idxType);

  const tasksPerLocale = 1;
  var pid: int = -1;

  def Cyclic(param rank: int,
             type idxType = int(64),
             low: rank*idxType = _tupleOfZero(rank, idxType),
             targetLocales: [] locale = thisRealm.Locales,
             tasksPerLocale=0) {
    if rank == 1  {
      targetLocDom = [0..#targetLocales.numElements];
      targetLocs = targetLocales;
    } else if targetLocales.rank == 1 {
      const factors = _factor(rank, targetLocales.numElements);
      var ranges: rank*range;
      for param i in 1..rank {
        ranges(i) = 0..#factors(i);
      }
      targetLocDom = [(...ranges)];
      for (loc1, loc2) in (targetLocs, targetLocales) {
        loc1 = loc2;
      }
    } else {
      if targetLocales.rank != rank then
        compilerError("locales array rank must be one or match distribution rank");
      var ranges: rank*range;
      for param i in 1..rank do {
        var thisRange = targetLocales.domain.dim(i);
        ranges(i) = 0..#thisRange.length;
      }
      targetLocDom = [(...ranges)];
      targetLocs = reshape(targetLocales, targetLocDom);
    }

    for param i in 1..rank do
      lowIdx(i) = mod(low(i), targetLocDom.dim(i).length);

    if tasksPerLocale == 0 then
      this.tasksPerLocale = min reduce targetLocales.numCores;
    else
      this.tasksPerLocale = tasksPerLocale;
    coforall locid in targetLocDom {
      on targetLocs(locid) {
        locDist(locid) = new LocCyclic(rank, idxType, locid, this);
      }
    }
    if debugCyclicDist then
      for loc in locDist do writeln(loc);
  }

  def Cyclic(param rank, type idxType, other: Cyclic(rank, idxType)) {
    targetLocDom = other.targetLocDom;
    targetLocs = other.targetLocs;
    lowIdx = other.lowIdx;
    locDist = other.locDist;
    tasksPerLocale = other.tasksPerLocale;
  }

  def dsiClone() return new Cyclic(rank=rank, idxType=idxType, other=this);
}

def Cyclic.getChunk(inds, locid) {
  var sliceBy: rank*range(eltType=idxType, stridable=true);
  var locidtup: rank*int;
  if rank == 1 then
    locidtup(1) = locid;
  else
    locidtup = locid;
  for param i in 1..rank {
    var distStride = targetLocDom.dim(i).length;
    var loclowidx = mod(lowIdx(i) + locidtup(i), distStride);
    var lowmod = mod(inds.dim(i).low, distStride);
    var offset = loclowidx - lowmod;
    if offset < 0 then
      sliceBy(i) = inds.dim(i).low + (distStride + offset)..inds.dim(i).high by distStride;
    else
      sliceBy(i) = inds.dim(i).low + offset..inds.dim(i).high by distStride;
  }
  return inds((...sliceBy));
  //
  // WANT:
  //var distWhole = locDist(locid).myChunk.getIndices();
  //return inds((...distWhole));
}

def Cyclic.dsiDisplayRepresentation() {
  writeln("lowIdx = ", lowIdx);
  writeln("targetLocDom = ", targetLocDom);
  writeln("targetLocs = ", for tl in targetLocs do tl.id);
  writeln("tasksPerLocale = ", tasksPerLocale);
  for tli in targetLocDom do
    writeln("locDist[", tli, "].myChunk = ", locDist[tli].myChunk);
}

def Cyclic.dsiSupportsPrivatization() param return true;

def Cyclic.dsiGetPrivatizeData() return 0;

def Cyclic.dsiPrivatize(privatizeData) {
  return new Cyclic(rank=rank, idxType=idxType, other=this);
}

def Cyclic.dsiGetReprivatizeData() return 0;

def Cyclic.dsiReprivatize(other, reprivatizeData) {
  targetLocDom = other.targetLocDom;
  targetLocs = other.targetLocs;
  locDist = other.locDist;
  lowIdx = other.lowIdx;
  tasksPerLocale = other.tasksPerLocale;
}

def Cyclic.dsiNewArithmeticDom(param rank: int, type idxType, param stridable: bool) {
  if idxType != this.idxType then
    compilerError("Cyclic domain index type does not match distribution's");
  if rank != this.rank then
    compilerError("Cyclic domain rank does not match distribution's");
  var dom = new CyclicDom(rank=rank, idxType=idxType, dist = this, stridable=stridable);
  dom.setup();
  return dom;
}

def Cyclic.dsiCreateReindexDist(newSpace, oldSpace) {
  var newLow: rank*idxType;
  for param i in 1..rank {
    const oldDist = (oldSpace(i).low - lowIdx(i)) / oldSpace(i).stride;
    newLow(i) = newSpace(i).low - (newSpace(i).stride * oldDist);
  }
  var newDist = new Cyclic(rank=rank, idxType=idxType, low=newLow, targetLocales=targetLocs, tasksPerLocale=tasksPerLocale);
  return newDist;
}

//
// Given a tuple of scalars of type t or range(t) match the shape but
// using types rangeType and scalarType e.g. the call:
// _matchArgsShape(range(int), int, (1:int(64), 1:int(64)..5, 1:int(64)..5))
// returns the type: (int, range(int), range(int))
//
def _cyclic_matchArgsShape(type rangeType, type scalarType, args) type {
  def tuple(type t ...) type return t;
  def helper(param i: int) type {
    if i == args.size {
      if isCollapsedDimension(args(i)) then
        return tuple(scalarType);
      else
        return tuple(rangeType);
    } else {
      if isCollapsedDimension(args(i)) then
        return (scalarType, (... helper(i+1)));
      else
        return (rangeType, (... helper(i+1)));
    }
  }
  return helper(1);
}

def Cyclic.dsiCreateRankChangeDist(param newRank: int, args) {
  var collapsedDimInds: rank*idxType;
  var collapsedLocs: _cyclic_matchArgsShape(range(int), int, args);
  var newLow: newRank*idxType;

  var j: int = 1;
  for param i in 1..args.size {
    if isCollapsedDimension(args(i)) then
      collapsedDimInds(i) = args(i);
    else {
      newLow(j) = lowIdx(i);
      j += 1;
    }
  }
  const partialLocIdx = ind2locInd(collapsedDimInds);

  for param i in 1..args.size {
    if isCollapsedDimension(args(i)) {
      collapsedLocs(i) = partialLocIdx(i);
} else {
      collapsedLocs(i) = targetLocDom.dim(i);
    }
  }
  var newTargetLocales = targetLocs[(...collapsedLocs)];
  return new Cyclic(rank=newRank, idxType=idxType, low=newLow, targetLocales=newTargetLocales);
}

def Cyclic.writeThis(x: Writer) {
  x.writeln(typeToString(this.type));
  x.writeln("------");
  for locid in targetLocDom do
    x.writeln(" [", locid, "=", targetLocs(locid), "] owns chunk: ", locDist(locid).myChunk);
}

def Cyclic.ind2locInd(i: idxType) {
  const numLocs:idxType = targetLocDom.numIndices:idxType;
  // this is wrong if i is less than lowIdx
  //return ((i - lowIdx(1)) % numLocs):int;
  // this works even if i is less than lowIdx
  return mod(mod(i, numLocs) - mod(lowIdx(1), numLocs), numLocs):int;
}

def Cyclic.ind2locInd(ind: rank*idxType) {
  var x: rank*int;
  for param i in 1..rank {
    var dimLen = targetLocDom.dim(i).length;
    //x(i) = ((ind(i) - lowIdx(i)) % dimLen):int;
    x(i) = mod(mod(ind(i), dimLen) - mod(lowIdx(i), dimLen), dimLen):int;
  }
  if rank == 1 then
    return x(1);
  else
    return x;
}

def Cyclic.dsiIndexLocale(i: idxType) where rank == 1 {
  return targetLocs(ind2locInd(i));
}

def Cyclic.dsiIndexLocale(i: rank*idxType) {
  return targetLocs(ind2locInd(i));
}

class LocCyclic {
  param rank: int;
  type idxType;

  const myChunk: domain(rank, idxType, true);

  def LocCyclic(param rank, type idxType, locid, dist: Cyclic(rank, idxType)) {
    var locidx: rank*int;
    var lowIdx = dist.lowIdx;

    if rank == 1 then
      locidx(1) = locid;
    else
      locidx = locid;

    var tuple: rank*range(idxType, stridable=true);

    for param i in 1..rank {
      const lower = min(idxType)..(lowIdx(i)+locidx(i)) by -dist.targetLocDom.dim(i).length;
      const upper = (lowIdx(i) + locidx(i))..max(idxType) by dist.targetLocDom.dim(i).length;
      const lo = lower.low, hi = upper.high;
      tuple(i) = lo..hi by dist.targetLocDom.dim(i).length;
    }
    myChunk = [(...tuple)];
  }
}

class CyclicDom : BaseArithmeticDom {
  param rank: int;
  type idxType;
  param stridable: bool;

  const dist: Cyclic(rank, idxType);

  var locDoms: [dist.targetLocDom] LocCyclicDom(rank, idxType, stridable);

  const whole: domain(rank, idxType, stridable);

  var pid: int = -1;
}

def CyclicDom.setup() {
  if locDoms(dist.targetLocDom.low) == nil {
    coforall localeIdx in dist.targetLocDom {
      on dist.targetLocs(localeIdx) do
        locDoms(localeIdx) = new LocCyclicDom(rank, idxType, stridable, dist.getChunk(whole, localeIdx));
    }
  } else {
    coforall localeIdx in dist.targetLocDom {
      on dist.targetLocs(localeIdx) {
        var chunk = dist.getChunk(whole, localeIdx);
        locDoms(localeIdx).myBlock = chunk;
      }
    }
  }
}

def CyclicDom.dsiBuildArray(type eltType) {
  var arr = new CyclicArr(eltType=eltType, rank=rank, idxType=idxType, stridable=stridable, dom=this);
  arr.setup();
  return arr;
}

def CyclicDom.dsiDisplayRepresentation() {
  writeln("whole = ", whole);
  for tli in dist.targetLocDom do
    writeln("locDoms[", tli, "].myBlock = ", locDoms[tli].myBlock);
  dist.dsiDisplayRepresentation();
}

def CyclicDom.dsiLow return whole.low;

def CyclicDom.dsiHigh return whole.high;

def CyclicDom.dsiMember(i) return whole.member(i);

def CyclicDom.dsiIndexOrder(i) return whole.indexOrder(i);

def CyclicDom.dsiDims() return whole.dims();

def CyclicDom.dsiDim(d: int) return whole.dim(d);

def CyclicDom.getLocDom(localeIdx) return locDoms(localeIdx);

def CyclicDom.dsiGetIndices() {
  return whole.getIndices();
}

def CyclicDom.dsiSetIndices(x: domain) {
  whole = x;
  setup();
}

def CyclicDom.dsiSetIndices(x) {
  whole.setIndices(x);
  setup();
}

def CyclicDom.dsiSerialWrite(x: Writer) {
  if verboseCyclicDistWriters {
    x.writeln(typeToString(this.type));
    x.writeln("------");
    for loc in dist.targetLocDom {
      x.writeln("[", loc, "=", dist.targetLocs(loc), "] owns ", locDoms(loc).myBlock);
    }
  } else {
    x.write(whole);
  }
}

def CyclicDom.dsiNumIndices return whole.numIndices;

def CyclicDom.these() {
  for i in whole do
    yield i;
}

def CyclicDom.these(param tag: iterator) where tag == iterator.leader {
  const wholeLow = whole.low,
        precomputedNumTasks = dist.tasksPerLocale;
  coforall locDom in locDoms do on locDom {
    const numTasks = precomputedNumTasks;

    var result: rank*range(eltType=idxType, stridable=true);
    var zeroedLocalPart = whole((...locDom.myBlock.getIndices())) - whole.low;
    for param i in 1..rank {
      var dim = zeroedLocalPart.dim(i);
      var wholestride = whole.dim(i).stride;
      if dim.high >= dim.low then
        result(i) = (dim.low / wholestride)..(dim.high / wholestride) by (dim.stride / wholestride);
      else
        result(i) = 1..0 by 1;
    }
    if numTasks == 1 {
      if debugCyclicDist then
        writeln(here.id, ": leader whole: ", whole,
                         " result: ", result,
                         " myblock: ", locDom.myBlock);
      yield result;
    } else {

      coforall taskid in 0..#numTasks {
        var splitRanges: rank*range(eltType=idxType, stridable=true) = result;
        const low = result(1).low, high = result(1).high;
        const (lo,hi) = _computeBlock(high - low + 1, numTasks, taskid,
                                      high, low, low);
        splitRanges(1) = result(1)(lo..hi);
        if debugCyclicDist then
          writeln(here.id, ": leader whole: ", whole,
                           " result: ", result,
                           " splitRanges: ", splitRanges);
        yield splitRanges;
      }
    }
  }
}

def CyclicDom.these(param tag: iterator, follower, param aligned: bool = false) where tag == iterator.follower {
  var t: rank*range(idxType, stridable=true);
  if debugCyclicDist then
    writeln(here.id, ": follower whole is: ", whole,
                     " follower is: ", follower);
  for param i in 1..rank {
    const wholestride = whole.dim(i).stride;
    t(i) = ((follower(i).low*wholestride)..(follower(i).high*wholestride) by (follower(i).stride*wholestride)) + whole.dim(i).low;
  }
  if debugCyclicDist then
    writeln(here.id, ": follower maps to: ", t);
  for i in [(...t)] do
    yield i;
}

def CyclicDom.dsiSupportsPrivatization() param return true;

def CyclicDom.dsiGetPrivatizeData() return 0;

def CyclicDom.dsiPrivatize(privatizeData) {
  // These two variables are actually necessary even though it looks like
  // dist and dist.pid could be passed to the primitive directly.
  var distpid = dist.pid;
  var thisdist = dist;
  var privdist = __primitive("chpl_getPrivatizedClass", thisdist, distpid);
  var c = new CyclicDom(rank=rank, idxType=idxType, stridable=stridable, dist=privdist);
  c.locDoms = locDoms;
  c.whole = whole;
  return c;
}

def CyclicDom.dsiGetReprivatizeData() return 0;

def CyclicDom.dsiReprivatize(other, reprivatizeData) {
  locDoms = other.locDoms;
  whole = other.whole;
}

def CyclicDom.dsiBuildArithmeticDom(param rank, type idxType,
                                    param stridable: bool,
                                    ranges: rank*range(idxType,
                                                       BoundedRangeType.bounded,
                                                       stridable)) {
  if idxType != dist.idxType then
    compilerError("Cyclic domain index type does not match distribution's");
  if rank != dist.rank then
    compilerError("Cyclic domain rank does not match distribution's");

  var dom = new CyclicDom(rank=rank, idxType=idxType,
                         dist=dist, stridable=stridable);
  dom.dsiSetIndices(ranges);
  return dom;

}

def CyclicDom.localSlice(param stridable: bool, ranges) {
  return whole((...ranges));
}

class LocCyclicDom {
  param rank: int;
  type idxType;
  param stridable: bool;

  var myBlock: domain(rank, idxType, true);
}

//
// Added as a performance stopgap to avoid returning a domain
//
def LocCyclicDom.member(i) return myBlock.member(i);

class CyclicArr: BaseArr {
  type eltType;
  param rank: int;
  type idxType;
  param stridable: bool;
  var dom: CyclicDom(rank, idxType, stridable);

  var locArr: [dom.dist.targetLocDom] LocCyclicArr(eltType, rank, idxType, stridable);
  var myLocArr: LocCyclicArr(eltType=eltType, rank=rank, idxType=idxType, stridable=stridable);
  var pid: int = -1;
}

def CyclicArr.dsiSlice(d: CyclicDom) {
  var alias = new CyclicArr(eltType=eltType, rank=rank, idxType=idxType,
                            stridable=d.stridable, dom=d, pid=pid);
  var thisid = this.locale.uid;
  for i in dom.dist.targetLocDom {
    on dom.dist.targetLocs(i) {
      alias.locArr[i] =
        new LocCyclicArr(eltType=eltType, rank=rank, idxType=idxType,
                         stridable=d.stridable, locDom=d.locDoms[i],
                         myElems=>locArr[i].myElems[d.locDoms[i].myBlock]);
      if thisid == here.uid then
        alias.myLocArr = alias.locArr[i];
    }
  }
  return alias;
}

def CyclicArr.dsiRankChange(d, param newRank: int, param stridable: bool, args) {
  var alias = new CyclicArr(eltType=eltType, rank=newRank, idxType=idxType, stridable=stridable, dom = d);
  var thisid = this.locale.uid;
  coforall ind in d.dist.targetLocDom {
    on d.dist.targetLocs(ind) {
      const locDom = d.getLocDom(ind);
      var collapsedDims: rank*idxType;
      var locSlice: _cyclic_matchArgsShape(range(eltType=idxType, stridable=true), idxType, args);
      var locArrInd: rank*int;
      var j = 1;
      for param i in 1..args.size {
        if isCollapsedDimension(args(i)) {
          locSlice(i) = args(i);
          collapsedDims(i) = args(i);
        } else {
          locSlice(i) = locDom.myBlock.dim(j)(args(i));
          j += 1;
        }
      }
      locArrInd = dom.dist.ind2locInd(collapsedDims);
      j = 1;
      // Now that the locArrInd values are known for the collapsed dimensions
      // Pull the rest of the dimensions values from ind
      for param i in 1..args.size {
        if !isCollapsedDimension(args(i)) {
          if newRank > 1 then
            locArrInd(i) = ind(j);
          else
            locArrInd(i) = ind;
          j += 1;
        }
      }
      alias.locArr[ind] =
        new LocCyclicArr(eltType=eltType, rank=newRank, idxType=d.idxType,
                         stridable=d.stridable, locDom=locDom,
                         myElems=>locArr[(...locArrInd)].myElems[(...locSlice)]);
      if here.uid == thisid then
        alias.myLocArr = alias.locArr[ind];
    }
  }
  return alias;
}

def CyclicArr.dsiReindex(d: CyclicDom) {
  var alias = new CyclicArr(eltType=eltType, rank=rank, idxType=d.idxType,
                            stridable=d.stridable, dom=d);
  var thisid = this.locale.uid;
  coforall i in dom.dist.targetLocDom {
    on dom.dist.targetLocs(i) {
      const locDom = d.getLocDom(i);
      const locAlias => locArr[i].myElems;
      alias.locArr[i] = new LocCyclicArr(eltType=eltType, idxType=idxType,
                                         locDom=locDom, stridable=d.stridable,
                                         rank=rank, myElems=>locAlias);
      if thisid == here.uid then
        alias.myLocArr = alias.locArr[i];
    }
  }
  return alias;
}

def CyclicArr.localSlice(ranges) {
  var low: rank*idxType;
  for param i in 1..rank {
    low(i) = ranges(i).low;
  }

  var A => locArr(dom.dist.ind2locInd(low)).myElems((...ranges));
  return A;
}

def CyclicArr.dsiDisplayRepresentation() {
  for tli in dom.dist.targetLocDom do
    writeln("locArr[", tli, "].myElems = ", for e in locArr[tli].myElems do e);
  dom.dsiDisplayRepresentation();
}

def CyclicArr.dsiGetBaseDom() return dom;

def CyclicArr.setup() {
  coforall localeIdx in dom.dist.targetLocDom {
    on dom.dist.targetLocs(localeIdx) {
      locArr(localeIdx) = new LocCyclicArr(eltType, rank, idxType, stridable, dom.locDoms(localeIdx));
      if this.locale == here then
        myLocArr = locArr(localeIdx);
    }
  }
}

def CyclicArr.dsiSupportsPrivatization() param return true;

def CyclicArr.dsiGetPrivatizeData() return 0;

def CyclicArr.dsiPrivatize(privatizeData) {
  // These two variables are actually necessary even though it looks like
  // dist and dist.pid could be passed to the primitive directly.
  var dompid = dom.pid;
  var thisdom = dom;
  var privdom = __primitive("chpl_getPrivatizedClass", thisdom, dompid);
  var c = new CyclicArr(eltType=eltType, rank=rank, idxType=idxType, stridable=stridable, dom=privdom);
  c.locArr = locArr;
  for localeIdx in dom.dist.targetLocDom do
    if c.locArr(localeIdx).locale == here then
      c.myLocArr = c.locArr(localeIdx);
  return c;
}

def CyclicArr.dsiAccess(i: idxType) var where rank == 1 {
  local {
    if myLocArr != nil && myLocArr.locDom.myBlock.member(i) then
      return myLocArr.this(i);
  }
  return locArr(dom.dist.ind2locInd(i))(i);
}

def CyclicArr.dsiAccess(i:rank*idxType) var {
  return locArr(dom.dist.ind2locInd(i))(i);
}

def CyclicArr.these() var {
  for i in dom do
    yield dsiAccess(i);
}

def CyclicArr.these(param tag: iterator) where tag == iterator.leader {
  const wholeLow = dom.whole.low,
        precomputedNumTasks = dom.dist.tasksPerLocale;
  coforall locDom in dom.locDoms do on locDom {
    const numTasks = precomputedNumTasks;

    var result: rank*range(eltType=idxType, stridable=true);
    var zeroedLocalPart = dom.whole((...locDom.myBlock.getIndices())) - wholeLow;
    for param i in 1..rank {
      var dim = zeroedLocalPart.dim(i);
      var wholestride = dom.whole.dim(i).stride;
      if dim.high >= dim.low then
        result(i) = (dim.low / wholestride)..(dim.high / wholestride) by (dim.stride / wholestride);
      else
        result(i) = 1..0 by 1;
    }
    if numTasks == 1 {
      if debugCyclicDist then
        writeln(here.id, ": leader whole: ", dom.whole,
                         " result: ", result,
                         " myblock: ", locDom.myBlock);
      yield result;
    } else {

      coforall taskid in 0..#numTasks {
        var splitRanges: rank*range(eltType=idxType, stridable=true) = result;
        const low = result(1).low, high = result(1).high;
        const (lo,hi) = _computeBlock(high - low + 1, numTasks, taskid,
                                      high, low, low);
        splitRanges(1) = result(1)(lo..hi);
        if debugCyclicDist then
          writeln(here.id, ": leader whole: ", dom.whole,
                           " result: ", result,
                           " splitRanges: ", splitRanges);
        yield splitRanges;
      }
    }
  }
}

def CyclicArr.dsiSupportsAlignedFollower() param return true;

def CyclicArr.these(param tag: iterator, follower, param aligned: bool = false) var where tag == iterator.follower {
  var t: rank*range(eltType=idxType, stridable=true);
  for param i in 1..rank {
    const wholestride = dom.whole.dim(i).stride;
    t(i) = ((follower(i).low*wholestride)..(follower(i).high*wholestride) by (follower(i).stride*wholestride)) + dom.whole.dim(i).low;
  }
  const followThis = [(...t)];
  const arrSection = locArr(dom.dist.ind2locInd(followThis.low));
  if aligned {
    local {
      for i in followThis {
        yield arrSection.this(i);
      }
    }
  } else {
    def accessHelper(i) var {
      local {
        if myLocArr.locDom.member(i) then
          return myLocArr.this(i);
      }
      return dsiAccess(i);
    }

    for i in followThis {
      yield accessHelper(i);
    }
  }
}

def CyclicArr.dsiSerialWrite(f: Writer) {
  if verboseCyclicDistWriters {
    writeln(typeToString(this.type));
    writeln("------");
  }
  if dom.dsiNumIndices == 0 then return;
  var i : rank*idxType;
  for dim in 1..rank do
    i(dim) = dom.dsiDim(dim)._low;
  label next while true {
    f.write(dsiAccess(i));
    if i(rank) <= (dom.dsiDim(rank)._high - dom.dsiDim(rank)._stride:idxType) {
      f.write(" ");
      i(rank) += dom.dsiDim(rank)._stride:idxType;
    } else {
      for dim in 1..rank-1 by -1 {
        if i(dim) <= (dom.dsiDim(dim)._high - dom.dsiDim(dim)._stride:idxType) {
          i(dim) += dom.dsiDim(dim)._stride:idxType;
          for dim2 in dim+1..rank {
            f.writeln();
            i(dim2) = dom.dsiDim(dim2)._low;
          }
          continue next;
        }
      }
      break;
    }
  }
}

def CyclicArr.dsiReallocate(d: domain) {
  // The reallocation happens when the LocCyclicDom.myBlock field is changed
  // in CyclicDom.setup(). Nothing more needs to happen here.
}

class LocCyclicArr {
  type eltType;
  param rank: int;
  type idxType;
  param stridable: bool;

  const locDom: LocCyclicDom(rank, idxType, stridable);

  var myElems: [locDom.myBlock] eltType;

}

def LocCyclicArr.this(i:idxType) var {
  return myElems(i);
}

def LocCyclicArr.this(i:rank*idxType) var {
  return myElems((...i));
}

def LocCyclicArr.these() var {
  for elem in myElems {
    yield elem;
  }
}

use DSIUtil;

def _tupleOfZero(param size: int, type t) {
  var tmp: size * t;
  return tmp;
}

config param debugCyclicDist = false;
config param verboseCyclicDistWriters = false;

class Cyclic: BaseDist {
  param rank: int = 1;
  type idxType = int(64);

  const lowIdx: rank*idxType;
  const targetLocDom: domain(rank);
  const targetLocs: [targetLocDom] locale;

  const locDist: [targetLocDom] LocCyclic(rank, idxType);

  const tasksPerLocale = 1;
  var pid: int = -1;

  def Cyclic(param rank: int,
             type idxType = int(64),
             low: rank*idxType = _tupleOfZero(rank, idxType),
             targetLocales: [] locale = thisRealm.Locales,
             tasksPerLocale=0) {
    if rank == 1  {
      targetLocDom = [0..#targetLocales.numElements];
      targetLocs = targetLocales;
    } else if targetLocales.rank == 1 {
      const factors = _factor(rank, targetLocales.numElements);
      var ranges: rank*range;
      for param i in 1..rank {
        ranges(i) = 0..#factors(i);
      }
      targetLocDom = [(...ranges)];
      for (loc1, loc2) in (targetLocs, targetLocales) {
        loc1 = loc2;
      }
    } else {
      if targetLocales.rank != rank then
        compilerError("locales array rank must be one or match distribution rank");
      var ranges: rank*range;
      for param i in 1..rank do {
        var thisRange = targetLocales.domain.dim(i);
        ranges(i) = 0..#thisRange.length;
      }
      targetLocDom = [(...ranges)];
      targetLocs = reshape(targetLocales, targetLocDom);
    }

    for param i in 1..rank do
      lowIdx(i) = mod(low(i), targetLocDom.dim(i).length);

    if tasksPerLocale == 0 then
      this.tasksPerLocale = min reduce targetLocales.numCores;
    else
      this.tasksPerLocale = tasksPerLocale;
    coforall locid in targetLocDom {
      on targetLocs(locid) {
        locDist(locid) = new LocCyclic(rank, idxType, locid, this);
      }
    }
    if debugCyclicDist then
      for loc in locDist do writeln(loc);
  }

  def Cyclic(param rank, type idxType, other: Cyclic(rank, idxType)) {
    targetLocDom = other.targetLocDom;
    targetLocs = other.targetLocs;
    lowIdx = other.lowIdx;
    locDist = other.locDist;
    tasksPerLocale = other.tasksPerLocale;
  }

  def dsiClone() return new Cyclic(rank=rank, idxType=idxType, other=this);
}

def Cyclic.getChunk(inds, locid) {
  var sliceBy: rank*range(eltType=idxType, stridable=true);
  var locidtup: rank*int;
  if rank == 1 then
    locidtup(1) = locid;
  else
    locidtup = locid;
  for param i in 1..rank {
    var distStride = targetLocDom.dim(i).length;
    var loclowidx = mod(lowIdx(i) + locidtup(i), distStride);
    var lowmod = mod(inds.dim(i).low, distStride);
    var offset = loclowidx - lowmod;
    if offset < 0 then
      sliceBy(i) = inds.dim(i).low + (distStride + offset)..inds.dim(i).high by distStride;
    else
      sliceBy(i) = inds.dim(i).low + offset..inds.dim(i).high by distStride;
  }
  return inds((...sliceBy));
  //
  // WANT:
  //var distWhole = locDist(locid).myChunk.getIndices();
  //return inds((...distWhole));
}

def Cyclic.dsiDisplayRepresentation() {
  writeln("lowIdx = ", lowIdx);
  writeln("targetLocDom = ", targetLocDom);
  writeln("targetLocs = ", for tl in targetLocs do tl.id);
  writeln("tasksPerLocale = ", tasksPerLocale);
  for tli in targetLocDom do
    writeln("locDist[", tli, "].myChunk = ", locDist[tli].myChunk);
}

def Cyclic.dsiSupportsPrivatization() param return true;

def Cyclic.dsiGetPrivatizeData() return 0;

def Cyclic.dsiPrivatize(privatizeData) {
  return new Cyclic(rank=rank, idxType=idxType, other=this);
}

def Cyclic.dsiGetReprivatizeData() return 0;

def Cyclic.dsiReprivatize(other, reprivatizeData) {
  targetLocDom = other.targetLocDom;
  targetLocs = other.targetLocs;
  locDist = other.locDist;
  lowIdx = other.lowIdx;
  tasksPerLocale = other.tasksPerLocale;
}

def Cyclic.dsiNewArithmeticDom(param rank: int, type idxType, param stridable: bool) {
  if idxType != this.idxType then
    compilerError("Cyclic domain index type does not match distribution's");
  if rank != this.rank then
    compilerError("Cyclic domain rank does not match distribution's");
  var dom = new CyclicDom(rank=rank, idxType=idxType, dist = this, stridable=stridable);
  dom.setup();
  return dom;
}

def Cyclic.dsiCreateReindexDist(newSpace, oldSpace) {
  var newLow: rank*idxType;
  for param i in 1..rank {
    const oldDist = (oldSpace(i).low - lowIdx(i)) / oldSpace(i).stride;
    newLow(i) = newSpace(i).low - (newSpace(i).stride * oldDist);
  }
  var newDist = new Cyclic(rank=rank, idxType=idxType, low=newLow, targetLocales=targetLocs, tasksPerLocale=tasksPerLocale);
  return newDist;
}

//
// Given a tuple of scalars of type t or range(t) match the shape but
// using types rangeType and scalarType e.g. the call:
// _matchArgsShape(range(int), int, (1:int(64), 1:int(64)..5, 1:int(64)..5))
// returns the type: (int, range(int), range(int))
//
def _cyclic_matchArgsShape(type rangeType, type scalarType, args) type {
  def tuple(type t ...) type return t;
  def helper(param i: int) type {
    if i == args.size {
      if isCollapsedDimension(args(i)) then
        return tuple(scalarType);
      else
        return tuple(rangeType);
    } else {
      if isCollapsedDimension(args(i)) then
        return (scalarType, (... helper(i+1)));
      else
        return (rangeType, (... helper(i+1)));
    }
  }
  return helper(1);
}

def Cyclic.dsiCreateRankChangeDist(param newRank: int, args) {
  var collapsedDimInds: rank*idxType;
  var collapsedLocs: _cyclic_matchArgsShape(range(int), int, args);
  var newLow: newRank*idxType;

  var j: int = 1;
  for param i in 1..args.size {
    if isCollapsedDimension(args(i)) then
      collapsedDimInds(i) = args(i);
    else {
      newLow(j) = lowIdx(i);
      j += 1;
    }
  }
  const partialLocIdx = ind2locInd(collapsedDimInds);

  for param i in 1..args.size {
    if isCollapsedDimension(args(i)) {
      collapsedLocs(i) = partialLocIdx(i);
} else {
      collapsedLocs(i) = targetLocDom.dim(i);
    }
  }
  var newTargetLocales = targetLocs[(...collapsedLocs)];
  return new Cyclic(rank=newRank, idxType=idxType, low=newLow, targetLocales=newTargetLocales);
}

def Cyclic.writeThis(x: Writer) {
  x.writeln(typeToString(this.type));
  x.writeln("------");
  for locid in targetLocDom do
    x.writeln(" [", locid, "=", targetLocs(locid), "] owns chunk: ", locDist(locid).myChunk);
}

def Cyclic.ind2locInd(i: idxType) {
  const numLocs:idxType = targetLocDom.numIndices:idxType;
  // this is wrong if i is less than lowIdx
  //return ((i - lowIdx(1)) % numLocs):int;
  // this works even if i is less than lowIdx
  return mod(mod(i, numLocs) - mod(lowIdx(1), numLocs), numLocs):int;
}

def Cyclic.ind2locInd(ind: rank*idxType) {
  var x: rank*int;
  for param i in 1..rank {
    var dimLen = targetLocDom.dim(i).length;
    //x(i) = ((ind(i) - lowIdx(i)) % dimLen):int;
    x(i) = mod(mod(ind(i), dimLen) - mod(lowIdx(i), dimLen), dimLen):int;
  }
  if rank == 1 then
    return x(1);
  else
    return x;
}

def Cyclic.dsiIndexLocale(i: idxType) where rank == 1 {
  return targetLocs(ind2locInd(i));
}

def Cyclic.dsiIndexLocale(i: rank*idxType) {
  return targetLocs(ind2locInd(i));
}

class LocCyclic {
  param rank: int;
  type idxType;

  const myChunk: domain(rank, idxType, true);

  def LocCyclic(param rank, type idxType, locid, dist: Cyclic(rank, idxType)) {
    var locidx: rank*int;
    var lowIdx = dist.lowIdx;

    if rank == 1 then
      locidx(1) = locid;
    else
      locidx = locid;

    var tuple: rank*range(idxType, stridable=true);

    for param i in 1..rank {
      const lower = min(idxType)..(lowIdx(i)+locidx(i)) by -dist.targetLocDom.dim(i).length;
      const upper = (lowIdx(i) + locidx(i))..max(idxType) by dist.targetLocDom.dim(i).length;
      const lo = lower.low, hi = upper.high;
      tuple(i) = lo..hi by dist.targetLocDom.dim(i).length;
    }
    myChunk = [(...tuple)];
  }
}

class CyclicDom : BaseArithmeticDom {
  param rank: int;
  type idxType;
  param stridable: bool;

  const dist: Cyclic(rank, idxType);

  var locDoms: [dist.targetLocDom] LocCyclicDom(rank, idxType, stridable);

  const whole: domain(rank, idxType, stridable);

  var pid: int = -1;
}

def CyclicDom.setup() {
  coforall localeIdx in dist.targetLocDom {
    on dist.targetLocs(localeIdx) {
      var chunk = dist.getChunk(whole, localeIdx);
      locDoms(localeIdx) = new LocCyclicDom(rank, idxType, stridable, chunk);
    }
  }
}

def CyclicDom.dsiBuildArray(type eltType) {
  var arr = new CyclicArr(eltType=eltType, rank=rank, idxType=idxType, stridable=stridable, dom=this);
  arr.setup();
  return arr;
}

def CyclicDom.dsiDisplayRepresentation() {
  writeln("whole = ", whole);
  for tli in dist.targetLocDom do
    writeln("locDoms[", tli, "].myBlock = ", locDoms[tli].myBlock);
  dist.dsiDisplayRepresentation();
}

def CyclicDom.dsiLow return whole.low;

def CyclicDom.dsiHigh return whole.high;

def CyclicDom.dsiMember(i) return whole.member(i);

def CyclicDom.dsiIndexOrder(i) return whole.indexOrder(i);

def CyclicDom.dsiDims() return whole.dims();

def CyclicDom.dsiDim(d: int) return whole.dim(d);

def CyclicDom.getLocDom(localeIdx) return locDoms(localeIdx);

def CyclicDom.dsiGetIndices() {
  return whole.getIndices();
}

def CyclicDom.dsiSetIndices(x: domain) {
  whole = x;
  setup();
}

def CyclicDom.dsiSetIndices(x) {
  whole.setIndices(x);
  setup();
}

def CyclicDom.dsiSerialWrite(x: Writer) {
  if verboseCyclicDistWriters {
    x.writeln(typeToString(this.type));
    x.writeln("------");
    for loc in dist.targetLocDom {
      x.writeln("[", loc, "=", dist.targetLocs(loc), "] owns ", locDoms(loc).myBlock);
    }
  } else {
    x.write(whole);
  }
}

def CyclicDom.dsiNumIndices return whole.numIndices;

def CyclicDom.these() {
  for i in whole do
    yield i;
}

def CyclicDom.these(param tag: iterator) where tag == iterator.leader {
  const wholeLow = whole.low,
        precomputedNumTasks = dist.tasksPerLocale;
  coforall locDom in locDoms do on locDom {
    const numTasks = precomputedNumTasks;

    var result: rank*range(eltType=idxType, stridable=true);
    var zeroedLocalPart = whole((...locDom.myBlock.getIndices())) - whole.low;
    for param i in 1..rank {
      var dim = zeroedLocalPart.dim(i);
      var wholestride = whole.dim(i).stride;
      if dim.high >= dim.low then
        result(i) = (dim.low / wholestride)..(dim.high / wholestride) by (dim.stride / wholestride);
      else
        result(i) = 1..0 by 1;
    }
    if numTasks == 1 {
      if debugCyclicDist then
        writeln(here.id, ": leader whole: ", whole,
                         " result: ", result,
                         " myblock: ", locDom.myBlock);
      yield result;
    } else {

      coforall taskid in 0..#numTasks {
        var splitRanges: rank*range(eltType=idxType, stridable=true) = result;
        const low = result(1).low, high = result(1).high;
        const (lo,hi) = _computeBlock(high - low + 1, numTasks, taskid,
                                      high, low, low);
        splitRanges(1) = result(1)(lo..hi);
        if debugCyclicDist then
          writeln(here.id, ": leader whole: ", whole,
                           " result: ", result,
                           " splitRanges: ", splitRanges);
        yield splitRanges;
      }
    }
  }
}

def CyclicDom.these(param tag: iterator, follower, param aligned: bool = false) where tag == iterator.follower {
  var t: rank*range(idxType, stridable=true);
  if debugCyclicDist then
    writeln(here.id, ": follower whole is: ", whole,
                     " follower is: ", follower);
  for param i in 1..rank {
    const wholestride = whole.dim(i).stride;
    t(i) = ((follower(i).low*wholestride)..(follower(i).high*wholestride) by (follower(i).stride*wholestride)) + whole.dim(i).low;
  }
  if debugCyclicDist then
    writeln(here.id, ": follower maps to: ", t);
  for i in [(...t)] do
    yield i;
}

def CyclicDom.dsiSupportsPrivatization() param return true;

def CyclicDom.dsiGetPrivatizeData() return 0;

def CyclicDom.dsiPrivatize(privatizeData) {
  // These two variables are actually necessary even though it looks like
  // dist and dist.pid could be passed to the primitive directly.
  var distpid = dist.pid;
  var thisdist = dist;
  var privdist = __primitive("chpl_getPrivatizedClass", thisdist, distpid);
  var c = new CyclicDom(rank=rank, idxType=idxType, stridable=stridable, dist=privdist);
  c.locDoms = locDoms;
  c.whole = whole;
  return c;
}

def CyclicDom.dsiGetReprivatizeData() return 0;

def CyclicDom.dsiReprivatize(other, reprivatizeData) {
  locDoms = other.locDoms;
  whole = other.whole;
}

def CyclicDom.dsiBuildArithmeticDom(param rank, type idxType,
                                    param stridable: bool,
                                    ranges: rank*range(idxType,
                                                       BoundedRangeType.bounded,
                                                       stridable)) {
  if idxType != dist.idxType then
    compilerError("Cyclic domain index type does not match distribution's");
  if rank != dist.rank then
    compilerError("Cyclic domain rank does not match distribution's");

  var dom = new CyclicDom(rank=rank, idxType=idxType,
                         dist=dist, stridable=stridable);
  dom.dsiSetIndices(ranges);
  return dom;

}

def CyclicDom.localSlice(param stridable: bool, ranges) {
  return whole((...ranges));
}

class LocCyclicDom {
  param rank: int;
  type idxType;
  param stridable: bool;

  var myBlock: domain(rank, idxType, true);
}

//
// Added as a performance stopgap to avoid returning a domain
//
def LocCyclicDom.member(i) return myBlock.member(i);

class CyclicArr: BaseArr {
  type eltType;
  param rank: int;
  type idxType;
  param stridable: bool;
  var dom: CyclicDom(rank, idxType, stridable);

  var locArr: [dom.dist.targetLocDom] LocCyclicArr(eltType, rank, idxType, stridable);
  var myLocArr: LocCyclicArr(eltType=eltType, rank=rank, idxType=idxType, stridable=stridable);
  var pid: int = -1;
}

def CyclicArr.dsiSlice(d: CyclicDom) {
  var alias = new CyclicArr(eltType=eltType, rank=rank, idxType=idxType,
                            stridable=d.stridable, dom=d, pid=pid);
  var thisid = this.locale.uid;
  for i in dom.dist.targetLocDom {
    on dom.dist.targetLocs(i) {
      alias.locArr[i] =
        new LocCyclicArr(eltType=eltType, rank=rank, idxType=idxType,
                         stridable=d.stridable, locDom=d.locDoms[i],
                         myElems=>locArr[i].myElems[d.locDoms[i].myBlock]);
      if thisid == here.uid then
        alias.myLocArr = alias.locArr[i];
    }
  }
  return alias;
}

def CyclicArr.dsiRankChange(d, param newRank: int, param stridable: bool, args) {
  var alias = new CyclicArr(eltType=eltType, rank=newRank, idxType=idxType, stridable=stridable, dom = d);
  var thisid = this.locale.uid;
  coforall ind in d.dist.targetLocDom {
    on d.dist.targetLocs(ind) {
      const locDom = d.getLocDom(ind);
      var collapsedDims: rank*idxType;
      var locSlice: _cyclic_matchArgsShape(range(eltType=idxType, stridable=true), idxType, args);
      var locArrInd: rank*int;
      var j = 1;
      for param i in 1..args.size {
        if isCollapsedDimension(args(i)) {
          locSlice(i) = args(i);
          collapsedDims(i) = args(i);
        } else {
          locSlice(i) = locDom.myBlock.dim(j)(args(i));
          j += 1;
        }
      }
      locArrInd = dom.dist.ind2locInd(collapsedDims);
      j = 1;
      // Now that the locArrInd values are known for the collapsed dimensions
      // Pull the rest of the dimensions values from ind
      for param i in 1..args.size {
        if !isCollapsedDimension(args(i)) {
          if newRank > 1 then
            locArrInd(i) = ind(j);
          else
            locArrInd(i) = ind;
          j += 1;
        }
      }
      alias.locArr[ind] =
        new LocCyclicArr(eltType=eltType, rank=newRank, idxType=d.idxType,
                         stridable=d.stridable, locDom=locDom,
                         myElems=>locArr[(...locArrInd)].myElems[(...locSlice)]);
      if here.uid == thisid then
        alias.myLocArr = alias.locArr[ind];
    }
  }
  return alias;
}

def CyclicArr.dsiReindex(d: CyclicDom) {
  var alias = new CyclicArr(eltType=eltType, rank=rank, idxType=d.idxType,
                            stridable=d.stridable, dom=d);
  var thisid = this.locale.uid;
  coforall i in dom.dist.targetLocDom {
    on dom.dist.targetLocs(i) {
      const locDom = d.getLocDom(i);
      const locAlias => locArr[i].myElems;
      alias.locArr[i] = new LocCyclicArr(eltType=eltType, idxType=idxType,
                                         locDom=locDom, stridable=d.stridable,
                                         rank=rank, myElems=>locAlias);
      if thisid == here.uid then
        alias.myLocArr = alias.locArr[i];
    }
  }
  return alias;
}

def CyclicArr.localSlice(ranges) {
  var low: rank*idxType;
  for param i in 1..rank {
    low(i) = ranges(i).low;
  }

  var A => locArr(dom.dist.ind2locInd(low)).myElems((...ranges));
  return A;
}

def CyclicArr.dsiDisplayRepresentation() {
  for tli in dom.dist.targetLocDom do
    writeln("locArr[", tli, "].myElems = ", for e in locArr[tli].myElems do e);
  dom.dsiDisplayRepresentation();
}

def CyclicArr.dsiGetBaseDom() return dom;

def CyclicArr.setup() {
  coforall localeIdx in dom.dist.targetLocDom {
    on dom.dist.targetLocs(localeIdx) {
      locArr(localeIdx) = new LocCyclicArr(eltType, rank, idxType, stridable, dom.locDoms(localeIdx));
      if this.locale == here then
        myLocArr = locArr(localeIdx);
    }
  }
}

def CyclicArr.dsiSupportsPrivatization() param return true;

def CyclicArr.dsiGetPrivatizeData() return 0;

def CyclicArr.dsiPrivatize(privatizeData) {
  // These two variables are actually necessary even though it looks like
  // dist and dist.pid could be passed to the primitive directly.
  var dompid = dom.pid;
  var thisdom = dom;
  var privdom = __primitive("chpl_getPrivatizedClass", thisdom, dompid);
  var c = new CyclicArr(eltType=eltType, rank=rank, idxType=idxType, stridable=stridable, dom=privdom);
  c.locArr = locArr;
  for localeIdx in dom.dist.targetLocDom do
    if c.locArr(localeIdx).locale == here then
      c.myLocArr = c.locArr(localeIdx);
  return c;
}

def CyclicArr.dsiAccess(i: idxType) var where rank == 1 {
  local {
    if myLocArr != nil && myLocArr.locDom.myBlock.member(i) then
      return myLocArr.this(i);
  }
  return locArr(dom.dist.ind2locInd(i))(i);
}

def CyclicArr.dsiAccess(i:rank*idxType) var {
  return locArr(dom.dist.ind2locInd(i))(i);
}

def CyclicArr.these() var {
  for i in dom do
    yield dsiAccess(i);
}

def CyclicArr.these(param tag: iterator) where tag == iterator.leader {
  const wholeLow = dom.whole.low,
        precomputedNumTasks = dom.dist.tasksPerLocale;
  coforall locDom in dom.locDoms do on locDom {
    const numTasks = precomputedNumTasks;

    var result: rank*range(eltType=idxType, stridable=true);
    var zeroedLocalPart = dom.whole((...locDom.myBlock.getIndices())) - wholeLow;
    for param i in 1..rank {
      var dim = zeroedLocalPart.dim(i);
      var wholestride = dom.whole.dim(i).stride;
      if dim.high >= dim.low then
        result(i) = (dim.low / wholestride)..(dim.high / wholestride) by (dim.stride / wholestride);
      else
        result(i) = 1..0 by 1;
    }
    if numTasks == 1 {
      if debugCyclicDist then
        writeln(here.id, ": leader whole: ", dom.whole,
                         " result: ", result,
                         " myblock: ", locDom.myBlock);
      yield result;
    } else {

      coforall taskid in 0..#numTasks {
        var splitRanges: rank*range(eltType=idxType, stridable=true) = result;
        const low = result(1).low, high = result(1).high;
        const (lo,hi) = _computeBlock(high - low + 1, numTasks, taskid,
                                      high, low, low);
        splitRanges(1) = result(1)(lo..hi);
        if debugCyclicDist then
          writeln(here.id, ": leader whole: ", dom.whole,
                           " result: ", result,
                           " splitRanges: ", splitRanges);
        yield splitRanges;
      }
    }
  }
}

def CyclicArr.dsiSupportsAlignedFollower() param return true;

def CyclicArr.these(param tag: iterator, follower, param aligned: bool = false) var where tag == iterator.follower {
  var t: rank*range(eltType=idxType, stridable=true);
  for param i in 1..rank {
    const wholestride = dom.whole.dim(i).stride;
    t(i) = ((follower(i).low*wholestride)..(follower(i).high*wholestride) by (follower(i).stride*wholestride)) + dom.whole.dim(i).low;
  }
  const followThis = [(...t)];
  const arrSection = locArr(dom.dist.ind2locInd(followThis.low));
  if aligned {
    local {
      for i in followThis {
        yield arrSection.this(i);
      }
    }
  } else {
    def accessHelper(i) var {
      local {
        if myLocArr.locDom.member(i) then
          return myLocArr.this(i);
      }
      return dsiAccess(i);
    }

    for i in followThis {
      yield accessHelper(i);
    }
  }
}

def CyclicArr.dsiSerialWrite(f: Writer) {
  if verboseCyclicDistWriters {
    writeln(typeToString(this.type));
    writeln("------");
  }
  if dom.dsiNumIndices == 0 then return;
  var i : rank*idxType;
  for dim in 1..rank do
    i(dim) = dom.dsiDim(dim)._low;
  label next while true {
    f.write(dsiAccess(i));
    if i(rank) <= (dom.dsiDim(rank)._high - dom.dsiDim(rank)._stride:idxType) {
      f.write(" ");
      i(rank) += dom.dsiDim(rank)._stride:idxType;
    } else {
      for dim in 1..rank-1 by -1 {
        if i(dim) <= (dom.dsiDim(dim)._high - dom.dsiDim(dim)._stride:idxType) {
          i(dim) += dom.dsiDim(dim)._stride:idxType;
          for dim2 in dim+1..rank {
            f.writeln();
            i(dim2) = dom.dsiDim(dim2)._low;
          }
          continue next;
        }
      }
      break;
    }
  }
}

class LocCyclicArr {
  type eltType;
  param rank: int;
  type idxType;
  param stridable: bool;

  const locDom: LocCyclicDom(rank, idxType, stridable);

  var myElems: [locDom.myBlock] eltType;

}

def LocCyclicArr.this(i:idxType) var {
  return myElems(i);
}

def LocCyclicArr.this(i:rank*idxType) var {
  return myElems((...i));
}

def LocCyclicArr.these() var {
  for elem in myElems {
    yield elem;
  }
}