/* Copyright (C) 2012-2019 IBM Corp.
 * This program is Licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. See accompanying LICENSE file.
 */

#include <helib/permutations.h>

namespace helib {

const NTL::Vec<long> SubDimension::dummyBenes; // global variable

// Apply a permutation to a vector, out[i]=in[p1[i]]
// Unfortunately we need to implement for both NTL::Vec<long> and std::vector<long>
template<class T>
void applyPermToVec(NTL::Vec<T>& out, const NTL::Vec<T>& in, const Permut& p1)
{
  //OLD: assert(&out != &in); // NOT an in-place procedure
  helib::assertNeq<helib::InvalidArgument>(static_cast<const NTL::Vec<T>*>(&out), &in, "Cannot have equal in and out addresses (Not an in-place procedure)");
  out.SetLength(p1.length());
  for (long i=0; i<p1.length(); i++)
    out[i] = in.at(p1[i]);
}
template<class T>
void applyPermToVec(std::vector<T>& out, const std::vector<T>& in, const Permut& p1)
{
  out.resize(p1.length());
  for (long i=0; i<p1.length(); i++)
    out[i] = in.at(p1[i]);
}

// Apply two permutations to a vector out[i]=in[p2[p1[i]]]
template<class T>
void applyPermsToVec(NTL::Vec<T>& out, const NTL::Vec<T>& in,
		      const Permut& p2, const Permut& p1)
{
  //OLD: assert(&out != &in); // NOT an in-place procedure
  helib::assertNeq<helib::InvalidArgument>(static_cast<const NTL::Vec<T>*>(&out), &in, "Cannot have equal in and out addresses (Not an in-place procedure)");
  //OLD: assert(p1.length() == p2.length());
  helib::assertEq(p1.length(), p2.length(), "Permutation p1 and p2 sizes differ");
  out.SetLength(p1.length());
  for (long i=0; i<p1.length(); i++)
    out[i] = in.at(p2.at(p1[i]));
}
template<class T>
void applyPermsToVec(std::vector<T>& out, const std::vector<T>& in,
		      const Permut& p2, const Permut& p1)
{
  //OLD: assert(p1.length() == p2.length());
  helib::assertEq(p1.length(), p2.length(), "Permutation p1 and p2 sizes differ");
  out.resize(p1.length());
  for (long i=0; i<p1.length(); i++)
    out[i] = in.at(p2.at(p1[i]));
}
// explicit instantiations
template void applyPermToVec<long>(NTL::Vec<long>& out, const NTL::Vec<long>& in,
			       const Permut& p1);
template void applyPermToVec<long>(std::vector<long>& out, const std::vector<long>& in,
			       const Permut& p1);
template void applyPermToVec<NTL::ZZX>(std::vector<NTL::ZZX>& out, const std::vector<NTL::ZZX>& in,
			       const Permut& p1);

template void applyPermsToVec<long>(NTL::Vec<long>& out, const NTL::Vec<long>& in,
				const Permut& p2, const Permut& p1);
template void applyPermsToVec<long>(std::vector<long>& out, const std::vector<long>& in,
				const Permut& p2, const Permut& p1);


// Generate a random permutation on [0..n-1]
void randomPerm(Permut& perm, long n)
{
  perm.SetLength(n);
  for (long j = 0; j < n; j++)
     perm[j] = j;
   
  // random shuffle
  for (long m = n; m > 0; m--) {
     long p = NTL::RandomBnd(m);
     // swap positions p and m-1 of perm
     long tmp = perm[p];
     perm[p] = perm[m-1];
     perm[m-1] = tmp;
  }
}


// Write a column permutation explicitly
void ColPerm::makeExplicit(Permut& out) const
{
  long sz = getSize();
  out.SetLength(sz);

  for (long k = 0; k < sz; k++) {
    long i = getCoord(k, dim);
    long pi_i = at(k);
    out.at(k) = addCoord(k, dim, pi_i-i);
  }
}




// For each position in the data vector, compute how many slots it should be
// shifted inside its small permutation.
// Return value is zero if all the shift amounts are zero, nonzero otherwise.
long ColPerm::getShiftAmounts(NTL::Vec<long>& out) const
{
  long sz = getSize();
  out.SetLength(sz);
  long nonZero = 0;

  for (long k = 0; k < sz; k++) {
    long i = getCoord(k, dim);
    long pi_i = at(k);
    if (i != pi_i) nonZero = 1;
    out.at(addCoord(k, dim, pi_i-i)) = i - pi_i;
  }

  return nonZero;
 
}

// Compute the shift amounts corresponding to collapsing 'numLvls' levels
// of the Benes network in 'net' starting from 'startLvl'. Resulting shift
// amounts are returned in 'out', and the return value indicates if the
// resulting permutation from collapsing these levels is the identity.
static bool
collapseBenesLevels(Permut& out, const GeneralBenesNetwork &net,
		    long startLvl, long numLvls)
{
  bool noChange = true;
  // Compute partial permutation computed by the next few Benes levels
  for (long i=0; i<net.getSize(); i++) { // go over all slots
    long i2 = i;
    // compute where this slot is mapped to
    for (long l=startLvl; l<startLvl+numLvls; l++) {
      // how much to shift slot i2 between this level and next
      i2 +=  net.shamt(l) * net.getLevel(l)[i2];
    }
    out[i] = i2 - i; // shift amount for all levels together
    noChange = noChange && (i == i2);
  }
  return noChange;
}

// Get multiple layers of a Benes permutation network. Returns in out[i][j]
// the shift amount to move item j in the i'th layer. Also isID[i]=true if
// the i'th layer is the identity (i.e., contains only 0 shift amounts).
void ColPerm::getBenesShiftAmounts(NTL::Vec<Permut>& out, NTL::Vec<bool>& isID,
				   const NTL::Vec<long>& benesLvls) const
{
  // Go over the columns one by one. For each column extract the columns
  // permutation, prepare a Benes network for it, and then for every layer
  // compute the shift amounts for this columns.

  long n = getDim(dim);     // the permutations are over [0,n-1]

  // Allocate space
  out.SetLength(benesLvls.length());
  isID.SetLength(benesLvls.length());
  for (long k=0; k<benesLvls.length(); k++) {
    out[k].SetLength(getSize());
    isID[k] = true;
  }

  NTL::Vec<long> col;
  col.SetLength(n);

  for (long slice_index = 0; slice_index < numSlices(dim); slice_index++) {
    ConstCubeSlice<long> slice(*this, slice_index, dim);
    for (long col_index = 0; col_index < slice.numCols(); col_index++) {
      getHyperColumn(col, slice, col_index);

      GeneralBenesNetwork net(col); // build a Benes network for this column

      // Sanity checks: width of network == n,
      //                and sum of benesLvls entries == # of levels
      //OLD: assert(net.getSize()==n);
      helib::assertEq(net.getSize(), n, "Network width is different to n");
      {long sum=0;
       for (long k=0; k<benesLvls.length(); k++) sum+=benesLvls[k];
        //OLD: assert(net.getNumLevels()==sum);
        helib::assertEq(net.getNumLevels(), sum, "Sum of benesLvls entries is different to number of levels");
      }

      // Compute the layers of the collapased network for this column
      for (long lvl=0,k=0; k<benesLvls.length(); lvl += benesLvls[k], k++) {

	// Returns in col the shift amounts for this layer in the network,
	// restricted to this column. Also returns true if the returned
	// permutation is the idendity, false otherwise.
	bool id = collapseBenesLevels(col, net, lvl, benesLvls[k]);
	isID[k] = isID[k] && id;

        CubeSlice<long> oslice(out[k], getSig());
        CubeSlice<long> osubslice(oslice, slice_index, dim);
        setHyperColumn(col, osubslice, col_index);
      }  // next collapsed layer
    }  // next column
  } // next slice
}


// Break a permutation into column-row-column format. The input pi permutes
// each dimension-i subcube, and in the output rho1,rho3 permute only along
// the i'th dimension and rho2 permutes each dimension-i+1 subcube.
// This routine cannot permute in-place, it is assumed that pi and rho2 point
// to disjoint vectors.
void breakPermTo3(const HyperCube<long>& pi, long dim, 
                  ColPerm& rho1, HyperCube<long>& rho2, ColPerm& rho3)
{
  //OLD: assert(&rho1.getSig()==&pi.getSig());
  helib::assertEq(&rho1.getSig(), &pi.getSig(), "rho1 and pi signatures differ");
  //OLD: assert(&rho2.getSig()==&pi.getSig());
  helib::assertEq(&rho2.getSig(), &pi.getSig(), "rho2 and pi signatures differ");
  //OLD: assert(&rho3.getSig()==&pi.getSig());
  helib::assertEq(&rho3.getSig(), &pi.getSig(), "rho3 and pi signatures differ");

  // pi consists of separate permutations over [0,n-1], and each
  // of these is viewed as a permutation over an n1 x n2 cube

  long n1 = pi.getDim(dim); // Size of this dimension
  long n2 = pi.getProd(dim+1); 
  long n = pi.getProd(dim); // = n1*n2;

  // representing I_n as I_n1 x I_n2: i == n2*rep[i].first + rep[i].second
  std::vector< std::pair<long,long> > rep(n);
  for (long ind=0,i=0; i<n1; i++) for (long j=0; j<n2; j++,ind++) {
      rep[ind].first = i;
      rep[ind].second = j;
    }

  for (long slice_index = 0; slice_index < pi.numSlices(dim); slice_index++) {
    ConstCubeSlice<long> pi_slice(pi, slice_index, dim);
    CubeSlice<long> rho1_slice(rho1, slice_index, dim);
    CubeSlice<long> rho2_slice(rho2, slice_index, dim);
    CubeSlice<long> rho3_slice(rho3, slice_index, dim);

    // Construct a bipartite n2-by-n2 graph for pi (cf. Lemma 1 in [GHS12a]).
    // For each j=pi(i) with representations i=(i1,i2) and j=(j1,j2), we put
    // in the bipartite graph an edge i2->j2 and label it by i.
    BipartitleGraph bg;
    for (long i=0; i<n; i++) {
      long j = pi_slice[i];  // the image of i under the permutation
      // when i = (i1,i2) and j=(j1,j2), add an edge from i2 to j2 labeled i
      bg.addEdge(rep.at(i).second, rep.at(j).second, i);
    }
    // The bipartite graph is n1-regular, so we can break its edges into
    // n1 perfect matchings, which are numbered 1,2,...,n1.
    bg.partitionToMatchings();

    // The output permutations are defined by the representation i<->(i1,i2),
    // the target permutation pi, and the coloring of the bipartite graph.
    // Denote by sigma(i) the color of the edge labeled by i, sigma(i) \in
    // {1,...,n1}. Also denote by (pi^1(i),pi^2(i)) the representation of
    // pi(i). Then:
    //
    // + rho_1 is defined by           (i1, i2) -> (sigma(i), i2)
    // + rho_2 is defined by     (sigma(i), i2) -> (sigma(i),pi^2(i))
    // + rho_3 is defined by (sigma(i),pi^2(i)) -> (pi^1(i), pi^2(i))
    //
    // rho_1 is a permutation because for every value of i2 (corresponding to
    // a left node in the graph), all the edges leaving that node have different
    // colors.
    //
    // rho_2 is a permutation since the edges of each color form a perfect
    // matching, so for every left-node j2 and color c there is a single
    // c-colored edge going into j2. The label of that edge determines a unique
    // origin index i=(i1,i2), and therefore also the pre-image of (c,j2) under
    // rho_2, which is (sigma(i),i2)=(c,i2).
    //
    // rho_3 is a permutation because rho_1,rho_2,pi are permutations, and
    // so is pi = rho_3 o rho_2 o rho_1.
    //
    // Note that the edges are colored 1..n2 while our represenation above
    // has the second digits in the range 0..n2-1, so below we use sigma(i)-1
    // rather than sigma(i).

    for (long i2=0; i2<n2; i2++) // go over all edges in the bipartite graph
      for (LNeighborList::iterator it=bg.left[i2].neighbors.begin(); 
	   it!=bg.left[i2].neighbors.end(); ++it) {
	LabeledEdge& e = it->second; // An edge in the bipartite graph;
	long i = e.label; // labeled by i
	long c = e.color -1; // colored by c (after the -1 adjustment)
	// i2 = e.from = rep[i].second;
	long j = pi_slice[i];
	long j1 = rep[j].first;
	long j2 = e.to;   // = it->first = rep[j].second;

	long tmp1 = c*n2 + i2;  // the image of i under rho1 = (c,i2)
	rho3_slice[i] = c;
	long tmp2 = c*n2 + j2;  // the image of tmp1 under rho2 =(c,j2)
	rho2_slice[tmp1] = j2;
	rho1_slice[tmp2] = j1; // image of tmp2 under rho3 =(j1,j2)=pi(i)
      }
    // FIXME: The comments above do not match the code, the roles
    //        of rho1,rho3 are switched. Why is this code working??
  }
  rho1.setPermDim(dim);
  rho3.setPermDim(dim);
}



/* Takes a permutation pi over an m-dimensional cube C=Z_{n1} x ... x Z_{nm}
 * and expresses pi as a product pi = rho_{2m-1} o ... o rho_2 o rho_1 where
 * each rho_i is a column permutation along one dimension. Specifically for
 * i<m, the permutations rho_i and rho_{2(m-1)-i} permute the i'th dimension
 ************************************************************************/
void breakPermByDim(std::vector<ColPerm>& out, 
		    const Permut &pi, const CubeSignature& sig)
{
  //OLD: assert(sig.getSize()==pi.length());
  helib::assertEq(sig.getSize(), pi.length(), "Signature sig size is different to pi.length");

  HyperCube<long> tmp1(sig);  
  tmp1.getData() = pi;

  HyperCube<long> tmp2(sig);  

  HyperCube<long>* tp1 = &tmp1;
  HyperCube<long>* tp2 = &tmp2;

  // Allocate the output permutations
  long m = sig.getNumDims();
  ColPerm dummy(sig);
  out.assign(2*m-1, dummy); // allocate space and initialize

  if (m == 1) { // special case, no need to break
    HyperCube<long>& out0 = out[0];
    out0 = tmp1;
    out[0].setPermDim(0);
    return;
  }

  for (long i=0; i<m-2; i++) {
    breakPermTo3(*tp1, i, out[i], *tp2, out[2*m-i-2]);
    std::swap(tp1,tp2);
  }

  breakPermTo3(*tp1, m-2, out[m-2], out[m-1], out[m]);
  out[m-1].setPermDim(m-1);
}


/********************************************************************/
/**********     MAPPING BETWEEN CUBE AND LINEAR ARRAY      **********/
/********************************************************************/

// Get the "crude" cube dimensions corresponding to a vector of trees,
// the ordered vector with one dimension per tree
void GeneratorTrees::getCubeDims(NTL::Vec<long>& dims) const
{
  dims.SetLength(trees.length());

  // copy dims from the trees
  for (long i=0; i<trees.length(); i++) {
    const OneGeneratorTree& T = trees[i];
    dims[T.getAuxKey()] = T.DataOfNode(T.rootIdx()).size;
    // getAuxKey() returns the generator number associated with this tree
  }
}

// Get the "fine" cube dimensions corresponding to a vector of trees,
// the ordered vector with one dimension per leaf in any of the trees.
void GeneratorTrees::getCubeSubDims(NTL::Vec<long>& dims) const
{
  // how many dimensions do we need
  long nDims = 0;
  for (long i=0; i<trees.length(); i++)
    nDims += trees[i].getNleaves();
  dims.SetLength(nDims); // set the size

  // copy dims from the leaves in all the trees
  long idx = 0;
  for (long i=0; i<trees.length(); i++) {
    const OneGeneratorTree& T = trees[i];
    for (long leaf=T.firstLeaf(); leaf>=0; leaf=T.nextLeaf(leaf))
      dims[idx++] = T[leaf].getData().size;
  }
}

// Adds one to the little-endian representation of an integer in base digits,
// returns true if there was an overflow
static bool addOne(NTL::Vec<long>& rep, const NTL::Vec<long> digits)
{
  for (long i=rep.length()-1; i>=0; --i) {
    rep[i]++;
    if (rep[i] >= digits[i])
      rep[i] -= digits[i];
    else
      return false;
  }
  return true;
}

// Compute the mapping between linear array and a hypercube corresponding
/// to a single generator tree
void ComputeOneGenMapping(Permut& genMap, const OneGeneratorTree& T)
{
  NTL::Vec<long> dims(NTL::INIT_SIZE, T.getNleaves());
  NTL::Vec<long> coefs(NTL::INIT_SIZE,T.getNleaves());
  for (long i=T.getNleaves()-1, leaf=T.lastLeaf(); i>=0;
                                i--, leaf=T.prevLeaf(leaf)) {
    dims[i] = T[leaf].getData().size;
    coefs[i] = T[leaf].getData().e;
  }

  // A representation of an integer with digits from dims
  NTL::Vec<long> rep(NTL::INIT_SIZE, T.getNleaves());
  for (long i=0; i<rep.length(); i++) rep[i]=0; // initialize to zero

  // initialize to all zero
  long sz = T[0].getData().size;
  genMap.SetLength(sz);
  for (long i=0; i<sz; i++) genMap[i]=0;

  // compute the permutation
  for (long i=1; i<sz; i++) {
    addOne(rep, dims); // representation of i in base dims
    for (long j=0; j<coefs.length(); j++) {
      long tmp = NTL::MulMod(rep[j], coefs[j], sz);
      genMap[i] = NTL::AddMod(genMap[i], tmp, sz);
    }
  }
}

// Compute the mapping between linear array and the hypercube
// corresponding to all the trees.
void GeneratorTrees::ComputeCubeMapping()
{
  //OLD: assert(trees.length()>=1);
  helib::assertTrue(trees.length()>=1, "Trees length is less than 1");

  if (trees.length()==1) {  // A single tree
    ComputeOneGenMapping(map2array, trees[0]);
  }
  else { // more than one generator
    // Compute the sub-mapping for every generator. Also prepare two hypercube
    // signature objects for the index calculations, with the two ordering of
    // the generators: one for the generators ordered by their index 0,1,2...
    // and the other odred the generators by the order of their trees

    NTL::Vec<long> dims1(NTL::INIT_SIZE,trees.length()), dims2(NTL::INIT_SIZE,trees.length());
    NTL::Vec<Permut> genMappings(NTL::INIT_SIZE, trees.length());
    for (long i=0; i<trees.length(); i++) {
      dims1[i] = trees[i][0].getData().size;
      ComputeOneGenMapping(genMappings[i], trees[i]);
    }
    getCubeDims(dims2);
    CubeSignature sig1(dims1), sig2(dims2);


    // Allocate space for the mapping
    map2array.SetLength(sig1.getSize());

    // Combine the generator perms to a single permutation over the cube
    for (long i=0; i<map2array.length(); i++) {
      long t=0;
      for (long j1=0; j1<trees.length(); j1++) {
	long j2 = trees[j1].getAuxKey();
	long digit = sig1.getCoord(i,j1); // the j1 digit of i in base dims
	digit = genMappings[j1][digit];   // apply the j1 permutation to it
	t += digit * sig2.getProd(j2+1);  // adds the permuted digit

      }
      map2array[i] = t;
    }
  }

  // Compute the inverse permutation
  map2cube.SetLength(map2array.length());
  for (long i=0; i<map2array.length(); i++) map2cube[ map2array[i] ] = i;
}

// Prints out a column permutation
std::ostream& operator<< (std::ostream &s, const ColPerm& p)
{
  Permut pp;
  p.makeExplicit(pp);
  return s << pp;
}

// Prints out a sub-dimension
std::ostream& operator<< (std::ostream &s, const SubDimension& sd)
{
  s << (sd.good? "(g ": "(b ") << sd.size << " " << sd.e << ")";
  if (sd.frstBenes.length()>0 || sd.scndBenes.length()>0)
    s  << sd.frstBenes << sd.scndBenes;
  return s;
}

// Prints out the vector of trees
std::ostream& operator<< (std::ostream &s, const GeneratorTrees &trees)
{
  s << "[" << trees.depth << "\n";
  for (long g=0; g<trees.numTrees(); g++) {
    const OneGeneratorTree &T = trees[g];
    s << " ["; T.printout(s); s<<"]\n";
  }
  return s << "]";
}

}
