import sys
# !conda install --yes --prefix {sys.prefix} numpy
# !conda install --yes --prefix {sys.prefix} qiskit
# !conda install --yes --prefix {sys.prefix} matplotlib
# !{sys.executable} -m pip install pyppeteer

import numpy as np
from qiskit import *
from matplotlib import pyplot
from math import pi

thetas = [i*pi/10 for i in range(0,11)]
shots = 100

import matplotlib.pyplot as plt
import numpy as np

def plot_results(counts_arr):

    zeros_count = [count[0] for count in counts_arr] #+ [sum([count[0] for count in counts_arr])/len(counts_arr)]
    ones_count = [count[1] for count in counts_arr] #+ [sum([count[1] for count in counts_arr])/len(counts_arr)]
    ones_pct = [count/shots for count in ones_count]

    fig, ax = plt.subplots()
    labels = [f'{i}/10*pi' for i in range(len(thetas))] #+ ['Average']
    x = np.arange(len(labels))
    width=0.35

#     c1 = ax.bar(x - width/2, zeros_count, width, label='Output = 0')
    c2 = ax.bar(x + width/2, ones_pct, width, label='Output = 1')

    ax.set_ylabel('Output 1 Average')
    ax.set_xlabel('Theta')
    ax.set_title('Output 1 Averages vs. Theta')
    ax.set_xticks(x)
    ax.set_xticklabels(labels)
    ax.legend()

#     ax.bar_label(c1, padding=3)
    ax.bar_label(c2, padding=3)

    fig.set_size_inches(12, 9)

    plt.show()

counts_arr = [ [] for i in range(len(thetas)) ]

for i, theta in enumerate(thetas):
    
    # Build circuit
    qc = QuantumCircuit(1,1)
    qc.h(0)
    qc.rz(theta,0)
    qc.h(0)
    qc.measure(0,0)
    if i==0:
        qc.draw('mpl')
        pyplot.show()

    # Simulate circuit using Aer's qasm_simulator
    backend_sim = Aer.get_backend('qasm_simulator')
    job_sim = backend_sim.run(transpile(qc, backend_sim), shots=shots)
    result_sim = job_sim.result()
    counts = result_sim.get_counts(qc)
    
    # collect counts of simulation output
    counts_arr[i].append( counts['0'] if '0' in counts else 0 )
    counts_arr[i].append( counts['1'] if '1' in counts else 0 )
    
    print(f'theta={i}/10*pi\tcounts: {counts}\tavg: {counts_arr[i][1]/shots*100:.1f}%')
    
plot_results(counts_arr)

counts_arr = [ [] for i in range(len(thetas)) ]
for i, theta in enumerate(thetas):
    
    # Build circuit
    qc = QuantumCircuit(1,1)
    qc.h(0)
    qc.rz(theta/2,0)
    qc.x(0)
    qc.rz(theta/2,0)
    qc.x(0)
    qc.h(0)
    qc.measure(0,0)
    if i==0:
        qc.draw('mpl')
        pyplot.show()

    # Simulate circuit using Aer's qasm_simulator
    backend_sim = Aer.get_backend('qasm_simulator')
    job_sim = backend_sim.run(transpile(qc, backend_sim), shots=shots)
    result_sim = job_sim.result()
    counts = result_sim.get_counts(qc)
    
    # collect counts of simulation output
    counts_arr[i].append( counts['0'] if '0' in counts else 0 )
    counts_arr[i].append( counts['1'] if '1' in counts else 0 )
    
    print(f'theta={i}/10*pi\tcounts: {counts}\tavg: {counts_arr[i][1]/shots:.2f}%')
    
plot_results(counts_arr)


deltas = [0.1, 0.2]

for delta in deltas:
    counts_arr = [ [] for i in range(len(thetas)) ]
    print(f"\n-----\nDelta: {delta}")
    for i, theta in enumerate(thetas):

        # Build circuit
        qc = QuantumCircuit(1,1)
        qc.h(0)
        qc.rz(theta/2+delta,0)
        qc.x(0)
        qc.rz(theta/2-delta,0)
        qc.x(0)
        qc.h(0)
        qc.measure(0,0)
        if i==0:
            qc.draw('mpl')
            pyplot.show()

        # Simulate circuit using Aer's qasm_simulator
        backend_sim = Aer.get_backend('qasm_simulator')
        job_sim = backend_sim.run(transpile(qc, backend_sim), shots=shots)
        result_sim = job_sim.result()
        counts = result_sim.get_counts(qc)

        # collect counts of simulation output
        counts_arr[i].append( counts['0'] if '0' in counts else 0 )
        counts_arr[i].append( counts['1'] if '1' in counts else 0 )

        print(f'theta={i}/10*pi\tcounts: {counts}\tavg: {counts_arr[i][1]/shots:.2f}%')
    
    plot_results(counts_arr)

from qiskit.quantum_info import Operator
from qiskit.extensions import HamiltonianGate
from qiskit.opflow import I, X, Y, Z

counts_arr = [ [] for i in range(len(thetas)) ]
for i, theta in enumerate(thetas):    

    # Define operator in exponent
    op = X + Y + Z

    # Define coefficient in exponent
    c = (theta/4)/2*((1/3)**(1/2))

    # Create HamiltonianGate, which is defined as exp(-i*op*c)
    V_gate = HamiltonianGate(op, c)

    # Define circuit
    qc = QuantumCircuit(1,1)
    qc.append(V_gate, [0])
    qc.x(0)
    qc.append(V_gate, [0])
    qc.x(0)
    qc.append(V_gate, [0])
    qc.x(0)
    qc.append(V_gate, [0])
    qc.x(0)
    qc.measure(0,0)

    # Draw circuit once
    if i==0:
        qc.draw('mpl')
        pyplot.show()

    # Simulate circuit using Aer's qasm_simulator
    backend_sim = Aer.get_backend('qasm_simulator')
    job_sim = backend_sim.run(transpile(qc, backend_sim), shots=shots)
    result_sim = job_sim.result()
    counts = result_sim.get_counts(qc)
    
    # collect counts of simulation output
    counts_arr[i].append( counts['0'] if '0' in counts else 0 )
    counts_arr[i].append( counts['1'] if '1' in counts else 0 )
    
    print(f'theta={i}/10*pi\tcounts: {counts}\tavg: {counts_arr[i][1]/shots:.2f}%')
    
plot_results(counts_arr)

counts_arr = [ [] for i in range(len(thetas)) ]
for i, theta in enumerate(thetas):    

    # Define operator in exponent
    op = X + Y + Z

    # Define coefficient in exponent
    c = (theta/4)/2*((1/3)**(1/2))

    # Create HamiltonianGate, which is defined as exp(-i*op*c)
    V_gate = HamiltonianGate(op, c)

    # Define circuit
    qc = QuantumCircuit(1,1)
    qc.append(V_gate, [0])
    qc.x(0)
    qc.append(V_gate, [0])
    qc.y(0)
    qc.append(V_gate, [0])
    qc.x(0)
    qc.append(V_gate, [0])
    qc.y(0)
    qc.measure(0,0)

    # Draw circuit once
    if i==0:
        qc.draw('mpl')
        pyplot.show()

    # Simulate circuit using Aer's qasm_simulator
    backend_sim = Aer.get_backend('qasm_simulator')
    job_sim = backend_sim.run(transpile(qc, backend_sim), shots=shots)
    result_sim = job_sim.result()
    counts = result_sim.get_counts(qc)
    
    # collect counts of simulation output
    counts_arr[i].append( counts['0'] if '0' in counts else 0 )
    counts_arr[i].append( counts['1'] if '1' in counts else 0 )
    
    print(f'theta={i}/10*pi\tcounts: {counts}\tavg: {counts_arr[i][1]/shots:.2f}%')
    
plot_results(counts_arr)