from functools import reduce
from itertools import product

from qiskit import BasicAer, QuantumRegister
from qiskit_aqua import QuantumInstance
from qiskit_aqua import Operator, run_algorithm
from qiskit.quantum_info import Pauli
from qiskit_aqua.components.optimizers import COBYLA

from constrainedqaoa import constrainedQAOA

import numpy as np

edges = [(0, 0)]
vertices = 1
colors   = 2
n_qubits = vertices * colors

zr = np.zeros(n_qubits)
ws = np.eye(n_qubits)

cost_operator = reduce(
            lambda x, y: x + y,
            [
                Operator([[1, (Pauli(ws[colors*v1 + j, :], zr)
                               *Pauli(ws[colors*v2 + j, :], zr))]])
                for (v1, v2), j in product(edges, range(colors))
            ]
)

mixer_operator = reduce(
            lambda x, y: x + y,
            [
                Operator([[1, (Pauli(zr, ws[colors*i + j, :])
                               *Pauli(zr, ws[colors*i + (j+1) % colors, :]))]]) +
                Operator([[1, (Pauli(ws[colors*i + j % colors, :], ws[colors*i + j % colors, :])
                               *Pauli(ws[colors*i + (j+1) % colors, :], ws[colors*i + (j+1) % colors, :]))]])
                for i, j in product(range(vertices), range(colors))
            ]
)
if colors == 2:
    mixer_operator.scaling_coeff(1/2)

print(mixer_operator.print_operators())

what = 'circuit'
circ = mixer_operator.evolve(np.eye(4), np.pi*0.35, what, 1, QuantumRegister(2))
if what == 'matrix':
    print(np.round(circ, 4))
elif what == 'circuit':
    print(circ.draw())

cobyla = COBYLA()
cobyla.set_options(maxiter=250)

constrained = constrainedQAOA(cost_operator, cobyla, mixer_operator, 1)

backend = BasicAer.get_backend('statevector_simulator')
seed = 50
constrained.random_seed = seed

quantum_instance = QuantumInstance(backend=backend, seed=seed, seed_mapper=seed)

result = constrained.run(quantum_instance)

result['eigvecs']

