import qiskit
from qiskit import QuantumCircuit
from qiskit.visualization import plot_histogram

import math
import numpy as np

from qiskit import execute, Aer

simulator = Aer.get_backend('qasm_simulator')

import matplotlib.pyplot as plt

import warnings
warnings.filterwarnings('ignore')

state_0 = np.array([1, 0])
state_1 = np.array([0, 1])

from functools import reduce

Dag = lambda matrix: matrix.conj().T
Kron = lambda *matrices: reduce(np.kron, matrices)

def getMeasurements(n):
    psi_0 = np.array([1.0, 0.0])
    psi_1 = np.array([0.0, 1.0])
    I = np.eye(2)

    M_0 = psi_0.reshape([2, 1]) @ psi_0.reshape([1, 2]).conj()
    M_1 = psi_1.reshape([2, 1]) @ psi_1.reshape([1, 2]).conj()
    
    M = [M_0, M_1]
    
    measurements = []
    for i in range(2 ** n):
        binnum = bin(i)[2:].rjust(n, '0')

        temp = []
        indices = map(lambda x: int(x), list(binnum))
        for index in indices:
            temp.append(M[index])
        
        measurements.append( Kron(*temp) )
    
    return measurements

def measure(M, state):
    return Dag(state) @ Dag(M) @ M @ state

initState = Kron(
    state_0, state_0, state_0, state_0, state_0, 
    state_0, state_0, state_0, state_0
)

H = np.array([[1,1], [1,-1]]) / (np.sqrt(2))
H_8 = Kron(H, H, H, H, H, H, H, H)
H_9 = Kron(H, H, H, H, H, H, H, H, H)

N = 2 ** 9

def createTargetState(n):
    states = []
    s = bin(n)[2:].rjust(9, "0")
    for i in list(s):
        if i == '0':
            states.append(state_0)
        else:
            states.append(state_1)
    return Kron(*states)

def createTargetStates(p):
    n = math.floor(p * N)
    
    state = createTargetState(0)
    for i in range(1, n):
        state = state + createTargetState(i)
        
    return state / (np.sqrt(n))

target_states = [
    [createTargetStates(0.05), 0.05], # 0.05
    [createTargetStates(0.1), 0.1], # 0.1
    [createTargetStates(0.15), 0.15], # 0.15
    [createTargetStates(0.2), 0.2], # 0.2
    [createTargetStates(0.25), 0.25], # 0.25
    [createTargetStates(0.3), 0.3], # 0.3
    [createTargetStates(0.35), 0.35], # 0.35
    [createTargetStates(0.4), 0.4], # 0.4
    [createTargetStates(0.45), 0.45], # 0.45
    [createTargetStates(0.5), 0.5], # 0.5
    [createTargetStates(0.55), 0.55], # 0.55
    [createTargetStates(0.6), 0.6], # 0.6
    [createTargetStates(0.65), 0.65], # 0.65
    [createTargetStates(0.7), 0.7], # 0.7
    [createTargetStates(0.75), 0.75], # 0.75
    [createTargetStates(0.8), 0.8], # 0.8
    [createTargetStates(0.85), 0.85], # 0.85
    [createTargetStates(0.9), 0.9], # 0.9
    [createTargetStates(0.95), 0.95], # 0.95
    [createTargetStates(1), 1] # 1
]

def createOriginalOracle(target_state):
    O = np.eye(N) - 2 * ( Dag(target_state).reshape(N, 1) @ target_state.reshape(1, N) )
    
    return O

def createOriginalDiffuser():
    R = 2 * ( 
        Dag(initState).reshape(N, 1) @ 
        initState.reshape(1, N) 
    ) - np.eye(N)
    
    return  H_9 @ R @ Dag(H_9)

def createHalfPIOracle(target_state):
    alpha = np.pi / 2
    alpha_phase = np.exp(alpha * 1j)

    # Oracle = I - (1 - e^(αi))|target><target|

    O = np.eye(N) - (1 - alpha_phase) * \
        ( Dag(target_state).reshape(N, 1) @ target_state.reshape(1, N) )
    
    return O

def createHalfPIDiffuser():
    beta = -np.pi / 2
    beta_phase = np.exp(beta * 1j)

    # R = (1 - e^(βi))|0><0| + e^(βi)I
    R = (1 - beta_phase) * ( 
        Dag(initState).reshape(N, 1) @ 
        initState.reshape(1, N) 
    ) + beta_phase * np.eye(N)
    
    return  H_9 @ R @ Dag(H_9)

def createOneTenthPIOracle(target_state):
    alpha = np.pi / 10
    alpha_phase = np.exp(alpha * 1j)

    # Oracle = I - (1 - e^(αi))|target><target|

    O = np.eye(N) - (1 - alpha_phase) * \
        ( Dag(target_state).reshape(N, 1) @ target_state.reshape(1, N) )
    
    return O

def createOneTenthPIDiffuser():
    beta = -np.pi / 10
    beta_phase = np.exp(beta * 1j)

    # R = (1 - e^(βi))|0><0| + e^(βi)I
    R = (1 - beta_phase) * ( 
        Dag(initState).reshape(N, 1) @ 
        initState.reshape(1, N) 
    ) + beta_phase * np.eye(N)
    
    return  H_9 @ R @ Dag(H_9)

def createYouneOracle(n):
    O = np.eye(N)
    for i in range(n):
        O[i][i] = 0
        O[i][i + N // 2] = 1
        O[i + N // 2][i + N // 2] = 0
        O[i + N // 2][i] = 1
    
    return O

def createYouneDiffuser():
    R = (2 * Dag(initState).reshape(N, 1) @ 
        initState.reshape(1, N) - np.eye(N))
    
    return  Kron(np.eye(2), H_8) @ R @ Dag(Kron(np.eye(2), H_8))

measurements = getMeasurements(9)

# test: one grover iteration
x_o = []
y_o = []
for target_state, p in target_states:
    n = math.floor(p * N)

    O = createOriginalOracle(target_state)
    diff = createOriginalDiffuser()
    final_state = diff @ O @ H_9 @ initState

    probability = 0.0
    for i in range(N):
        if i < n:
            probability = probability + measure(measurements[i], final_state)
    
#     print(f'p={p}, probability={probability}')
    x_o.append(p)
    y_o.append(probability)

plt.title("Original Grover - apply one iterator")
plt.scatter(x_o, y_o, color="#66CCCC")

# test: best probability
x_o_best = []
y_o_best = []
for target_state, p in target_states:
    n = math.floor(p * N)

    O = createOriginalOracle(target_state)
    diff = createOriginalDiffuser()
    final_state = H_9 @ initState

    theta = math.degrees( math.asin(math.sqrt(p)) )
    k_tilde = math.floor( 180 / (4 * theta) - 0.5 )
    for i in range(k_tilde):
        final_state = diff @ O @ final_state

    probability = 0.0
    for i in range(N):
        if i < n:
            probability = probability + measure(measurements[i], final_state)
    
    final_state = diff @ O @ final_state
    probability2 = 0.0
    for i in range(N):
        if i < n:
            probability2 = probability2 + measure(measurements[i], final_state)
            
#     print(f'p={p}, probability={probability}')
    x_o_best.append(p)
    y_o_best.append(max(probability, probability2))

plt.title("Original Grover - best Probability")
x = np.arange(-0,1,0.01)
y = 1 - x
plt.plot(x, y, ls='--', color="#FFCC00")

plt.ylim(0, 1.05)
plt.scatter(x_o_best, y_o_best, color="#66CCCC")

# test: one grover iteration
x_half_pi = []
y_half_pi= []
for target_state, p in target_states:
    n = math.floor(p * N)

    O = createHalfPIOracle(target_state)
    diff = createHalfPIDiffuser()
    final_state = diff @ O @ H_9 @ initState

    probability = 0.0
    for i in range(N):
        if i < n:
            probability = probability + measure(measurements[i], final_state)
    
#     print(f'p={p}, probability={probability}')
    x_half_pi.append(p)
    y_half_pi.append(probability)

plt.title("pi/2 Grover - apply one iterator")
plt.scatter(x_half_pi, y_half_pi, color="#66CCCC")

# one grover iteration - pi/2 vs origin

fig, ax = plt.subplots()
plt.title("Apply one iteration")
plt.axvline(1/3, linestyle='--', color="#FFCCCC", label="t/N = 1/3")
plt.axhline(0.925, linestyle='--', color="#FFCC00", label="P=0.925")
plt.scatter(x_o, y_o, color="#66CCCC", label="Original Grover")
plt.scatter(x_half_pi, y_half_pi, color="#FF6666", label="pi/2 phase based Grover")
ax.legend()

# test: one grover iteration
x_ot_pi = []
y_ot_pi= []
for target_state, p in target_states:
    n = math.floor(p * N)

    O = createOneTenthPIOracle(target_state)
    diff = createOneTenthPIDiffuser()
    final_state = diff @ O @ H_9 @ initState

    probability = 0.0
    for i in range(N):
        if i < n:
            probability = probability + measure(measurements[i], final_state)
    
#     print(f'p={p}, probability={probability}')
    x_ot_pi.append(p)
    y_ot_pi.append(probability)

plt.title("0.1pi Grover - apply one iteration")
plt.scatter(x_ot_pi, y_ot_pi, color="#66CCCC")

# test: best probability
x_ot_best = []
y_ot_best = []
diff = createOneTenthPIDiffuser()
for target_state, p in target_states:
    n = math.floor(p * N)

    O = createOneTenthPIOracle(target_state)
#     diff = createOneTenthPIDiffuser()
    final_state = H_9 @ initState

    p_best = 0.0
    k_tilde = math.floor( 5 * np.sqrt(N) )
    for i in range(k_tilde):
        final_state = diff @ O @ final_state

        probability = 0.0
        for i in range(N):
            if i < n:
                probability = probability + measure(measurements[i], final_state)
                
        p_best = max(p_best, probability)
    
#     print(f'p={p}, probability={probability}')
    x_ot_best.append(p)
    y_ot_best.append(p_best)

plt.title("0.1pi Grover - best Probability")
plt.ylim(0, 1)
plt.scatter(x_ot_best, y_ot_best, color="#66CCCC")

# best probability - 0.1pi vs origin

fig, ax = plt.subplots()
plt.title("Best Probability")
x = np.arange(-0,1,0.01)
y = 1 - x
plt.plot(x, y, ls='--', color="#FFCC00")

plt.ylim(0, 1.05)
plt.scatter(x_o_best, y_o_best, color="#66CCCC", label="Original Grover")
plt.scatter(x_ot_best, y_ot_best, color="#FF6666", label="0.1pi phase based Grover")
ax.legend()

# test: best probability
x_y_best = []
y_y_best = []
for _, p in target_states:
    n = math.floor(p * (N // 2))

    O = createYouneOracle(n)
    diff = createYouneDiffuser()
    final_state = Kron( np.eye(2), H_8) @ initState
    
    p_best = 0.0
    k_tilde = math.floor(np.pi / (2 * np.sqrt(2)) * np.sqrt((N // 2) / n))
    for i in range(k_tilde):
        final_state = diff @ O @ final_state
        
    probability = 0.0
    for i in range(N // 2):
        if i < n:
            probability = probability + measure(measurements[i], final_state) + measure(measurements[i + N // 2], final_state)

    #     print(f'p={p}, probability={probability}')
    x_y_best.append(p)
    y_y_best.append(probability)

fig, ax = plt.subplots()
plt.title("Best Probability")

x = np.arange(-0, 1, 0.01)
y = 1 - x
plt.plot(x, y, ls='--', color="#FFCC00")

plt.axhline(0.8472, linestyle='--', color="#99CCFF", label="P=0.8472")
plt.axvline(0.308, linestyle='--', color="#FFCCCC", label="P=0.308")

plt.ylim(0, 1.05)
plt.scatter(x_y_best, y_y_best, color="#FF6666", label="Youne")
plt.scatter(x_o_best, y_o_best, color="#66CCCC", label="Original Grover")
ax.legend()

