from qiskit import QuantumCircuit, execute, Aer
from qiskit.opflow import One, Zero, Plus, Minus
from qiskit.opflow import I, X, Y, Z, S, H
from qiskit.visualization import plot_bloch_multivector, plot_bloch_vector, plot_histogram
from math import pi, sqrt, acos, asin, e, log, sin, cos
from numpy import angle

Zero

One

Plus

Minus

1.5 * Zero

-3 * Plus

Zero^Zero^One

Plus^Minus^Zero^Zero

Zero + One

Plus + Minus

(Zero^Zero^One) + (Zero^One^Zero) + (One^Zero^Zero)

(Zero^Zero^One) - 1.5*(Zero^One^Zero) + 3*(One^Zero^Zero)

def arbitrary_state(beta):
    ##### ==================================
    # Write your solution in here.
    alpha = sqrt(1 - abs(beta)**2) 
    return alpha*Zero + beta*One
    ##### ==================================

Zero.eval('0')

Zero.eval('1')

Minus.eval('0')

Minus.eval('1')

~Zero

~One

~Plus

print("< 0 | 0 > =",(~Zero @ Zero).eval())

print("< 1 | 0 > =",(~One @ Zero).eval())

beta = 0.232 + 0.341j 
arbitrary_state(beta)

beta = 0.232 + 0.341j 

print("< 1 | Ψ > =",(~One @ arbitrary_state(beta)).eval())

print("< 0 | Ψ > =",(~Zero @ arbitrary_state(beta)).eval())

print("< Ψ | Ψ > =", (~arbitrary_state(beta) @ arbitrary_state(beta)).eval())

def normality_check(beta):
    ##### ==================================
    # Write your solution in here.
    return abs((~One @ arbitrary_state(beta)).eval())**2 + abs((~Zero @ arbitrary_state(beta)).eval())**2
    ##### ==================================

# Play with your beta parameter to see if this value changes.
beta = 0.232 + 0.341j 
print("|< 1 | Ψ >|^2 + |< 0 | Ψ >|^2 =", normality_check(beta))

def convert_braket_to_Bloch(beta):
    if beta == 0: # special case
        return [1.0, 0, 0]
    ##### ==================================
    # Write your solution in here.
    alpha = sqrt(1 - abs(beta)**2) 
    theta = acos(alpha)*2
    beta = beta / sin(theta/2)
    phi = angle(beta)
    ##### ==================================
    return [1.0, theta, phi]

# Check 
beta = -1/sqrt(2)*1j
plot_bloch_vector(convert_braket_to_Bloch(beta), coord_type='spherical', 
                  title=r"$|\psi\rangle = \frac{1}{\sqrt{2}}|0\rangle-\frac{1}{\sqrt{2}}i|0\rangle$")

# Play with more values of beta to see if it matches with your expectation
beta = -1/sqrt(10)*1j
plot_bloch_vector(convert_braket_to_Bloch(beta), coord_type='spherical', 
                  title=r"$|\psi\rangle = \sqrt{1-|\beta|^2}|0\rangle+\beta|0\rangle$ in Spherical")

X

Y

Z

def expectation_Z(beta):
    ##### ==================================
    # Write your solution in here. There are more than one solution here...
    return (~arbitrary_state(beta) @ Z @ arbitrary_state(beta)).eval()
    ##### ==================================

print("<Z> =", expectation_Z(beta))

def expectation_X(beta):
    ##### ==================================
    # Write your solution in here. There are more than one solution here...
    return (~arbitrary_state(beta) @ ~H @ Z @ H @ arbitrary_state(beta)).eval()
    ##### ==================================
    
def expectation_Y(beta):
    ##### ==================================
    # Write your solution in here. There are more than one solution here...
    return (~arbitrary_state(beta) @ ~(H @ ~S) @ Z @ (H @ ~S) @ arbitrary_state(beta)).eval()
    ##### ==================================

def get_cartesian_coordinate(beta):
    ##### ==================================
    # Write your solution in here. There are more than one solution here...
    x = expectation_X(beta).real
    y = expectation_Y(beta).real
    z = expectation_Z(beta).real
    ##### ==================================
    return [x, y, z]

# Play with more values of beta to see if it matches with your previous Bloch sphere
beta = -1/sqrt(10)*1j
plot_bloch_vector(get_cartesian_coordinate(beta), coord_type='cartesian', 
                  title=r"$|\psi\rangle = \sqrt{1-|\beta|^2}|0\rangle+\beta|0\rangle$ in Cartesian")

# Let's define the quantum simulator

sv_simulator = Aer.get_backend("statevector_simulator")

qasm_simulator = Aer.get_backend("qasm_simulator")

this_circ = QuantumCircuit(1)
this_circ.initialize([1/sqrt(2), -1/sqrt(2)])
this_circ.measure_all()
this_circ.draw('mpl')

shots = 1024

counts = execute(this_circ, backend=qasm_simulator, shots=shots).result().get_counts()

print(counts)

plot_histogram(counts, title=r"Measuring the $|+\rangle$ state")

def contruct_arbitrary_state_on_circuit(beta):
    ##### ==================================
    # Write your solution in here. Here you do not need to append the measurement part in this function
    circuit = QuantumCircuit(1)
    circuit.initialize([sqrt(1-abs(beta)**2), beta])
    ##### ==================================
    return circuit

shots = 81920

beta = 0.213
this_circuit = contruct_arbitrary_state_on_circuit(beta)
this_circuit.measure_all()

counts = execute(this_circuit, backend=qasm_simulator, shots=shots).result().get_counts()

print("Probability in |1> basis =", beta**2)
plot_histogram(counts, title=r"Measuring the arbitrary $|\psi(\beta)\rangle$ state")


def expectation_Z_with_QASM(beta, shots):
    ##### ==================================
    # Write your solution in here. Here
    circuit = contruct_arbitrary_state_on_circuit(beta)
    circuit.measure_all()
    counts = execute(circuit, backend=qasm_simulator, shots=shots).result().get_counts()
    expectation = (counts["0"] - counts["1"])/(counts["0"] + counts["1"])
    ##### ==================================
    return expectation

beta = -1/sqrt(2)*1j
shots = 81920

print("Expectation <Z> with QASM =", expectation_Z_with_QASM(beta, shots=shots))
print("Expectation <Z> with opflow =", expectation_Z(beta))

import qiskit.tools.jupyter
%qiskit_version_table