from qiskit import QuantumCircuit, transpile, QuantumRegister, ClassicalRegister, AncillaRegister
from qiskit_aer import AerSimulator
from qiskit.circuit.library import MCMT, GroverOperator
from qiskit.quantum_info import Statevector
from math import pi
from qiskit.visualization import plot_histogram, plot_bloch_multivector, array_to_latex
from qiskit.circuit import Parameter
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from utils import get_image_path
import seaborn as sns
from qutip import Bloch, Qobj
from itertools import combinations, product

def get_diffuser(total_qubits):
    mct = MCMT('x', total_qubits-1, 1)
    qc = QuantumCircuit(total_qubits)
    qc.h(range(total_qubits))
    qc.x(range(total_qubits))

    qc.h(-1)
    qc.compose(mct, range(total_qubits), inplace=True)
    qc.h(-1)
    
    qc.x(range(total_qubits))
    qc.h(range(total_qubits))
    return qc

def classical_grover(oracle):
    total_qubits = oracle.num_qubits
    qc = QuantumCircuit(total_qubits)
    grover = GroverOperator(oracle).decompose()
    qc.h(range(total_qubits))
    qc.compose(grover, range(total_qubits), inplace=True)
    return qc

def phase_oracle(states):
    total_qubits = len(states[0])
    mct = MCMT('z', total_qubits-1, 1)
    
    qc = QuantumCircuit(total_qubits, name="oracle")
    for state in states:
        zeros = [i for i, bit in enumerate(state[::-1]) if bit == '0']
        if(zeros): qc.x(zeros)
        qc.compose(mct,range(total_qubits),inplace=True)
        if(zeros): qc.x(zeros)
    return qc

oracle = phase_oracle(['100'])
oracle.draw('mpl')

classical_grover(oracle).draw('mpl')

def get_counts(qc, sim=AerSimulator(), shots=1000):
    qc_c = qc.copy()
    qc_c.measure_all()
    transpiled = transpile(qc_c, sim)
    return sim.run(transpiled, shots=shots).result().get_counts()

def new_grover(oracle):
    theta = Parameter('theta')
    total_qubits = oracle.num_qubits
    
    qc = QuantumCircuit(total_qubits)
    qc.ry(theta, range(total_qubits))
    qc.barrier()
    qc.append(oracle, range(total_qubits))
    qc.barrier()
    qc.compose(get_diffuser(total_qubits), range(total_qubits), inplace=True)
    
    return qc, theta

n_grover,theta = new_grover(oracle) 
n_grover.draw('mpl')

def find_best_angle(theta, qc, values_space, target_bit_string, debug=False):
    best_angle = 0
    best_total = 0
    
    for angle in values_space:
        bc = qc.assign_parameters({theta:angle})
        counts = get_counts(bc)

        value = 0
        if(counts.get(target_bit_string)):
            value = counts[target_bit_string]

        if(value > best_total):
            best_total = value
            best_angle = angle
        
        if(debug):
            print(angle, value)

    return best_angle, best_total

best_angle, best_total = find_best_angle(theta, n_grover, np.linspace(0, pi, num=1000), '001', debug=True)

best_angle, best_total

BIT_STRINGS_3 = ['000', '001', '010', '011', '100', '101', '110', '111']

def compare_old_new_1_bit_string(bit_strings, values_space, debug=False):
    old = {}
    new = {}
    angles = {}
    
    for bit_string in bit_strings:
        oracle = phase_oracle([bit_string])

        old_ = classical_grover(oracle.copy())
        if(debug):
            display(old_.draw('mpl'))
        
        old_counts = get_counts(old_)
        old_total = 0
        if(old_counts.get(bit_string)):
            old[bit_string] = old_counts[bit_string]

        new_, theta = new_grover(oracle.copy())
        angle, max_ = find_best_angle(theta, new_, values_space, bit_string)
        new[bit_string] = max_
        angles[bit_string] = angle

    data = {"old":old, "new":new, "angles":angles}
    return pd.DataFrame(data, columns=("old", "new", "angles"), index=bit_strings)
        
        
        

results_1 = compare_old_new_1_bit_string(BIT_STRINGS_3, np.linspace(0, pi, num=500))
results_1

plt.plot(results_1.index.tolist(), results_1['new'], label="new")
plt.plot(results_1.index.tolist(), results_1['old'], label="old")
plt.title('old vs new Grover improvement (range [0, $\pi$])')
plt.xlabel("bit string")
plt.ylabel("counts")
plt.legend(loc="upper right")

plt.savefig(get_image_path("superposition-improvement-grover-algorithm-range-0-to-pi.png"))
plt.show()

results_2 = compare_old_new_1_bit_string(BIT_STRINGS_3, np.linspace(0, 2*pi, num=500))
results_2

plt.plot(results_2.index.tolist(), results_2['new'], label="new")
plt.plot(results_2.index.tolist(), results_2['old'], label="old")
plt.title('old vs new Grover improvement (range [0, $2\pi$])')
plt.xlabel("bit string")
plt.ylabel("counts")
plt.legend(loc="upper right")

plt.savefig(get_image_path("superposition-improvement-grover-algorithm-range-0-to-2pi.png"))
plt.show()

plt.plot(BIT_STRINGS_3, results_1['angles'], label="first")
plt.plot(BIT_STRINGS_3, results_2['angles'], label="second")
plt.title('angles comparation ([0, $\pi$] vs [0, $2\pi$])')
plt.xlabel("bit string")
plt.ylabel("angle")
plt.legend(loc="upper right")
plt.savefig(get_image_path("improvement-grover-algorithm-angles-comparation-0-pi-and-2pi.png"))
plt.show()

results_3 = compare_old_new_1_bit_string([bin(i)[2:].zfill(4) for i in range(2**4)], np.linspace(0, 2*pi, num=500))
results_3

plt.plot(results_3.index.tolist(), results_3['new'], label="new")
plt.plot(results_3.index.tolist(), results_3['old'], label="old")
plt.title('old vs new Grover improvement (range [0, $2\pi$])')
plt.xlabel("bit string")
plt.ylabel("counts")
plt.xticks(rotation = 45)
plt.legend(loc="upper right")
plt.savefig(get_image_path("improvement-grover-algorithm-4bits-0-to-2pi.png"))
plt.show()

pi/2

mid_values = results_3["angles"].drop(["0000", "1111"])
mid_values

mid_values.describe()

mid_values.mean(), pi/2

bloch_data = results_3
bloch_data['alfa'] = np.cos(bloch_data["angles"]/2)
bloch_data['beta'] = np.sin(bloch_data["angles"]/2)
bloch_data

def show_bloch(states):
    bloch = Bloch()
    bloch.make_sphere()
    bloch.add_states(states)
    display(bloch.show())

for i in bloch_data.index.tolist():
    print(i, i.count('0'))

states = []
for i in range(len(bloch_data)):
    row = bloch_data.iloc[i]
    alfa,beta = row['alfa'], row['beta']
    states.append(Qobj([[alfa], [beta]]))
show_bloch(states)

states = []
for i in range(len(bloch_data)):
    row = bloch_data.iloc[i]
    
    if(row.name.count('0') >= 3):
        alfa,beta = row['alfa'], row['beta']
        states.append(Qobj([[alfa], [beta]]))
show_bloch(states)

states = []
for i in range(len(bloch_data)):
    row = bloch_data.iloc[i]
    
    if(row.name.count('0') == 2):
        alfa,beta = row['alfa'], row['beta']
        states.append(Qobj([[alfa], [beta]]))
show_bloch(states)

states = []
for i in range(len(bloch_data)):
    row = bloch_data.iloc[i]
    
    if(row.name.count('0') <= 1):
        alfa,beta = row['alfa'], row['beta']
        states.append(Qobj([[alfa], [beta]]))
show_bloch(states)

def show_unitary(qc, sim=AerSimulator()):
    qc_c = qc.copy()
    qc_c.save_unitary()
    transpiled = transpile(qc_c, sim)
    u = sim.run(transpiled, shots=1000).result().get_unitary()
    display(array_to_latex(u, max_size=100000))

d = pd.DataFrame(columns=("a", "b", "c"))
d["a"] = 

BIT_STRINGS_2 = ['00', '01', '10', '11']

def test_angle_with_multiple_bit_strings(bit_strings):
    df = pd.DataFrame(columns=("angle", *bit_strings, "default"), index=bit_strings)
    
    for string in bit_strings:
        oracle = phase_oracle([string])
        n_grover,theta = new_grover(oracle.copy()) 
        best_angle, best_total = find_best_angle(theta, n_grover, np.linspace(0, pi, num=1000), string, debug=False)
        df.loc[string]['angle'] = best_angle
        df.loc[string][string] = best_total

        default = classical_grover(oracle.copy())
        default_counts = get_counts(default)
  
        default_total = 0
        if(default_counts.get(string)):
            default_total = default_counts[string]
        df.loc[string]['default'] = default_total
    
        for string2 in set(bit_strings)-{string}:
            oracle = phase_oracle([string2])
            n_grover_2,theta_2 = new_grover(oracle.copy())
            bc = n_grover_2.assign_parameters({theta_2: best_angle})
            counts = get_counts(bc)
    
            total = 0
            if(counts.get(string2)):
                total = counts[string2]

            df.loc[string2][string] = total
    
    return df

df = test_angle_with_multiple_bit_strings(BIT_STRINGS_2)
df

fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2)

ax1.plot(BIT_STRINGS_2, df['00'])
ax2.plot(BIT_STRINGS_2, df['01'])
ax3.plot(BIT_STRINGS_2, df['10'])
ax4.plot(BIT_STRINGS_2, df['11'])
plt.show()

df = test_angle_with_multiple_bit_strings(BIT_STRINGS_3)
df

plt.figure(figsize=(13, 7))

sns.set_palette("Blues")
for string in BIT_STRINGS_3:
    sns.lineplot(data=df, x=BIT_STRINGS_3, y=string, label=f"{string} angle")

sns.lineplot(data=df, x=BIT_STRINGS_3, y="default", label="classical Grover", color="red")

plt.xlabel("bit strings")
plt.ylabel("counts")
plt.legend(loc="upper right")
plt.show()

df.describe()

df.std()

angles = df['angle'][1:-2]
angles = angles.reset_index()
angles = angles.drop(columns=['index'])
angles.loc[5] = pi/2
angles

angles.std()

data = pd.DataFrame(columns=("+", "=", "-"), index=BIT_STRINGS_3)
default_theta = pi/2
deviation = angles.std()['angle']

for index, t_value in (('-', default_theta-deviation), ('=', default_theta), ('+', default_theta+deviation)): 
    for bit_string in BIT_STRINGS_3:
        oracle = phase_oracle([bit_string])
        ng, theta = new_grover(oracle)
        bc = ng.assign_parameters({theta:t_value})
        counts = get_counts(bc)
    
        total = 0
        if(counts.get(bit_string)):
            total = counts[bit_string]
    
        data.loc[bit_string][index] = total

data

data.describe()

data.std()

sns.set_palette("flare")
sns.lineplot(data=data, x=BIT_STRINGS_3, y='-', label="-deviation")
sns.lineplot(data=data, x=BIT_STRINGS_3, y='=', label="=")
sns.lineplot(data=data, x=BIT_STRINGS_3, y='+', label="+deviation")

plt.xlabel("bit strings")
plt.ylabel("counts")
plt.legend(loc="upper right")
plt.show()

oracle = phase_oracle(['010', '110'])
oracle.draw('mpl')

ng, theta = new_grover(oracle)
ng.draw('mpl')

def find_best_angle_multiple_bs(theta, qc, values_space, target_bit_strings, debug=False):
    best_angle = 0
    best_total = 0
    
    for angle in values_space:
        bc = qc.assign_parameters({theta:angle})
        counts = get_counts(bc)

        value = 0
        for bit_string in target_bit_strings:
            if(counts.get(bit_string)):
                value += counts[bit_string]

        if(value > best_total):
            best_total = value
            best_angle = angle
        
        if(debug):
            print(angle, value)

    return best_angle, best_total

angle, max_ = find_best_angle_multiple_bs(theta, ng, np.linspace(0, pi, num=1000), ['010', '110'])
angle,max_

ng, theta = new_grover(oracle)
bc = ng.assign_parameters({theta:angle})
bc.draw('mpl')

plot_histogram(get_counts(bc))

c = classical_grover(oracle)
plot_histogram(get_counts(c))

df = pd.DataFrame(columns=("bs_1", "bs_2", "bs_3", "angle","max_new", "max_old"))

for i, comb in enumerate(combinations(BIT_STRINGS_3, 3)):
    oracle = phase_oracle(comb)
    
    old_ = classical_grover(oracle.copy())
    new_,theta = new_grover(oracle.copy())


    #get old
    old_counts = get_counts(old_)
    total_old = 0
    for c in comb:
        if(old_counts.get(c)):
            total_old += old_counts[c]


    #get new
    angle, max_ = find_best_angle_multiple_bs(theta, new_, np.linspace(0, pi, num=200), comb)


    df.loc[i] = (comb[0], comb[1], comb[2], angle, max_, total_old)


df

data = df
data["data"] = df["bs_1"] + " " + df["bs_2"] + " " + df["bs_3"]
data

plt.figure(figsize=(20,7))
sns.lineplot(data=data, x="data", y="max_old", label="grover")
sns.lineplot(data=data, x="data", y="max_new", label="new grover")

plt.xlabel("")
plt.ylabel("counts")
plt.xticks(rotation=90)
plt.legend(loc="upper right")
plt.savefig(get_image_path("multiple_strings_new_grover_best_angles.png"))
plt.show()

data['alfa'] = np.cos(data["angle"]/2)
data['beta'] = np.sin(data["angle"]/2)
data

states = []
for i in range(len(data)):
    row = data.loc[i]
    alfa,beta = row['alfa'], row['beta']
    states.append(Qobj([[alfa], [beta]]))
show_bloch(states)

plt.figure(figsize=(20, 7))
sns.lineplot(data=data, x="data", y="alfa")
sns.lineplot(data=data, x="data", y="beta")
plt.xticks(rotation=90)
plt.show()

data["t_zeros"] = data["data"].str.count('0')
data

more_zeros = data[data['t_zeros'] >= 5]
more_zeros

more_zeros[['alfa', 'beta', 't_zeros']].describe()

more_zeros['alfa'].mean()

#alfa = cos(theta/2)
#cos^-1(alfa) = theta/2
#(cos^-1(alfa))*2 = theta

theta = np.arccos(more_zeros['alfa'].mean())*2
theta, np.degrees(theta)

melted = pd.melt(more_zeros, id_vars=['t_zeros'], value_vars=['alfa', 'beta'], var_name="amplitude")
melted

sns.barplot(data=melted, x="t_zeros", y="value", hue="amplitude")
plt.legend(loc="lower right")
plt.show()

data_2 = data
data_2 = data_2[['t_zeros', 'alfa', 'beta']]
data_2

melted_entire = pd.melt(data_2, id_vars=['t_zeros'], value_vars=['alfa', 'beta'], var_name="amplitude")
melted_entire

sns.barplot(data=melted_entire, x="t_zeros", y="value", hue="amplitude")
plt.legend(loc="lower right")
plt.ylabel("amplitude")
plt.savefig(get_image_path("amplitudes_based_on_total_zeros_new_grover.png"))
plt.show()

df = compare_old_new_1_bit_string(BIT_STRINGS_3, np.linspace(0, pi, num=200))
df

df['alfa'] = np.cos(df['angles']/2)
df['beta'] = np.sin(df['angles']/2)
df

df['t_zeros'] = df.index.str.count('0')
df

melted_entire_2 = pd.melt(df, id_vars=['t_zeros'], value_vars=['alfa', 'beta'], var_name="amplitude")
melted_entire_2

sns.barplot(data=melted_entire_2, x="t_zeros", y="value", hue="amplitude")
plt.legend(loc="lower right")
plt.ylabel("amplitude")
plt.savefig(get_image_path("amplitudes_based_on_total_zeros_new_grover_one_bit_string.png"))
plt.show()

q = QuantumCircuit(2)
q.ry(1.957575,0)
q.ry(1.294526,1)
plot_bloch_multivector(q)

np.degrees((3*pi)/4)

def create_grover(strings):

    total_zeros = sum([ string.count('0') for string in strings ])
    n_qubits = len(strings[0])
    N = 2**n_qubits
    t = N/2
    
    if(n_qubits == 2):
        angle = pi/2

    if(total_zeros == n_qubits):
        angle = 0.963001
    elif(total_zeros == 0):
        angle = 2.099657
    else:
        angle = 1.657624
    
    
    
    oracle = phase_oracle(strings)
    ng, theta = new_grover(oracle)
    bc = ng.assign_parameters({theta:angle})
    return bc

g = create_grover(['111'])
plot_histogram(get_counts(g))

o = phase_oracle(['000'])
c = classical_grover(o)
plot_histogram(get_counts(c))

plot_histogram(get_counts(create_grover(['100','111', '000'])))

o = phase_oracle(['100','111', '000'])
c = classical_grover(o)
plot_histogram(get_counts(c))

def run_for_2(bit_strings):
    df = pd.DataFrame(columns=("bs_1", "bs_2", "angle","max_new", "max_old"))
    
    for i, comb in enumerate(combinations(bit_strings, 2)):
        oracle = phase_oracle(comb)
        
        old_ = classical_grover(oracle.copy())
        new_,theta = new_grover(oracle.copy())
    
    
        #get old
        old_counts = get_counts(old_)
        total_old = 0
        for c in comb:
            if(old_counts.get(c)):
                total_old += old_counts[c]
    
    
        #get new
        angle, max_ = find_best_angle_multiple_bs(theta, new_, np.linspace(0, pi, num=200), comb)
    
    
        df.loc[i] = (comb[0], comb[1], angle, max_, total_old)
    return df
df = run_for_2(BIT_STRINGS_3)
df.head()

# all angles get around pi/2
df['angle'].unique()

data = df
data["data"] = df["bs_1"] + " " + df["bs_2"]
data.head()

sns.lineplot(data=data, x="data", y="max_new")
sns.lineplot(data=data, x="data", y="max_old")

plt.xticks(rotation=90)
plt.xlabel("bit strings")
plt.ylabel("counts")
plt.savefig(get_image_path("2_bit_strings_l3_new_grover.png"))
plt.show()

plt.plot(df['angle'])
plt.show()

df = run_for_2(BIT_STRINGS_2)
df.head()

data = df
data["data"] = df["bs_1"] + " " + df["bs_2"]
data.head()

sns.lineplot(data=data, x="data", y="max_new", label="new")
sns.lineplot(data=data, x="data", y="max_old", label="old")

plt.xticks(rotation=90)
plt.xlabel("bit strings")
plt.ylabel("counts")
plt.legend(loc="upper right")
plt.savefig(get_image_path("2_bit_strings_l2_new_grover.png"))
plt.show()

plt.plot(df['angle'])
plt.show()

BIT_STRINGS_4 = [''.join(i) for i in product('01', repeat=4)]
BIT_STRINGS_4

df = run_for_2(BIT_STRINGS_4)
df.head()

data = df
data["data"] = df["bs_1"] + " " + df["bs_2"]
data.head()

plt.figure(figsize=(20, 7))
sns.lineplot(data=data, x="data", y="max_new", label="new")
sns.lineplot(data=data, x="data", y="max_old", label="old")

plt.xticks(rotation=90)
plt.xlabel("bit strings")
plt.ylabel("counts")
plt.legend(loc="upper right")
plt.savefig(get_image_path("2_bit_strings_l4_new_grover.png"))
plt.show()

def get_statistics_for_bit_strings(bit_strings_l):

    for bs in bit_strings_l:
        
        oracle = phase_oracle(bs)
        ng, theta = new_grover(oracle.copy())
        angle, max_ = find_best_angle_multiple_bs(theta, ng, np.linspace(0, pi, num=200), bs)


        total_bit_strings = len(bs)
        l = len(bs[0])
        
        bc = ng.assign_parameters({theta:angle})
        print(f"----{total_bit_strings} bit strings with length {l}----")
        display(bc.draw('mpl', filename=get_image_path(f"new_grover_rotation_for_{total_bit_strings}_bit_strings_l{l}.png")))
        plot_histogram(get_counts(bc), filename=get_image_path(f"new_grover_for_{total_bit_strings}_bit_strings_l{l}_outcomes.png"))
        display(plot_histogram(get_counts(bc)))
        
        c = classical_grover(oracle.copy())
        display(c.draw('mpl', filename=get_image_path(f"classical_grover_for_{total_bit_strings}_bit_strings_l{l}.png")))
        plot_histogram(get_counts(c), filename=get_image_path(f"classical_grover_for_{total_bit_strings}_bit_strings_l{l}_outcomes.png"))
        display(plot_histogram(get_counts(c)))
    

get_statistics_for_bit_strings([
    ['01'],
    ['01', '10'],
    ['01', '10', '00'],

    ['010'],
    ['010', '101'],
    ['010', '101', '110'],

    ['0101'],
    ['0101', '1101'],
    ['0101', '1101', '1001'],

    ['01011'],
    ['01011', '01101'],
    ['01011', '01101', '10111'],
])