Numbas.addExtension('linalg',['jme'],function(extension) {
  var scope = extension.scope;
  var jme = Numbas.jme;
  var funcObj = jme.funcObj;
  var types = jme.types;
  var TNum = types.TNum;
  var TMatrix = types.TMatrix;
  var TList = types.TList;

  function pivot(M) {
  for (let j = 0; j < M[0].length; j++) {
    for (let i = 0; i < M.length; i++) {
      if (M[i][j] !== 0) {
        return [i,j];
      }
    }
  }
  return "0-matrix"
}

function copyMat(M) {
  let L = [];
  for (let i = 0; i < M.length; i++) {
    L[i] = M[i].slice();
  }
  return L;
}

// takes a matrix M and a pair of indices l. Returns the matrix with the rows specified in l switched.
function rowSwitch(M,l) {
  let L = copyMat(M);
  L[l[0]] = M[l[1]];
  L[l[1]] = M[l[0]];
  return L;
}

// take a submatrix out of a larger matrix
// First argument is the larger matrix, second and third argument are tuples which specify opposite corners of the submatrix of interest
function submat(M,l,k) {
  let L = [];
  for (let i2 = 0, i = Math.min(l[0],k[0]); i <= Math.max(l[0],k[0]); i++, i2++) {
    L[i2] = []
    for (let j2 = 0, j = Math.min(l[1],k[1]); j <= Math.max(l[1],k[1]); j++, j2++) {
      // L.push([i2,j2])
      L[i2][j2] = M[i][j];
    }
  }
  return L;
}

// replace a submatrix in a matrix
// first argument is the matrix, second argument is the replacement matrix, third argument is
// a tuple representing the left hand corner of the submatrix which is to replace.
function replmat(M,L,l) {
  N = copyMat(M);
  for (let i = 0, i2 = l[0]; i < L.length; i++, i2++) {
    for (let j = 0, j2 = l[1]; j < L[0].length; j++, j2++) {
      N[i2][j2] = L[i][j];
    }
  }
  return N
}

// takes a matrix returns true if matrix is zero-matrix
function iszero(M) {
  for (let i = 0; i < M.length; i++) {
    for (let j = 0; j < M[0].length; j++) {
      if (M[i][j] !== 0) {
        return false;
      }
    }
  }
  return true;
}

function bareiss(M) {
  let N = copyMat(M);
  let l = pivot(N);
  if (l === "0-Matrix" || N.length === 1) {
    return N;
  }

  let pivcol = l[1], pivval = N[l[0]][l[1]];
  N = rowSwitch(N,[0,l[0]]);
  for (let i = 1; i < N.length; i++) {
    let to_elim = N[i][l[1]];
    for (let j = pivcol; j < N[0].length; j++) {
      N[i][j] = pivval * N[i][j] - to_elim * N[0][j];
    }
  }
  if (pivcol !== N[0].length-1) {
    N = replmat(N,bareiss(submat(N,[1,pivcol],[N.length-1,N[0].length-1])),[1,pivcol])
  }
  return N;
}

// takes a matrix returns its ranks
function rank(M) {
  N = bareiss(M)
  for (let i = N.length - 1; i > -1; i = i-1) {
    if (!iszero(submat(N,[i,0],[N.length-1,N[0].length-1]))) {
      return i + 1;
    }
  }
  return 0;
}

function vectorpivot(v) {
  for (let i = 0; i < v.length; i++) {
    if (v[i] !== 0) {
      return i;
    }
  }
  return "0-vector"
}

function rowop(M,l,a) {
  N = copyMat(M);
  N[l[0]] = M[l[0]].map((x, index) => x + a * M[l[1]][index]);
  return N;
}

function rowmult(M,row,factor) {
  let N = copyMat(M);
  N[row] = N[row].map(x => x*factor);
  return N;
}

function echelon(M) {
  let N = copyMat(M);
  N = bareiss(M);
  for (i=0; i<N.length; i++) {
    pind = vectorpivot(N[i]);
    if (pind === "0-vector"){
      return N;
    }
    let p = N[i][pind];
    N[i] = N[i].map(x => x/p);
  }
  return N;
}

function redEchelon(M) {
  let N = echelon(M);
  for (let i = 0; i < N.length; i++) {
    let pivot = vectorpivot(N[i]);
    if (pivot !== "0-vector") {
      for (let j = i-1; j >= 0; j--) {
        N = rowop(N,[j,i],-N[j][pivot]);
      }
    }
  }
  return N;
}

function choosefrom(n, arr) {
    var result = new Array(n),
        len = arr.length,
        taken = new Array(len);
    if (n > len)
        throw new RangeError("getRandom: more elements taken than available");
    while (n--) {
        var x = Math.floor(Math.random() * len);
        result[n] = arr[x in taken ? taken[x] : x];
        taken[x] = --len in taken ? taken[len] : len;
    }
    return result;
}

function fromto(a,b) {
  return new Array(b-a+1).fill().map((v,i) => i+a);
}

function redEchelonRandom(size,pivots,range) {
  pivots = pivots.sort().map(x => x-1);
  N = new Array(size[0]).fill(0).map(() => new Array(size[1]).fill(0))
  for (let i = 0; i < size[0]; i++) {
    for (let j = 0; j < size[1]; j++) {
      if (j === pivots[i]) {
        N[i][j] = 1;
      }
      if (j > pivots[i] && pivots.includes(j)) {
        N[i][j] = 0;
      }
      if (j > pivots[i] && !pivots.includes(j)) {
        N[i][j] = choosefrom(1,range)[0];
      }
    }
  }
  return N;
}

// Takes a matrix and performs row operations on it. More precisely it shuffles the row order and afterwards performs random
// row operations of the other types. The second argument is the positive limit of the symmetric range from which the factor
// used in the other types of row operations is drawn. The third argument specifies how many other type row operations are
// respectively done.
// Default values are passed to second and third argument (2 and 20).
function rowshuffle(M,b=1,a=40) {
  let N = copyMat(M);
  let factorarr = fromto(-b,b)
  let rowarr = fromto(0,N.length-1)
  N = choosefrom(N.length,N);
  for (i = 0; i<a; i++) {
    N = rowop(N,choosefrom(2,rowarr),choosefrom(1,factorarr)[0]);
    N = rowmult(N,choosefrom(1,rowarr)[0],choosefrom(1,factorarr.filter(x => x !==0))[0]);
  }
  return N;
}

function random_row_operation(M, rate, factors) {
  let N = copyMat(M);
  let rowarr = fromto(0,N.length-1)
  let factorarr = fromto(-factors,factors)
  let ones = Array(rate[0]).fill(1), twos = Array(rate[1]).fill(2), threes = Array(rate[2]).fill(3);
  let prevalence = ones.concat(ones,twos,threes);
  // let randonrowop = choosefrom(1,prevalence);
  switch (choosefrom(1,prevalence)[0]) {
    case 1:
      N = rowSwitch(N, choosefrom(2,rowarr));
      return N;
    case 2:
      N = rowop(N,choosefrom(2,rowarr),choosefrom(1,factorarr)[0]);
      return N;
    case 3:
      N = rowmult(N,choosefrom(1,rowarr)[0],choosefrom(1,factorarr.filter(x => x !==0))[0]);
      return N;
  }
}
  
// Random Elementary matrix. Julia can't make it work just now
//function randomElementary(n,b=1) {
//  let N = id(n);
//  let factorarr = fromto(-b,b)
//  let rowarr = fromto(0,N.length-1)
//  N = choosefrom(N.length,N);
//  for (i = 0; i<a; i++) {
//    N = rowop(N,choosefrom(2,rowarr),choosefrom(1,factorarr)[0]);
//    N = rowmult(N,choosefrom(1,rowarr)[0],choosefrom(1,factorarr.filter(x => x !==0))[0]);
//  }
//  return N;
//}  

// function roundingerror(M) {
//   let N = copyMat(M);
//   for (i = 0; i<M.length; i++) {
//     for (j = 0; j<M[0].length; j++)
//     if (Math.abs(N[i][j]<0.005)) {
//       N[i][j] = 0;
//     }
//   }
//   return N;
// }

function gcd_2_arg(a, b) {
    if(b === 0) {
        return a;
    }
    const remainder = a % b;
    return gcd_2_arg(b, remainder)
}

function gcd(...args) {
  const gcd = args.reduce((memo, next) => {
      return gcd_2_arg(memo, next)}
  );
  return gcd;
}

function zero_matrix(rownum,colnum) {
  N = new Array(rownum).fill(0).map(() => new Array(colnum).fill(0))
  return N;
}
  
	extension.scope.addFunction(new funcObj('pivot',[TMatrix],TList,function(M) {
    	return pivot(M);
  		}
    ,{unwrapValues: true}));

	extension.scope.addFunction(new funcObj('rank',[TMatrix],TNum,function(M) {
    	return rank(M);
  		}
    ,{unwrapValues: true}));

  	extension.scope.addFunction(new funcObj('bareiss',[TMatrix],TMatrix,function(M) {
    	return bareiss(M);
  		}
    ,{unwrapValues: true}));

  	extension.scope.addFunction(new funcObj('echelon',[TMatrix],TMatrix,function(M) {
    	return echelon(M);
  		}
    ,{unwrapValues: true}));
  
  	extension.scope.addFunction(new funcObj('redEchelon',[TMatrix],TMatrix,function(M) {
    	return redEchelon(M);
  		}
    ,{unwrapValues: true}));
  
   	extension.scope.addFunction(new funcObj('redEchelonRandom',[TList,TList,TList],TMatrix,function(size,pivots,range) {
    	return redEchelonRandom(size,pivots,range);
  		}
    ,{unwrapValues: true}));
  
   	extension.scope.addFunction(new funcObj('fromto',[TNum,TNum],TList,function(a,b) {
    	return fromto(a,b);
  		}
    ,{unwrapValues: true}));
  
   	extension.scope.addFunction(new funcObj('choosefrom',[TNum,TList],TList,function(n,arr) {
    	return choosefrom(n,arr);
  		}
    ,{unwrapValues: true}));
  
   	extension.scope.addFunction(new funcObj('rowshuffle',[TMatrix],TMatrix,function(M) {
    	return rowshuffle(M);
  		}
    ,{unwrapValues: true}));
  
   	extension.scope.addFunction(new funcObj('rowshuffle',[TMatrix,TNum,TNum],TMatrix,function(M,a,b) {
    	return rowshuffle(M,a,b);
  		}
    ,{unwrapValues: true}));
  
  	extension.scope.addFunction(new funcObj('random_row_operation',[TMatrix,TList,TNum],TMatrix,function(M,a,b) {
    	return random_row_operation(M,a,b);
  		}
    ,{unwrapValues: true}));
  
  	extension.scope.addFunction(new funcObj('zero_matrix',[TNum,TNum],TMatrix,function(a,b) {
    	return zero_matrix(a,b);
  		}
    ,{unwrapValues: true}));
  
});