import numpy as np
import itertools
import qiskit 
from qiskit import QuantumRegister, QuantumCircuit
from qiskit import Aer
import qiskit.ignis.verification.tomography as tomo
from qiskit.quantum_info import state_fidelity

I = np.array([[1,0],[0,1]])
X = np.array([[0,1],[1,0]])
Y = np.array([[0,-1j],[1j,0]])
Z = np.array([[1,0],[0,-1]])

def HS_product(A,B):
    return 0.5*np.trace(np.conj(B).T @ A)

PX0 = 0.5*np.array([[1, 1], [1, 1]])
PX1 = 0.5*np.array([[1, -1], [-1, 1]])

PY0 = 0.5*np.array([[1, -1j], [1j, 1]])
PY1 = 0.5*np.array([[1, 1j], [-1j, 1]])

PZ0 = np.array([[1, 0], [0, 0]])
PZ1 = np.array([[0, 0], [0, 1]])

projectors = [PX0, PX1, PY0, PY1, PZ0, PZ1]

M = np.array([p.flatten() for p in projectors])
M

M_dg = np.conj(M).T
linear_inversion_matrix = np.linalg.inv(M_dg @ M) @ M_dg

projectors_2 = [np.kron(p1, p2) for (p1, p2) in itertools.product(projectors, repeat = 2)]
M_2 = np.array([p.flatten() for p in projectors_2])
M_dg_2 = np.conj(M_2).T
linear_inversion_matrix_2 = np.linalg.inv(M_dg_2 @ M_2) @ M_dg_2

q2 = QuantumRegister(2)
bell = QuantumCircuit(q2)
bell.h(q2[0])
bell.cx(q2[0], q2[1])
bell.qasm()

qst_bell = tomo.state_tomography_circuits(bell, q2)
job = qiskit.execute(qst_bell, Aer.get_backend('qasm_simulator'), shots=5000)

statefit = tomo.StateTomographyFitter(job.result(), qst_bell)

statefit.data

p, M, weights = statefit._fitter_data(True, 0.5)

M_dg = np.conj(M).T
linear_inversion_matrix = np.linalg.inv(M_dg @ M) @ M_dg
rho_bell = linear_inversion_matrix @ p
rho_bell = np.reshape(rho_bell, (4, 4))
print(rho_bell)

job = qiskit.execute(bell, Aer.get_backend('statevector_simulator'))
psi_bell = job.result().get_statevector(bell)
F_bell = state_fidelity(psi_bell, rho_bell)
print('Fit Fidelity linear inversion =', F_bell)

rho_cvx_bell = statefit.fit(method='cvx')
F_bell = state_fidelity(psi_bell, rho_cvx_bell)
print('Fit Fidelity CVX fit =', F_bell)

rho_mle_bell = statefit.fit(method='lstsq')
F_bell = state_fidelity(psi_bell, rho_mle_bell)
print('Fit Fidelity MLE fit =', F_bell)

