from qiskit import *
from qiskit.visualization import plot_histogram
from qiskit.visualization import plot_state_qsphere, plot_bloch_vector

from qiskit.quantum_info import Statevector

import numpy as np

sim = Aer.get_backend('aer_simulator')

qr = QuantumRegister(3)
qc_3qx = QuantumCircuit(qr)

### your code goes here. ###

qc_3qx.h(0)
qc_3qx.s(0)
qc_3qx.cx(0,1)
qc_3qx.cx(0,2)

#######
qc_3qx.draw('mpl')

from qiskit.visualization import array_to_latex

state_ve = Statevector.from_instruction(qc_3qx)
array_to_latex(state_ve)

### your code goes here ###

state_ve = Statevector.from_instruction(qc_3qx)
plot_state_qsphere(state_ve)


def apply_err(n, err):
    
    qc = QuantumCircuit(int(n), name='Error')
    which_qubit = np.random.randint(n)
    
    if err=='bit':
        qc.x(which_qubit)
    elif err=='phase':
        qc.z(which_qubit)
    else:
        pass
    
    err = qc.to_gate()
    
    return err, which_qubit

err, which_qubit = apply_err(3, 'bit')
print('Error applied to qubit: ', which_qubit)

qc_3qx.append(err,range(3))
qc_3qx.draw('mpl')

# Execute this cell to add the extra registers

k = int(input('number of auxiliary qubits ( / syndrome bits): '))
ar = QuantumRegister(k, 'auxiliary')
cr = ClassicalRegister(k, 'syndrome')
qc_3qx.add_register(ar)
qc_3qx.add_register(cr)

# Apply the parity check gates and measure the parities on the syndrome bits to localize a single bit-flip ( X ) error on the code.

### your code goes here. ###

qc_3qx.draw('mpl')

qc_3qx.cx(qr[0], ar[0])
qc_3qx.cx(qr[1], ar[0])
qc_3qx.cx(qr[1], ar[1])
qc_3qx.cx(qr[2], ar[1])
qc_3qx.barrier()
qc_3qx.measure(ar, cr)

######

qc_3qx.draw('mpl')

#### complete the dictionary ###

table_syndrome = {'00': 'I[0]I[1]I[2]', '01':'I[0]I[1]X[2]',
                  '10': 'X[0]I[1]I[2]', '11':'I[0]X[1]I[2]'}

######
print(table_syndrome)

qc_3qx_trans = transpile(qc_3qx, sim)
syndrome = sim.run(qc_3qx_trans, shots=1, memory=True).result().get_memory()
print(syndrome)

your_answer = input('Enter the index of the code qubit that underwent bit-flip error: ')
print('\n')
print(which_qubit == int(your_answer))

### your code goes here ###
def get_logical_x():
    qr_xl = QuantumRegister(3, 'qubit')
    ar_xl = QuantumRegister(2, 'auxiliary')
    sr_xl = ClassicalRegister(2, 'syndrome')
    
    qc_xl = QuantumCircuit(qr_xl, ar_xl, sr_xl)
    qc_xl.x(range(3))
    
    qc_xl.cx([qr_xl[0], qr_xl[1]], [ar_xl[0], ar_xl[0]])
    qc_xl.cx([qr_xl[1], qr_xl[2]], [ar_xl[1], ar_xl[1]])
    
    qc_xl.barrier()
    qc_xl.measure(ar_xl, sr_xl)
    
    # correct errors
    qc_xl.x(qr_xl[2]).c_if(sr_xl, 2)
    qc_xl.x(qr_xl[0]).c_if(sr_xl, 1)
    qc_xl.x(qr_xl[1]).c_if(sr_xl, 3)

    return qc_xl

logical_x = get_logical_x()

bits_out = sim.run(logical_x, shots=1, memory=True).result().get_memory()
print(bits_out)

logical_x.draw('mpl')

from qiskit.providers.aer.noise import NoiseModel
from qiskit.providers.aer.noise.errors import pauli_error, depolarizing_error
from qiskit import QuantumCircuit, Aer, transpile, assemble
from qiskit.visualization import plot_histogram

aer_sim = Aer.get_backend('aer_simulator')

def get_noise(p_gate):
    error_gate1 = depolarizing_error(p_gate, 1)

    noise_model = NoiseModel()
    noise_model.add_all_qubit_quantum_error(error_gate1, ["x"]) # single qubit gate error is applied to x gates

    return noise_model

noise_model = get_noise(0.01)

qr = QuantumRegister(3, 'qubit')
ar = QuantumRegister(2, 'auxiliary')
sr = ClassicalRegister(2, 'syndrome')
cr = ClassicalRegister(3, 'measure')
qc_000 = QuantumCircuit(qr, ar, sr, cr)
qc_000 = qc_000.compose(logical_x)
qc_000.barrier()
qc_000.measure(qr, cr)

qc_000.draw('mpl')

# run the circuit with the noise model and extract the counts
qobj = assemble(qc_000)
counts = aer_sim.run(qobj, noise_model=noise_model).result().get_counts()
plot_histogram(counts)

qc_3qz = QuantumCircuit(3)

### your code goes here. ###





########

qc_3qz.draw('mpl')

err, which_qubit = apply_err(3, 'phase')
qc_3qz.append(err, range(3))

### your code goes here ###











##########
qc_3qz.draw('mpl')


qc_3qz_trans = transpile(qc_3qz, sim)
syndrome = sim.run(qc_3qz_trans, shots=1, memory=True).result().get_memory()
print(syndrome)

your_answer = input('Enter the index of the code qubit that underwent phase-flip error: ')
print('\n')
print(which_qubit == int(your_answer))

## Your answer goes here
t = int(input('The number of counting qubit: '))

qc0 = QuantumCircuit(t+3, 1)

qc0.h(-1)
qc0.barrier()

## your code goes here ##





######

qc0.measure(0, 0)
qc0.draw('mpl')

counts_qc0 = sim.run(qc0, shots=8192).result().get_counts()
plot_histogram(counts_qc0)

qc1 = QuantumCircuit(t+3, 1)

qc1.h(-1)
qc1.barrier()

## your code goes here ##




######

qc1.measure(0, 0)
qc1.draw('mpl')

counts_qc1 = sim.run(qc1, shots=8192).result().get_counts()
plot_histogram(counts_qc1)