%matplotlib inline
import hashlib
import numpy as np
import matplotlib.pyplot as plt
# Importing standard Qiskit libraries and configuring account
from qiskit import QuantumCircuit, execute
from qiskit.circuit import Parameter
from qiskit.providers.aer import QasmSimulator, StatevectorSimulator
from qiskit.visualization import *
from qiskit.quantum_info import *
from qiskit.circuit.library import HGate
success_msg = 'Your answer is correct and has been saved. Please continue to the next section.'
fail_msg = 'Your answer is not correct. Please try again.'

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

# Insert the necessary gates to change to the Hadamard basis below and measure

# Do not change below this line
qc1.draw('mpl')

basis_gates = ['id', 'x', 'y', 'z', 's', 't', 'sdg', 'tdg', 'h', 'p', 'sx' ,'r', 'rx', 'ry', 'rz', 'u', 'u1', 'u2', 'u3', 'barrier', 'measure']
assert list(qc1.count_ops()) != [], "Circuit cannot be empty"
assert set(qc1.count_ops().keys()).intersection(basis_gates) != set(), "Only basic gates are allowed"
job = execute(qc1, backend=QasmSimulator(), shots=1024, seed_simulator=0)
counts = job.result().get_counts()
sv_check = Statevector.from_instruction(qc1.remove_final_measurements(inplace=False)).evolve(HGate()).equiv(Statevector.from_label('r'))
op_check_dict = qc1.count_ops()
_ = op_check_dict.pop('measure', None)
_ = op_check_dict.pop('barrier', None)
op_check = len(op_check_dict) > 1
print(success_msg if (sv_check and op_check) else fail_msg)
answer1 = hashlib.sha256((str(counts)+str(sv_check and op_check)).encode()).hexdigest()
plot_histogram(counts)

from IPython.display import YouTubeVideo, display
polariser_exp = YouTubeVideo('6N3bJ7Uxpp0', end=93, height=405, width=720, modestbranding=1)
display(polariser_exp)

beta = Parameter('β')
qc2 = QuantumCircuit(1)
# Enter your code below this line

# Do not change below this line
qc2.draw(output='mpl')

def theoretical_prob(beta):
    '''
    Definition of theoretical transmission probability.
    The expression for transmitted probability between two polarisers 
    with a relative angle `beta` given in radians
       
    '''
    # Fill in the correct expression for this probability and assign it to the variable tp below
    # You may use numpy function like so: np.func_name()
    tp = 
    return tp

beta_range = np.linspace(0, np.pi, 50)

num_shots = 1024

basis_gates = ['id', 'x', 'y', 'z', 's', 't', 'sdg', 'tdg', 'h', 'p', 'sx' ,'r', 'rx', 'ry', 'rz', 'u', 'u1', 'u2', 'u3', 'barrier', 'measure']
assert set(qc2.count_ops().keys()).intersection(basis_gates) != set(), "Only basic gates are allowed"
job = execute(qc2,
              backend=QasmSimulator(),
              shots = num_shots,
              parameter_binds=[{beta: beta_val} for beta_val in beta_range], 
              seed_simulator=0) # For consistent results

counts = job.result().get_counts()
# Calculating the probability of photons passing through
probabilities = list(map(lambda c: c.get('0', 0)/num_shots, counts))

pol_checks = [Statevector.from_instruction(qc2.bind_parameters({beta: beta_val})
            .remove_final_measurements(inplace=False))
            .equiv(Statevector([np.cos(-beta_val), np.sin(-beta_val)]))
            for beta_val in beta_range]
print(success_msg if all(pol_checks) else fail_msg)
answer2 = hashlib.sha256((str(probabilities)+str(pol_checks)).encode()).hexdigest()

fig = plt.figure(figsize=(8,6))
ax = fig.add_subplot(111)

t_prob = np.vectorize(theoretical_prob)

ax.plot(beta_range, probabilities, 'o', label='Experimental')
ax.plot(beta_range, t_prob(beta_range), '-', label='Theoretical')

ax.set_xticks([i * np.pi / 4 for i in range(5)])
ax.set_xticklabels(['β', r'$\frac{\pi}{4}$', r'$\frac{\pi}{2}$', r'$\frac{3\pi}{4}$', r'$\pi$'], fontsize=14)
ax.set_xlabel('β', fontsize=14)
ax.set_ylabel('Probability of Transmission', fontsize=14)
ax.legend(fontsize=14)


print(f'Answer 1: {answer1}')
print(f'Answer 2: {answer2}')