import qiskit
qiskit.__version__

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 qiskit_aqua.components.initial_states import Custom

from constrainedqaoa import constrainedQAOA

import numpy as np
import qutip as qt

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

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

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

reduce(
            lambda x, y: x + y, [list(edge) for edge in edges])

def W(size, copies):
    initial_list = [dn] + [up] * (size - 1)
    cycles = [[initial_list[i - j] for i in range(size)] for j in range(size)]
    W_1copy = sum([qt.tensor(states) for states in cycles])
    return qt.tensor([W_1copy] * copies)

amplitudes = W(colors, vertices).full().T.tolist()[0]
init_state = Custom(n_qubits, state_vector=amplitudes)

W(colors, vertices)

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))
            ]
)

# Fix redundancies
if colors == 2:
    mixer_operator.scaling_coeff(1/2)

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

from qiskit import BasicAer


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

np.round(result['eigvecs'], 4)

result['eval_count']

