from functools import reduce
from itertools import product

from qiskit import BasicAer, QuantumRegister, ClassicalRegister
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 qiskit_aqua.algorithms import QAOA

import numpy as np
import qutip as qt

from qiskit_aqua.components.initial_states import Custom

edges = [(0, 1), (1, 2), (2, 3)]
colors   = 4
vertices = max(reduce(lambda x, y: x + y, [list(edge) for edge in edges])) + 1
logcolors = int(np.log2(colors))
n_qubits = int(vertices * logcolors)

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

up = qt.basis(2, 0)
dn = qt.basis(2, 1)

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

cobyla = COBYLA()
cobyla.set_options(maxiter=250)
p = 2 # steps of QAOA
constrained = QAOA(cost_operator, cobyla, p)

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['eigvals']

result['opt_params']

result['eigvecs']

