# importing the QISKit
from qiskit import QuantumCircuit, QuantumProgram
import Qconfig

# import tomography library
import qiskit.tools.qcvv.tomography as tomo

#visualization packages
from qiskit.tools.visualization import plot_wigner_function, plot_wigner_data

density_matrix = np.matrix([[0.5, 0, 0, 0.5],[0, 0, 0, 0],[0, 0, 0, 0],[0.5, 0, 0, 0.5]])
print(density_matrix)

plot_wigner_function(density_matrix, res=200)

import sympy as sym
from sympy.physics.quantum import TensorProduct

num = int(np.log2(len(density_matrix)))
harr = sym.sqrt(3)
Delta_su2 = sym.zeros(2)   
Delta = sym.ones(1)
for qubit in range(num):
    phi = sym.Indexed('phi', qubit)
    theta = sym.Indexed('theta', qubit)
    costheta  = harr*sym.cos(2*theta)
    sintheta  = harr*sym.sin(2*theta)
    Delta_su2[0,0] =  (1+costheta)/2
    Delta_su2[0,1] = -(sym.exp(2j*phi)*sintheta)/2
    Delta_su2[1,0] = -(sym.exp(-2j*phi)*sintheta)/2
    Delta_su2[1,1] =  (1-costheta)/2
    Delta = TensorProduct(Delta,Delta_su2)
W = sym.trace(density_matrix*Delta)
print(sym.latex(W))

Q_program = QuantumProgram()

number_of_qubits = 2
backend = 'local_qasm_simulator'
shots = 1024
bell_qubits = [0, 1]
qr = Q_program.create_quantum_register('qr',2)
cr = Q_program.create_classical_register('cr',2)
bell = Q_program.create_circuit('bell', [qr], [cr])
bell.h(qr[0])
bell.cx(qr[0],qr[1])

bell_tomo_set = tomo.state_tomography_set([0, 1])
bell_tomo_circuits = tomo.create_tomography_circuits(Q_program, 'bell', qr, cr, bell_tomo_set)
bell_tomo_result = Q_program.execute(bell_tomo_circuits, backend=backend, shots=shots)
bell_tomo_data = tomo.tomography_data(bell_tomo_result, 'bell', bell_tomo_set)
rho_fit_sim = tomo.fit_tomography_data(bell_tomo_data)
plot_wigner_function(np.matrix(rho_fit_sim),res=200)

Delta_su2 = sym.zeros(2)   
Delta = sym.ones(1)
for qubit in range(num):
    phi = sym.Indexed('phi', qubit)
    theta = sym.Indexed('theta', qubit)
    costheta  = harr*sym.cos(2*theta)
    sintheta  = harr*sym.sin(2*theta)
    Delta_su2[0,0] =  (1+costheta)/2
    Delta_su2[0,1] = -(sym.exp(2j*phi)*sintheta)/2
    Delta_su2[1,0] = -(sym.exp(-2j*phi)*sintheta)/2
    Delta_su2[1,1] =  (1-costheta)/2
    Delta = TensorProduct(Delta,Delta_su2)
W = sym.trace(np.matrix(rho_fit_sim)*Delta)
print(sym.latex(W))

Q_program.set_api(Qconfig.APItoken, Qconfig.config['url'])
backend = 'ibmqx2'
max_credits = 8
shots = 1024
bell_qubits = [0, 1]
bell_tomo_set = tomo.state_tomography_set(bell_qubits)
bell_tomo_circuits = tomo.create_tomography_circuits(Q_program, 'bell', qr, cr, bell_tomo_set)
bell_tomo_result = Q_program.execute(bell_tomo_circuits, backend=backend, shots=shots, 
                                     max_credits=max_credits, timeout=300)
bell_tomo_data = tomo.tomography_data(bell_tomo_result, 'bell', bell_tomo_set)
rho_fit_ibmqx = tomo.fit_tomography_data(bell_tomo_data)
plot_wigner_function(np.matrix(rho_fit_ibmqx), res=100)

print(rho_fit_ibmqx)

Delta_su2 = sym.zeros(2)   
Delta = sym.ones(1)
for qubit in range(num):
    phi = sym.Indexed('phi', qubit)
    theta = sym.Indexed('theta', qubit)
    costheta  = harr*sym.cos(2*theta)
    sintheta  = harr*sym.sin(2*theta)
    Delta_su2[0,0] =  (1+costheta)/2
    Delta_su2[0,1] = -(sym.exp(2j*phi)*sintheta)/2
    Delta_su2[1,0] = -(sym.exp(-2j*phi)*sintheta)/2
    Delta_su2[1,1] =  (1-costheta)/2
    Delta = TensorProduct(Delta,Delta_su2)
W = sym.trace(np.matrix(rho_fit_ibmqx)*Delta)
print(sym.latex(W))

theta1_points = 8
theta2_points = 8
number_of_points = theta1_points*theta2_points
the1 = [0]*number_of_points
the2 = [0]*number_of_points #initialize theta values
phis = [[0]*number_of_points]*number_of_qubits #set phi values to 0

point = 0
for i in range(theta1_points):
    for k in range(theta2_points):
        
        the1[point] = 2*i*np.pi/theta1_points
        the2[point] = 2*k*np.pi/theta2_points #create the values of theta for all points on plot
        point += 1
        
thetas = np.vstack((the1,the2))
bell_circuits = tomo.build_wigner_circuits(Q_program, 'bell', phis, thetas, 
                                      bell_qubits, qr, cr)

backend = 'local_qasm_simulator'
shots = 1024
bell_result = Q_program.execute(bell_circuits, backend=backend, shots=shots)
print(bell_result)

wdata = tomo.wigner_data(bell_result, bell_qubits,
                        bell_circuits, shots=shots)
wdata = np.matrix(wdata)
wdata = wdata.reshape(theta1_points,theta2_points) 
plot_wigner_data(wdata, method='plaquette')

equator_points = 64
theta = [np.pi/2]*equator_points
phi = [0]*equator_points

point = 0
for i in range(equator_points):
    phi[i] = 2*i*np.pi/equator_points
        
thetas = np.vstack((theta,theta))
phis = np.vstack((phi,phi))
bell_eq_circuits = tomo.build_wigner_circuits(Q_program, 'bell', phis, thetas, 
                                      bell_qubits, qr, cr)
bell_eq_result = Q_program.execute(bell_eq_circuits, backend=backend, shots=shots)
wdata_eq = tomo.wigner_data(bell_eq_result,  bell_qubits,
                        bell_eq_circuits, shots=shots)
plot_wigner_data(wdata_eq, method='curve')

Q_program = QuantumProgram()
number_of_qubits = 5
backend = 'local_qasm_simulator'
shots = 1024
ghz_qubits = [0, 1, 2, 3, 4]
qr = Q_program.create_quantum_register('qr',5)
cr = Q_program.create_classical_register('cr',5)
ghz = Q_program.create_circuit('ghz', [qr], [cr])
ghz.h(qr[0])
ghz.h(qr[1])
ghz.x(qr[2])
ghz.h(qr[3])
ghz.h(qr[4])
ghz.cx(qr[0],qr[2])
ghz.cx(qr[1],qr[2])
ghz.cx(qr[3],qr[2])
ghz.cx(qr[4],qr[2])
ghz.h(qr[0])
ghz.h(qr[1])
ghz.h(qr[2])
ghz.h(qr[3])
ghz.h(qr[4])

equator_points = 64
thetas = [[np.pi/2]*equator_points]*number_of_qubits
phi = [0]*equator_points

point = 0
for i in range(equator_points):
    phi[i] = 2*i*np.pi/equator_points
phis = np.vstack((phi,phi,phi,phi,phi))
ghz_eq_circuits = tomo.build_wigner_circuits(Q_program, 'ghz', phis, thetas, 
                                      ghz_qubits, qr, cr)
ghz_eq_result = Q_program.execute(ghz_eq_circuits, backend=backend, shots=shots, timeout = 300)
wghzdata_eq = tomo.wigner_data(ghz_eq_result, ghz_qubits,
                        ghz_eq_circuits, shots=shots)
plot_wigner_data(wghzdata_eq, method='curve')

import matplotlib.pyplot as plt
plt.plot(phi, wghzdata_eq, 'o')
plt.axis([0, 2*np.pi, -0.6, 0.6])
plt.show()

density_matrix = np.zeros((32,32))
density_matrix[0][0] = 0.5
density_matrix[0][31] = -0.5
density_matrix[31][0] = -0.5
density_matrix[31][31] = 0.5
plot_wigner_function(density_matrix, res=200)

