#
# qc.py - Quantum Computing Library for Python
#
#         by Pius Fischer, February 13-20, 2016
#
#         Various functions for mathematically simulating the quantum circuit model of computation.
#
#         Example 1 - Superdense coding (sending two classical bits a1 and a2 from Alice to Bob via
#                     an entangled pair of qubits)
#
#         from qc import *              <-- Import this library.
#         qA, qB = 'A', 'B'             <-- Define the names to be used for the qubits.
#         a1, a2 = 1, 0                 <-- Initialize the bits to be sent.
#         prepareBell(qA, qB)           <-- Create the entangled pair qA and qB.
#         encodeBell(a1, a2, qA)        <-- Alice encodes a1 and a2 onto qA (which also affects qB).
#         b1, b2 = measureBell(qA, qB)  <-- Bob recovers b1 and b2 by Bell measurement of qA and qB.
#
#         See the sendSuperdense() function below.
#
#         Example 2 - Quantum teleportation (transferring the state of a qubit qC from Alice to Bob
#                     via an entangled pair of qubits qA and qB and two classical bits b1 and b2)
#
#         from qc import *              <-- Import this library.
#         qA, qB, qC = 'A', 'B', 'C'    <-- Define the names to be used for the qubits.
#         createQubit(qC, 0.8, 0.6)     <-- Initialize the qubit to be teleported.
#         prepareBell(qA, qB)           <-- Create the entangled pair qA and qB.
#         b1, b2 = measureBell(qC, qA)  <-- Alice gets b1 and b2 by Bell measurement of qC and qA.
#         encodeBell(b1, b2, qB)        <-- Bob encodes b1 and b2 onto qB. Now qB is in the same
#                                           state that qC was in before Alice's Bell measurement.
#
#         See the teleportQubit() function below.
#
import math
import random

__all__ = (
	'IdentityMatrix',
	'IdentityMatrixN',
	'HadamardGate',
	'XGate',
	'YGate',
	'ZGate',
	'PhaseShiftGate',
	'ControlledGate',
	'ControlledNotGate',
	'SwapGate',
	'SqrtSwapGate',
	'FourierTransform',

	'multiplyMatrixByScalar',
	'multiplyMatrixByMatrix',
	'combineTransforms',
	'roundedMatrix',
	'compareMatrices',
	'compareVectors',
	'roundedStateVector',
	'compareStateVectors',

	'clearSystem',
	'printSystem',
	'printQubit',
	'createQubit',
	'removeQubit',
	'measureQubit',
	'applyGate',
	'qubitArray',

	'prepareBell',
	'encodeBell',
	'measureBell',

	'sendSuperdense',
	'teleportQubit',
	'quantumFourierTransform',
)

def validState(V, minSize=2):
	vlen = len(V)
	assert vlen & (vlen - 1) == 0
	assert vlen >= minSize
	return vlen

def validMatrix(U, size=None):
	numRows = len(U)

	if size is None:
		assert numRows & (numRows - 1) == 0
	else:
		assert numRows == size

	for row in U:
		assert len(row) == numRows

	return numRows

def multiplyMatrixByScalar(scalar, U):
	validMatrix(U)

	return [[scalar * element for element in row] for row in U]

def multiplyMatrixByMatrix(U1, U2):
	validMatrix(U2, validMatrix(U1))

	U2 = zip(*U2)
	return [[sum([e1 * e2 for e1, e2 in zip(row1, row2)]) for row2 in U2] for row1 in U1]

def simplify(N):
	real = round(N.real, 14)
	imag = round(N.imag, 14)

	if real == 0: # Sometimes real is rounded to -0 ... Really set it to 0.
		real = 0
	if imag == 0:
		return real

	return complex(real, imag)

def roundedMatrix(U):
	return [[simplify(col) for col in row] for row in U]

def compareMatrices(X, Y, tolerance=1.5e-14, verbose=True):
	if len(X) != len(Y):
		if verbose:
			print "numRows(X) ({}) != numRows(Y) ({})".format(len(X), len(Y))
		return False

	rX = roundedMatrix(X)
	rY = roundedMatrix(Y)

	equal = True
	for row, (rowX, rowY) in enumerate(zip(rX, rY)):
		if len(rowX) != len(rowY):
			if verbose:
				print "numCols(X[{0}]) ({1}) != numCols(Y[{0}]) ({2})".format(row,
					len(rowX), len(rowY))
			return False
		for col, (a, b) in enumerate(zip(rowX, rowY)):
			if a != b and abs(a - b) > tolerance:
				equal = False
				if verbose:
					print "X[{0},{1}] != Y[{0},{1}]:".format(row, col)
					print "X[{},{}] = {: .16f}".format(row, col, a)
					print "Y[{},{}] = {: .16f}".format(row, col, b)
					print
	return equal

IdentityMatrix = ((1, 0), (0, 1))

def IdentityMatrixN(N):
	return [[1 if row == col else 0 for col in xrange(N)] for row in xrange(N)]

HadamardGate = multiplyMatrixByScalar(1/math.sqrt(2), ((1, 1), (1, -1)))

XGate = ((0, 1), (1, 0))
YGate = ((0, -1J), (1J, 0))
ZGate = ((1, 0), (0, -1))

def PhaseShiftGate(phi):
	return ((1, 0), (0, complex(math.cos(phi), math.sin(phi))))

def ControlledGate(U):
	validMatrix(U, 2)

	return ((1, 0, 0, 0),
		(0, 1, 0, 0),
		(0, 0, U[0][0], U[0][1]),
		(0, 0, U[1][0], U[1][1]))

ControlledNotGate = ControlledGate(XGate)

SwapGate = (
		(1, 0, 0, 0),
		(0, 0, 1, 0),
		(0, 1, 0, 0),
		(0, 0, 0, 1))

SqrtSwapGate = (
		(1,     0,          0,      0),
		(0, 0.5 + 0.5J, 0.5 - 0.5J, 0),
		(0, 0.5 - 0.5J, 0.5 + 0.5J, 0),
		(0,     0,          0,      1))

def FourierTransform(N):