%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
# Importing standard Qiskit libraries and configuring account
import qiskit
from qiskit import QuantumCircuit, execute
from qiskit.providers.aer import QasmSimulator, StatevectorSimulator
from qiskit.visualization import *
from qiskit.quantum_info import *
success_msg = 'Your answer is correct.'
fail_msg = 'Your answer is not correct. Please try again.'
basis_gates = ['id', 'x', 'y', 'z', 's', 't', 'sdg', 'tdg', 'h', 'p', 'sx' ,'r', 'rx', 'ry', 'rz', 'u', 'u1', 'u2', 'u3', 'cx', 'barrier', 'measure', 'snapshot']

qc1 = QuantumCircuit(3) 
# Insert gates below to create the state 

# Do not change below this line. You do not need to add an additional measurement. 
qc1.measure_all()
qc1.draw('mpl')

try:
    assert list(qc1.count_ops()) != [], "Circuit cannot be empty"
    assert set(qc1.count_ops().keys()).difference(basis_gates) == set(), f"Only the following basic gates are allowed: {basis_gates}"
    assert all([type(gate[0]) == qiskit.circuit.measure.Measure for gate in qc1.data[-3:len(qc1.data)]]), "Measurement must be the last operation in a circuit."
    sv_check = Statevector.from_instruction(qc1.remove_final_measurements(inplace=False)).equiv((Statevector.from_label('000')+Statevector.from_label('111'))/np.sqrt(2))
    assert sv_check, "You did not prepare the correct state."
    job = execute(qc1, backend=QasmSimulator(), shots=1024, seed_simulator=0)
    counts = job.result().get_counts()
    print(success_msg if (sv_check) else fail_msg)  
except AssertionError as e:
    print(f'Your code has an error:  {e.args[0]}')
    
except Exception as e:
    print(f'This error occured: {e.args[0]}')
plot_histogram(counts)

def superdense_coding_circ(bitstr):
    '''
    bitstr is a 2-character string with possible values '00', '01', '10', '11'
    '''
    qc2 = QuantumCircuit(2)
    # Enter your code below this line
    # Prepare a Bell state B_00 below
    
    # Do not remove the following line
    qc2.snapshot('1')
    # Apply single-qubit gates only on q_0 to encode the bitstring

    # Do not remove the following line
    qc2.snapshot('2')
    # Apply gates to change to the Bell basis

    # Do not remove the following line   
    qc2.snapshot('3')
    # Do not change below this line. You do not need to add an additional measurement. 
    qc2.measure_all()
    return qc2

b00_sv = (Statevector.from_label('00')+Statevector.from_label('11'))/np.sqrt(2)
b01_sv = (Statevector.from_label('01')+Statevector.from_label('10'))/np.sqrt(2)
b10_sv = (Statevector.from_label('00')-Statevector.from_label('11'))/np.sqrt(2)
b11_sv = (Statevector.from_label('01')-Statevector.from_label('10'))/np.sqrt(2)
bell_states_sv = [b00_sv, b01_sv, b10_sv, b11_sv]
basis_states = ['00', '01', '10', '11']
bell_dict = dict(zip(basis_states, bell_states_sv))
enc_dec_mat = np.zeros((len(basis_states),len(basis_states)))
try:
    for state in basis_states:
        sv_checks = []
        qc = superdense_coding_circ(state)
        assert list(qc.count_ops()) != [], f" Bitstring {state}: Circuit cannot be empty"
        assert qc.count_ops()['measure'] <= 2, f" Bitstring {state}: Please remove extra measurements"
        assert qc.count_ops()['snapshot'] == 3, f" Bitstring {state}: Please do not add or remove snapshots"
        assert set(qc.count_ops().keys()).difference(basis_gates) == set(), f" Bitstring {state}: Only the following basic gates are allowed: {basis_gates}"
        assert all([type(gate[0]) == qiskit.circuit.measure.Measure for gate in qc.data[-2:len(qc.data)]]), f" Bitstring {state}: Measurement must be the last operation in a circuit."
        result = execute(qc.reverse_bits(),backend=QasmSimulator(), shots=1).result().get_counts().most_frequent()
        assert state == result, f" Bitstring {state}: Your encoding is not correct"
        enc_dec_mat[int(state,2), int(result,2)] = 1
        sv = execute(qc.reverse_bits(), backend=StatevectorSimulator()).result().data()['snapshots']['statevector']
        sv_checks.append(b00_sv.equiv(sv['1'][0]))
        sv_checks.append(bell_dict[state].equiv(sv['2'][0]))
        sv_checks.append(Statevector.from_label(state).equiv(sv['3'][0]))
    assert all(np.diag(enc_dec_mat) == np.ones(4)), "Your encoding is not correct"
    plt.matshow(enc_dec_mat, cmap='binary')
    plt.title('Superdense Coding Matrix')
    plt.xlabel("Alice's encoded bits", fontsize=14)
    plt.ylabel("Bob's decoded bits", fontsize=14)
    plt.xticks(ticks=range(4), labels=basis_states, fontsize=12)
    plt.yticks(ticks=range(4), labels=basis_states, fontsize=12)
    circ_check = all(sv_checks)
    assert circ_check, "Your circuit does not work for all bitstrings"

    print(success_msg if circ_check else fail_msg)
except AssertionError as e:
    print(f'Your code has an error:  {e.args[0]}')
    
except Exception as e:
    print(f'This error occured: {e.args[0]}')


