import numpy as np

from IPython.display import display, HTML

from docplex.mp.model import Model


from qiskit.aqua import aqua_globals, QuantumInstance
from qiskit.aqua.algorithms import QAOA, NumPyMinimumEigensolver
from qiskit.optimization.algorithms import MinimumEigenOptimizer, RecursiveMinimumEigenOptimizer
from qiskit import Aer, BasicAer
from qiskit.circuit.library import TwoLocal
from qiskit.algorithms import VQE
from qiskit.aqua.algorithms import NumPyMinimumEigensolver
from qiskit.algorithms.optimizers import COBYLA, L_BFGS_B, SPSA, SLSQP
from qiskit.opflow import Z, X, I, Gradient 
from qiskit.optimization.applications.ising import docplex
from qiskit.optimization.problems import QuadraticProgram
from qiskit.opflow.primitive_ops.pauli_sum_op import PauliSumOp

from pandas import DataFrame

from origin import Origin
from doublets import Doublet
from triplets import Triplet

C:\Users\david\.conda\envs\Qiskit\lib\site-packages\qiskit\optimization\__init__.py:92: DeprecationWarning: The package qiskit.optimization is deprecated. It was moved/refactored to qiskit_optimization (pip install qiskit-optimization). For more information see <https://github.com/Qiskit/qiskit-aqua/blob/master/README.md#migration-guide>
  warn_package('optimization', 'qiskit_optimization', 'qiskit-optimization')

data_set = np.load('simplified_experiment.npy', allow_pickle=True)

hits_per_plane = [len(k) for k in data_set]
for i, hits in enumerate(hits_per_plane):
    print(f'Number of hits in plane {i}: {hits}')


# constants
LIMIT_SIGMA_X = 3
LIMIT_SIGMA_Y = 3


# testing with smaller samples
TEST_SAMPLE_SIZE = 4

sub_dict_0 = dict(list(data_set[0].items())[:TEST_SAMPLE_SIZE])
sub_dict_1 = dict(list(data_set[1].items())[:TEST_SAMPLE_SIZE])
sub_dict_2 = dict(list(data_set[2].items())[:TEST_SAMPLE_SIZE])
sub_dict_3 = dict(list(data_set[3].items())[:TEST_SAMPLE_SIZE])


origin_object = Origin(data_set[0])


print(f'Sigma x limit set to: {LIMIT_SIGMA_X}')
print(f'Sigma x limit set to: {LIMIT_SIGMA_Y}')
print(f'Test sample size: {TEST_SAMPLE_SIZE}')

Number of hits in plane 0: 2807
Number of hits in plane 1: 2801
Number of hits in plane 2: 2793
Number of hits in plane 3: 2783
Sigma x limit set to: 3
Sigma x limit set to: 3
Test sample size: 4

# Doublets
doublet_list_01 = []
doublet_list_12 = []
doublet_list_23 = []


# fill plane 0 -> 1 doublets
for p0_key, p0_value in zip(sub_dict_0.keys(), sub_dict_0.values()):
    for p1_key, p1_value in zip(sub_dict_1.keys(), sub_dict_1.values()):
        doublet = Doublet(p0_key, p0_value, p1_key, p1_value)
        if doublet.is_valid_doublet(origin_object, LIMIT_SIGMA_X, LIMIT_SIGMA_Y):
            doublet_list_01.append(doublet)

# fill plane 1 -> 2 doublets
for p1_key, p1_value in zip(sub_dict_1.keys(), sub_dict_1.values()):
    for p2_key, p2_value in zip(sub_dict_2.keys(), sub_dict_2.values()):
        doublet = Doublet(p1_key, p1_value, p2_key, p2_value)
        if doublet.is_valid_doublet(origin_object, LIMIT_SIGMA_X, LIMIT_SIGMA_Y):
            doublet_list_12.append(doublet)

# fill plane 2 -> 3 doublets
for p2_key, p2_value in zip(sub_dict_2.keys(), sub_dict_2.values()):
    for p3_key, p3_value in zip(sub_dict_3.keys(), sub_dict_3.values()):
        doublet = Doublet(p2_key, p2_value, p3_key, p3_value)
        if doublet.is_valid_doublet(origin_object, LIMIT_SIGMA_X, LIMIT_SIGMA_Y):
            doublet_list_23.append(doublet)

print(f'Number of doublets created from plane 0 -> 1: {len(doublet_list_01)}')
print(f'Number of doublets created from plane 1 -> 2: {len(doublet_list_12)}')
print(f'Number of doublets created from plane 2 -> 3: {len(doublet_list_23)}')

count_01 = 0
count_12 = 0
count_23 = 0

for d in doublet_list_01:
    if d.is_correct_match():
        count_01 += 1
for d in doublet_list_12:
    if d.is_correct_match():
        count_12 += 1
for d in doublet_list_23:
    if d.is_correct_match():
        count_23 += 1


print('---------------------------------')
print(f'Fraction of correct doublets recognized for set limits: {np.round(count_01 / len(sub_dict_1), 3)}')
print(f'Fraction of correct doublets recognized for set limits: {np.round(count_12 / len(sub_dict_2), 3)}')
print(f'Fraction of correct doublets recognized for set limits: {np.round(count_23 / len(sub_dict_3), 3)}')

Number of doublets created from plane 0 -> 1: 16
Number of doublets created from plane 1 -> 2: 16
Number of doublets created from plane 2 -> 3: 13
---------------------------------
Fraction of correct doublets recognized for set limits: 1.0
Fraction of correct doublets recognized for set limits: 1.0
Fraction of correct doublets recognized for set limits: 1.0

# Triplets
triplet_list_012 = []
triplet_list_123 = []

for d1 in doublet_list_01:
    for d2 in doublet_list_12:
        if d1.get_p2_key() == d2.get_p1_key():
            triplet = Triplet(d1, d2)
            if triplet.is_valid_triplet(origin_object, 1/20 * LIMIT_SIGMA_X, 1/20 * LIMIT_SIGMA_Y):
                triplet_list_012.append(triplet)

for d2 in doublet_list_12:
    for d3 in doublet_list_23:
        if d2.get_p2_key() == d3.get_p1_key():
            triplet = Triplet(d2, d3)
            if triplet.is_valid_triplet(origin_object, 1/20 * LIMIT_SIGMA_X, 1/20 * LIMIT_SIGMA_Y):
                triplet_list_123.append(triplet)


print(f'Number of triplets created from plane 0 -> 1 -> 2: {len(triplet_list_012)}')
print(f'Number of triplets created from plane 1 -> 2 -> 3: {len(triplet_list_123)}')


count_012 = 0
count_123 = 0

correct_state = ''


for i, t in enumerate(triplet_list_012):
    if t.is_correct_match():
        count_012 += 1
        correct_state += '1'
    else:
        correct_state += '0'
for j, t in enumerate(triplet_list_123):
    if t.is_correct_match():
        count_123 += 1
        correct_state += '1'
    else:
        correct_state += '0'

print('Correct triplet list state: ', correct_state)

print(f'Fraction of correct triplets recognized for set limits: {np.round(count_012 / len(sub_dict_2), 3)}')
print(f'Fraction of correct triplets recognized for set limits: {np.round(count_123 / len(sub_dict_3), 3)}')

Number of triplets created from plane 0 -> 1 -> 2: 6
Number of triplets created from plane 1 -> 2 -> 3: 4
Correct triplet list state:  1010111111
Fraction of correct triplets recognized for set limits: 1.0
Fraction of correct triplets recognized for set limits: 1.0

VQE for Solving QUBO's

\begin{equation} \tag{QUBO: general formula} H(q) = \sum_i a_i q_i + \sum_{i<j} b_{ij}q_iq_j \end{equation}

The q's are in this case triplets of consecutive hits through the detector (e.g hit on plane 1, 2, 3 or 2, 3, 4) So the objective transforms to:

\begin{equation} \tag{QUBO: triplets} H(T) = \sum_i a_i T_i + \sum_{i<j} b_{ij}T_iT_j \end{equation}

How to calculate: Let $P_i$ = {1, 2, 3} and $P_j$ = {2, 3, 4} be the sets of planes where the $T_i$ and $T_j$ come from.

a_i = []
b_ij = []

for t_i in triplet_list_012:
    a_i.append(max(t_i.get_sigma_diff_at_origin(origin_object)))
    b_fixed_i_variable_j = []
    for t_j in triplet_list_123:
        b_fixed_i_variable_j.append(t_i.conflict_with_other_triplet(origin_object, t_j))
    b_ij.append(b_fixed_i_variable_j)

# build model with docplex

mdl = Model('QUBO')
x_i = [mdl.binary_var() for i in range(len(triplet_list_012))] # creating binary variables
x_j = [mdl.binary_var() for j in range(len(triplet_list_123))] # creating binary variables


# Create the Hamiltonian,
objective = mdl.sum(a_i[i] * x_i[i] for i in range(len(x_i))) + \
            mdl.sum(b_ij[i][j] * x_i[i] * x_j[j] for i in range(len(x_i)) for j in range(len(x_j)))


mdl.minimize(objective)
print(mdl.export_as_lp_string())

\ This file has been generated by DOcplex
\ ENCODING=ISO-8859-1
\Problem name: QUBO

Minimize
 obj: 0.000784435655 x1 + 0.114230452797 x2 + 0.006672050882 x3
      + 0.114758190618 x4 + 0.004781701073 x5 + 0.017807995856 x6 + [
      - 1.939587965365 x1*x7 + 0.628073662404 x2*x8 + 0.995824239618 x2*x9
      - 1.986696526418 x3*x8 + 0.956346572371 x4*x7 + 0.614610212867 x4*x8
      - 1.990365011654 x5*x9 - 1.924500008322 x6*x10 ]/2
Subject To

Bounds
 0 <= x1 <= 1
 0 <= x2 <= 1
 0 <= x3 <= 1
 0 <= x4 <= 1
 0 <= x5 <= 1
 0 <= x6 <= 1
 0 <= x7 <= 1
 0 <= x8 <= 1
 0 <= x9 <= 1
 0 <= x10 <= 1

Binaries
 x1 x2 x3 x4 x5 x6 x7 x8 x9 x10
End

qubo = QuadraticProgram()
qubo.from_docplex(mdl)

aqua_globals.random_seed = 10598
quantum_instance = QuantumInstance(BasicAer.get_backend('qasm_simulator'),
                                   seed_simulator=aqua_globals.random_seed,
                                   seed_transpiler=aqua_globals.random_seed,
                                   shots=8192)
qaoa_mes = QAOA(quantum_instance=quantum_instance,initial_point=[0., 1.])
exact_mes = NumPyMinimumEigensolver()

qaoa = MinimumEigenOptimizer(qaoa_mes)   # using QAOA
exact = MinimumEigenOptimizer(exact_mes)  # using the exact classical numpy minimum eigen solver

exact_result = exact.solve(qubo)
print(exact_result)

optimal function value: -3.8905285724123955
optimal value: [1. 0. 1. 0. 1. 1. 1. 1. 1. 1.]
status: SUCCESS

qaoa_result = qaoa.solve(qubo)
print(qaoa_result)

optimal function value: -3.8905285724123955
optimal value: [1. 0. 1. 0. 1. 1. 1. 1. 1. 1.]
status: SUCCESS