import numpy as np
from numpy import random as rd

im = np.sqrt(-1 + 0j)

# Adjoint operator
def dag(a):
	return np.transpose(np.conj(a))

# Matrix product
def mp(ms):
	out = np.dot(ms[0], ms[1])
	if(len(ms) > 2):
		for m in ms[2:]:
			out = np.dot(out, m)
	return out

# Kronecker delta
def kron(a, b):
	return 1 if a == b else 0

# Generate 2-dimensional state
def genstate():
    theta = np.pi * rd.rand()
    phi = 2 * np.pi * rd.rand()
    return [[np.cos(theta / 2)], [np.exp(im * phi) * np.sin(theta / 2)]]

def genunitary(d):
	mat = np.zeros((d, d)) + im * np.zeros((d,d)) # Create blank complex matrix

	# Fill matrix with IID normally distributed complex numbers
	for (x, y), value in np.ndenumerate(mat):
		mat[x, y] = rd.normal(0, 1) + im * rd.normal(0, 1)

	U = np.linalg.qr(mat)[0] # Perform QR decomposition, take Q
	return U

# randfixedsum to generate probability distribution by Roger Stafford, ported by Björn Brandenburg
def genprob(n, u = 1, nsets = 1):
	#deal with n=1 case
    if n == 1:
        return np.tile(np.array([u]),[nsets,1])

    k = np.floor(u)
    s = u
    step = 1 if k < (k-n+1) else -1
    s1 = s - np.arange( k, (k-n+1)+step, step )
    step = 1 if (k+n) < (k-n+1) else -1
    s2 = np.arange( (k+n), (k+1)+step, step ) - s

    tiny = np.finfo(float).tiny
    huge = np.finfo(float).max

    w = np.zeros((n, n+1))
    w[0,1] = huge
    t = np.zeros((n-1,n))

    for i in np.arange(2, (n+1)):
        tmp1 = w[i-2, np.arange(1,(i+1))] * s1[np.arange(0,i)]/float(i)
        tmp2 = w[i-2, np.arange(0,i)] * s2[np.arange((n-i),n)]/float(i)
        w[i-1, np.arange(1,(i+1))] = tmp1 + tmp2;
        tmp3 = w[i-1, np.arange(1,(i+1))] + tiny;
        tmp4 = np.array( (s2[np.arange((n-i),n)] > s1[np.arange(0,i)]) )
        t[i-2, np.arange(0,i)] = (tmp2 / tmp3) * tmp4 + (1 - tmp1/tmp3) * (np.logical_not(tmp4))

    m = nsets
    x = np.zeros((n,m))
    rt = np.random.uniform(size=(n-1,m)) #rand simplex type
    rs = np.random.uniform(size=(n-1,m)) #rand position in simplex
    s = np.repeat(s, m);
    j = np.repeat(int(k+1), m);
    sm = np.repeat(0, m);
    pr = np.repeat(1, m);

    for i in np.arange(n-1,0,-1): #iterate through dimensions
        e = ( rt[(n-i)-1,...] <= t[i-1,j-1] ) #decide which direction to move in this dimension (1 or 0)
        sx = rs[(n-i)-1,...] ** (1/float(i)) #next simplex coord
        sm = sm + (1-sx) * pr * s/float(i+1)
        pr = sx * pr
        x[(n-i)-1,...] = sm + pr * e
        s = s - e
        j = j - e #change transition table column if required

    x[n-1,...] = sm + pr * s

    #iterated in fixed dimension order but needs to be randomised
    #permute x row order within each column
    for i in range(0,m):
        x[...,i] = x[np.random.permutation(n),i]

    return np.transpose(x)[0];

# Generate n x n density operator
def genrho(n):
	# Create standard basis of Rn
	basis = []
	for i in range(n):
		vec = []
		for j in range(n):
			vec.append(kron(i,j))
		basis.append(vec)
	# Convert to column vectors
	basis = [np.array(np.transpose(np.matrix(x))) for x in basis]

	p = genprob(n) # Generate uniform probability distribution using randfixedsum
	U = genunitary(n) # Generate unitary matrix

	# Create density operator
	rho = np.zeros((n,n)).astype(complex)
	for i in range(n):
		rho += p[i] * mp([U, basis[i], dag(basis[i]), dag(U)])
	return rho