import qiskit as qsk
import numpy as np
import matplotlib.pyplot as plt

def measure_X(circuit,n):    
    circuit.barrier(n)
    circuit.h(n)
    circuit.measure(n,n)
    return circuit

def measure_Y(circuit,n):    
    circuit.barrier(n)
    circuit.sdg(n)
    circuit.h(n)    
    circuit.measure(n,n)    
    return circuit

def tomography(circuit):
    """ Tomography of a one qubit Circuit.
    """
    qc_list = []
    base = ['X', 'Y', 'Z']
    for basis in base:
        Q = qsk.QuantumCircuit(1,1)
        Q.append(circuit, [0])
        if basis == 'X':
            measure_X(Q, 0)
            qc_list.append(Q)
        if basis == 'Y':
            measure_Y(Q,  0)
            qc_list.append(Q)
        if basis == 'Z':
            Q.measure(0,0)
            qc_list.append(Q)
    return qc_list, base

qc = qsk.QuantumCircuit(1)
qc.h(0)
qcs, bases = tomography(qc)

backend_sim = qsk.Aer.get_backend('qasm_simulator')
job = qsk.execute(qcs, backend_sim, shots=5000)
result = job.result()

for index, circuit in enumerate(qcs):
    print(result.get_counts(circuit))
    print(f'Base measured {bases[index]}\n')

def get_density_matrix(measurements,circuits):
    """Get density matrix from tomography measurements.

    """
    density_matrix = np.eye(2, dtype=np.complex128)
    sigma_x = np.array([[0,1],[1,0]])
    sigma_y = np.array([[0,-1j],[1j,0]])
    sigma_z = np.array([[1,0],[0,-1]])
    basis = [sigma_x, sigma_y, sigma_z]

    for index in range(len(circuits)):
        R = measurements.get_counts(index)
        
        if '0' in R.keys() and '1' in R.keys():
            zero = R['0']
            one = R['1']
        
        elif '1' in R.keys():
            zero = 0
            one = R['1']
        
        elif '0' in R.keys():
            zero = R['0']
            one = 0

        total = sum(list(R.values()))
        expected = (zero - one)/total        
        density_matrix += expected * basis[index]

    return 0.5*density_matrix

density = get_density_matrix(result,qcs)
print(density)

def plot_density_matrix(DM):
    """Helper function to plot density matrices.

    Parameters
    ----------------------------------------------
    DM(np.array): Density Matrix.
    
    """
    from matplotlib.ticker import MaxNLocator
    
    fig = plt.figure(figsize=(16,10))
    gs = fig.add_gridspec(1, 2)
    ax1 = fig.add_subplot(gs[0, 0])
    ax2 = fig.add_subplot(gs[0, 1])

    im = ax1.imshow(np.real(DM), cmap='Greys')
    ax1.yaxis.set_major_locator(MaxNLocator(integer=True))
    ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
    ax1.set_title('Real Part',size=16)
    plt.colorbar(im, ax=ax1)

    im = ax2.imshow(np.imag(DM), cmap='Greys')
    plt.colorbar(im, ax=ax2)
    ax2.yaxis.set_major_locator(MaxNLocator(integer=True))
    ax2.xaxis.set_major_locator(MaxNLocator(integer=True))
    ax2.set_title('Imaginary Part',size=16)
    plt.show()

plot_density_matrix(density)

from qiskit.ignis.verification.tomography import state_tomography_circuits

tomography_circuits = state_tomography_circuits(qc, [0], meas_labels='Pauli', meas_basis='Pauli')

tomography_circuits[1].draw('mpl')

backend_sim = qsk.Aer.get_backend('qasm_simulator')
job = qsk.execute(tomography_circuits, backend_sim, shots=5000)
result = job.result()

from qiskit.ignis.verification.tomography import StateTomographyFitter

state_fitter = StateTomographyFitter(result, tomography_circuits, meas_basis='Pauli')

density_matrix = state_fitter.fit(method='lstsq')
print(density_matrix)

plot_density_matrix(density_matrix)

qc = qsk.QuantumCircuit(3)
qc.h(0)
qc.cx(0,1)
qc.cx(0,2)
qc.draw('mpl')

from qiskit.ignis.verification.tomography import state_tomography_circuits

tomography_circuits = state_tomography_circuits(qc, [0,1,2], 
                                                meas_labels='Pauli', 
                                                meas_basis='Pauli')
backend_sim = qsk.Aer.get_backend('qasm_simulator')
job = qsk.execute(tomography_circuits, backend_sim, shots=5000)
result = job.result()

from qiskit.ignis.verification.tomography import StateTomographyFitter

state_fitter = StateTomographyFitter(result, tomography_circuits, meas_basis='Pauli')

density_matrix = state_fitter.fit(method='lstsq')

plot_density_matrix(density_matrix)

import qiskit.tools.jupyter
%qiskit_version_table