##

    '''
    Code for simulating secure key rate, twofolds, and quantum bit error rate
    Written in Python and QuTIP by Catherine Holloway (c2hollow@iqc.ca).

    Detector model and squashing functions by Catherine Holloway,
    based on code by Dr. Thomas Jennewein (tjennewe@iqc.ca).

    Contributed to the QuTiP project on June 06, 2012 by Catherine Holloway.
    '''

    #imports
    from qutip import *
    from numpy import *
    from pylab import *
    import matplotlib
    import matplotlib.pyplot as plt

    def choose(n, k):
    	"""
    	Binomial coefficient function for the detector model.

    	Parameters
    	----------
    	n : int
    	    Number of elements.
    	k : int
    	    Number of subelements.

    	Returns
    	-------
    	coeff : int
    	    Binomial coefficient.

    	"""
    	if 0 <= k <= n:
    		ntok = 1
    		ktok = 1
    		for t in xrange(1, min(k, n - k) + 1):
    			ntok *= n
    			ktok *= t
    			n -= 1
    		return ntok // ktok
    	else:
    		return 0

    def BucketDetector_realistic_detector(N,efficiency,n_factor):
    	"""
    	Bucket detector model based on H. Lee, U. Yurtsever, P. Kok, G. Hockney, C. Adami, S. Braunstein,
    	and J. Dowling, "Towards photostatistics from photon-number discriminating detectors,"
    	Journal of Modern Optics, vol. 51, p. 15171528, 2004.

    	Parameters
    	----------
    	N : int
    	    The Fock Space dimension.
    	efficiency : float
    	    The channel efficiency.
    	n_factor : float
    	    The average number of dark counts per detection window APD (Bucket Detector).

    	Returns
    	-------
    	[proj, un_proj] : list
    	    The projection and unprojection operators.

    	"""
    	proj=zeros((N,N))
    	#APD (Bucket Detector) un_detector (=gives probability for 0-detection)
    	un_proj=identity(N)
    	#n_factor = 0;
    	for i in range(N):
    	    probs = 0;
    	    for k in range (1,100):
    	        for d in range(k+1):
    	            if k-d<=i:
    	                probs= probs+ (exp(-n_factor)*(n_factor)**(d))/factorial(d)*choose(i,k-d)*efficiency**(k-d)*(1-efficiency)**(i-k+d)

    	    proj[i,i]=probs

    	un_proj = un_proj-proj
    	un_proj = Qobj(un_proj)
    	proj = Qobj(proj)
    	return [proj,un_proj]

    def measure_2folds_4modes_squashing(N,psi,proj,proj2):
    	"""
    	Determines the 2-fold count rate on the joint state
    	outputs for an array of double count probabilities.

    	Parameters
    	----------
    	N : int
    	    The Fock Space dimension.
    	psi : qobj
    	    The entangled state to analyze
    	proj1 : qobj
    	    1st projection operator for the Channel between Alice and
        	the Channel between Bob.
    	proj2 : qobj
    	    2nd projection operator for the Channel between Alice and
    	    the Channel between Bob.

    	Returns
    	-------
    	[HH,HV,VH,VV] : list
    	    Two-fold probabilities.

    	Notes
    	-----
    	The squashing (assigning double pairs to random bases) comes from two papers:

    	    T. Moroder, O. Guhne, N. Beaudry, M. Piani, and N. Lutkenhaus,
    	    "Entanglement verication with realistic measurement devices via squashing operations,"
    	    Phys. Rev. A, vol. 81, p. 052342, May 2010.

    	    N. Lutkenhaus, "Estimates for practical quantum cryptography," Phys. Rev.A,
    	    vol. 59, pp. 3301-3319, May 1999.

    	"""
    	ida=qeye(N)
    	final_state=psi
    	det_exp = zeros((2,2,2,2))

    	#i,j,k,l means Ha,Va,Hb,Vb, 0 means detector clicked, 1 means detector did not click
    	for i in range(2):
    		for j in range(2):
    			for k in range(2):
    				for l in range(2):
    					#expectation values for different detector configurations
    					det_exp[i][j][k][l] = abs(expect(tensor(proj[i],proj[j],proj2[k],proj[l]),final_state))
    	#two fold probabilities
    	HH = det_exp[0][1][0][1]+0.5*(det_exp[0][0][0][1]+det_exp[0][1][0][0])+0.25*det_exp[0][0][0][0]
    	VV = det_exp[1][0][1][0]+0.5*(det_exp[0][0][1][0]+det_exp[1][0][0][0])+0.25*det_exp[0][0][0][0]
    	HV = det_exp[0][1][1][0]+0.5*(det_exp[0][0][1][0]+det_exp[0][1][0][0])+0.25*det_exp[0][0][0][0]
    	VH = det_exp[1][0][0][1]+0.5*(det_exp[0][0][0][1]+det_exp[1][0][0][0])+0.25*det_exp[0][0][0][0]

    	return [HH,HV,VH,VV]

    def sim_qkd_entanglement(eps,loss_a,loss_b,n_factor_a,n_factor_b,N):
    	"""
    	Simulate skr with an SPDC state.

    	Parameters
    	----------
    	eps : float
    	    The squeezing factor, sort of analogous to the amount of
    	    pumping power to the spdc source, but not really.
    	loss_a : float
    	    Efficiency of the quantum channel going to Alice.
    	loss_b : float
    	    Efficiency of the quantum channel going to Bob.
    	n_factor_a : float
    	    Background noise in Alice's detection.
    	n_factor_b : float
    	    Background noise in Bob's detection.
    	N : int
    	    Size of the fock space that we allow for the states

    	Returns
    	-------
    	qber : float
    	    The Quantum Bit Error Rate
    	twofolds : float
    	    Probability of Alice and Bob getting a simultaneous detection
    	    of a photon pair (also referred to as coincidences) within a
    	    timing window.
    	skr : float
    	    Probability of getting a secure key bit within a timing window,
    	    assuming error correction and privacy amplification, in the
    	    limit of many coincidences.

        """
    	#make vaccuum state
    	vacc = basis(N,0)

    	#make squeezing operator for SPDC
    	H_sq = 1j*eps*(tensor(create(N),create(N))+tensor(destroy(N),destroy(N)))

    	#exponentiate hamiltonian and apply it to vaccuum state to make an SPDC state
    	U_sq = H_sq.expm()
    	spdc = U_sq*tensor(vacc,vacc)
    	psi = tensor(spdc,spdc)
    	#since qutip doesn't have a permute function,
    	#we have to do a couple of steps in between
    	#1. turn psi from a sparse matrix to a full matrix
    	out = psi.full()
    	#2. reshape psi into a 4-D matrix
    	out = reshape(out, (N,N,N,-1))
    	#3. permute the dimensions of our 4-D matrix
    	out = transpose(out,(0,3,2,1))
    	#4. turn the matrix back into a 1-D array
    	out = reshape(out,(N*N*N*N,-1))
    	#5. convert the matrix back into a quantum object
    	psi = Qobj(out,dims = [[N, N, N, N], [1, 1, 1, 1]])

    	# model detectors
    	a_det = BucketDetector_realistic_detector(N,loss_a,n_factor_a)
    	b_det = BucketDetector_realistic_detector(N,loss_b,n_factor_b)

    	#measure detection probabilities
    	probs2f=measure_2folds_4modes_squashing(N,psi,a_det,b_det)

    	#Rates returned are 'per pulse', so multiply by source rate
    	twofolds=probs2f[0]+probs2f[1]+probs2f[2]+probs2f[3]
    	#Determine QBER from returned detection probabilities
    	qber = (probs2f[0]+probs2f[3])/twofolds

    	#calculate the entropy of the qber
    	if qber>0:
    		H2=-qber*log2(qber) - (1-qber)*log2(1-qber)
    	else:
    		H2 = 0
    	# estimate error correction efficiency from the CASCADE algorithm
    	f_e = 1.16904371810274 + qber
    	#security analysis - calculate skr in infinite key limit
    	#See Chris Erven's PhD thesis or Xiongfeng Ma's paper
    	#to understand where this equation comes from
    	skr=real(twofolds*0.5*(1-(1+f_e)*H2))
    	return [qber, skr, twofolds]

    if __name__=='__main__':
    	#Lets look at what happens to the secure key rate and
    	#the quantum bit error rate as the loss gets worse.
    	#Analogous to distance with fiber optic links.

    	#define the fock space
    	N = 7
    	#define the squeezing paramter
    	eps = 0.2
    	#define the noise factor
    	n_factor = 4.0e-5
    	#define the length of the coincidence window (in s)
    	coinc_window = 2.0e-9
    	loss_db = arange(0,30)
    	skr = zeros(30)
    	qber = zeros(30)
    	twofolds = zeros(30)

        #run calculation
    	for i in range(30):
    		exp_loss = 10.0**(-loss_db[i]/10.0);
    		[qber[i], skr[i], twofolds[i]] = sim_qkd_entanglement(eps,exp_loss,exp_loss,n_factor,n_factor,N)
    	skr = skr/coinc_window
    	qber = qber*100

        #plot results
    	fig = plt.figure()
    	ax = fig.add_subplot(211)
    	ax.plot(loss_db, skr,lw=2)
    	ax.set_yscale('log')
    	ax.set_ylabel('Secure Key Rate (bits/s)')
    	ax.set_xlabel('Loss (dB)')
    	ax = fig.add_subplot(212)
    	ax.plot(loss_db, qber,lw=2)
    	ax.set_ylabel('Quantum Bit Error Rate (%)')
    	ax.set_ylim([0,15])
    	ax.set_xlabel('Loss (dB)')
    	plt.show()