# Canonical ordering of symplectic group elements
# according to: Robert Koenig and John A. Smolin,
#	"How to efficiently select an arbitrary Clifford group element",
#	Journal of Mathematical Physics, 55 (12), 2014.
# Edited by: Michel Barbeau, Carleton University
# Version: November 21, 2016 
#
from numpy import *
from time import clock

# TRANSVECTION

def directsum(m1, m2):
	"""
	Combination of two matrices.

	Parameters
	----------
	m1
	   A n1*n1 saquare matrix
	m2
	   A n2*n2 square matrix

	Returns
	-------
	Matrix
	   A (n1+n2)*(n1+n2) square matrix. Matrix occupies the upper left,
	   matrix m2 occupies lower right
	"""
	n1=len(m1[0])
	n2=len(m2[0])
	output=zeros((n1+n2,n1+n2),dtype=int8)
	for i in range(0,n1):
		for j in range(0,n1):
			output[i,j]=m1[i,j]
	for i in range(0,n2):
		for j in range(0,n2):
			output[i+n1,j+n1]=m2 [i,j]
	return output

def inner(v,w): 
	"""
	Symplectic inner product.

	Parameters
	----------
	v
	   A vector
	w
	   A vector

	Returns
	-------
	real
	   A 
	"""
	t=0
	for i in range(0,size(v)>>1):
		t+=v[2*i]*w[2*i+1]
		t+=w[2*i]*v[2*i+1]
	return t%2

def transvection(k,v):
	"""
	Applies transvection Z_k to v.

	Parameters
	----------
	k
	   A ... 
	v
	   A vector

	Returns
	-------
	real
	A ...
	"""
	return(v+inner(k,v)*k)%2

def int2bits(i,n):
	"""
	Converts integer i to a length n array of bits.

	Parameters
	----------
	i
	   An integer
	n
	   An integer

	Returns
	-------
	real
	   A ...
	"""
	output=zeros(n,dtype=int8)
	for j in range(0,n):
		output[j]=i&1
		i>>=1
	return output

def findtransvection (x,y): 
	"""
	Finds h1,h2 such that y=Z_h1 Z_h2 x (Lemma 2)
	Note that if only one transvection is required output[1] will be
	zero and applying the all-zero transvection does nothing.

	Parameters
	----------
	x
	   A ...
	y
 	   A ...

	Returns
	-------
	real
	   A ...
	"""
	output=zeros((2,size(x)),dtype=int8)
	if array_equal(x,y):
		return output
	if inner(x,y)==1:
		output[0]=(x+y)%2
		return output
	
	# find a pair where they are both not 00
	z=zeros(size(x))
	for i in range(0,size(x)>>1):
		ii=2*i
		if ((x[ii]+x[ii+1])!=0) and ((y[ii]+y [ii+1])!=0): # found the pair
			z[ii]=(x[ii]+y[ii])%2
			z[ii+1]=(x[ii+1]+y[ii+1])%2
			if (z[ii]+z[ii+1])==0: # they were the same so they added to 00
				z[ii+1]=1
				if x[ii]!=x[ii+1]:
					z[ii]=1
			output[0]=(x+z)%2
			output[1]=(y+z)%2
			return output
	# didn't find a pair
	# so look for two places where x has 00 and y doesn't, and vice versa

	# first y==00 and x doesn't
	for i in range(0,size(x)>>1):
		ii=2*i
		if ((x[ii]+x[ii+1])!=0) and ((y[ii]+y[ii+1])==0): # found the pair
			if x[ii]==x[ii+1]:
				z [ii+1]=1
			else:
				z[ii+1]=x[ii]
				z[ii]=x[ii+1]
			break

	# finally x==00 and y doesn't
	for i in range(0,size(x)>>1):
		ii=2*i
		if ((x[ii]+x[ii+1])==0) and ((y[ii]+y[ii+1])!=0): # found the pair
			if y[ii]==y[ii+1]:
				z[ii+1]=1
			else:
				z[ii+1]=y[ii]
				z[ii]=y[ii+1]
				break
	output[0]=(x+z)%2
	output[1]=(y+z)%2
	return output 

def symplectic(i,n):
	"""
	Maps integer i to an element of the Clifford group |Cn|.

	Examples
	   import SYMPLECTIC
	   SYMPLECTIC.symplectic(2,1)
	   SYMPLECTIC.symplectic(4,6)
	Parameters
	----------
	i
	   An integer in the range 0..|Cn|.
	n
	   An integer, the dimension of the Clifford.

	Returns
	-------
 	A 2^n * 2^n matrix  
	"""
	nn=2*n
	# step 1
	s=((1<<nn)-1)
	k=(i%s)+1
	i/=s
	# step 2
	f1=int2bits(k,nn)
	# step 3
	e1=zeros(nn,dtype=int8) # define first basis vectors
	e1[0]=1
	T=findtransvection(e1,f1) # use Lemma 2 to compute T
	# step 4
	# b [0]=b in the text, b[1]...b[2n-2] are b_3...b_2n in the text
	bits=int2bits(i%(1<<(nn-1)),nn-1)
	# step 5
	eprime=copy(e1)
	for j in range(2,nn):
		eprime[j]=bits[j-1]
	h0=transvection(T[0],eprime)
	h0=transvection(T[1],h0)
	# step 6
	if bits[0]==1:
		f1*=0
	# T' from the text will be Z_f1 Z_h0. If f1 has been set to zero
	# it doesn't do anything
	# We could now compute f2 as said in the text but step 7 is slightly
	# changed and will recompute f1 , f2 for us anyway

	# step 7
	# define the 2*2 identity matrix
	id2=zeros((2,2),dtype=int8)
	id2[0,0]=1
	id2[1,1]=1

	if n!=1:
		g=directsum(id2,symplectic(i>>(nn-1),n-1))
	else :
		g=id2

	for j in range(0,nn):
		g [j]=transvection(T[0],g[j])
		g [j]=transvection(T[1],g[j])
		g [j]=transvection(h0,g[j])
		g [j]=transvection(f1,g[j])

	return g

def bits2int(b,nn): 
	"""
	Converts an nn-bit string b to an integer between 0 and 2^n-1.
	"""
	output=0
	tmp=1
	for j in range(0,nn):
		if b[j]==1:
			output=output+tmp
		tmp=tmp*2
	return output

def numberofcosets(n):
	"""
	Returns the number of different cosets.
	"""
	x=power(2,2*n-1)*(power(2,2*n)-1)
	return x ;

def numberofsymplectic(n): 
	"""
	Returns the number of symplectic group elements.
	"""
	x=1;
	for j in range(1,n+1):
		x=x*numberofcosets(j);
	return x;