# Copied and modified from qiskit_aqua.algorithms.adaptive.qaoa.qaoa.py
# and qiskit_aqua.algorithms.adaptive.qaoa.varform.py
# =============================================================================

import logging

from qiskit_aqua.algorithms import QuantumAlgorithm
from qiskit_aqua import AquaError, PluggableType, get_pluggable_class
from qiskit_aqua.algorithms.adaptive import VQE
from qiskit_aqua.algorithms.adaptive.qaoa.varform import QAOAVarForm

logger = logging.getLogger(__name__)


class constrainedQAOA(VQE):
    """
    The Quantum Approximate Optimization Algorithm.
    See https://arxiv.org/abs/1411.4028
    """

    CONFIGURATION = {
        'name': 'QAOA.Variational',
        'description': 'Quantum Approximate Optimization Algorithm',
        'input_schema': {
            '$schema': 'http://json-schema.org/schema#',
            'id': 'qaoa_schema',
            'type': 'object',
            'properties': {
                'operator_mode': {
                    'type': 'string',
                    'default': 'matrix',
                    'oneOf': [
                        {'enum': ['matrix', 'paulis', 'grouped_paulis']}
                    ]
                },
                'p': {
                    'type': 'integer',
                    'default': 1,
                    'minimum': 1
                },
                'initial_point': {
                    'type': ['array', 'null'],
                    "items": {
                        "type": "number"
                    },
                    'default': None
                },
                'batch_mode': {
                    'type': 'boolean',
                    'default': False
                }
            },
            'additionalProperties': False
        },
        'problems': ['ising'],
        'depends': ['optimizer'],
        'defaults': {
            'optimizer': {
                'name': 'COBYLA'
            },
        }
    }

    def __init__(self, cost, optimizer, mixer, p=1, initial_state=None, operator_mode='matrix', initial_point=None,
                 batch_mode=False, aux_operators=None):
        """
        Args:
            operator (Operator): Qubit operator
            operator_mode (str): operator mode, used for eval of operator
            p (int) : the integer parameter p as specified in https://arxiv.org/abs/1411.4028
            optimizer (Optimizer) : the classical optimization algorithm.
            initial_point (numpy.ndarray) : optimizer initial point.
        """
        self.validate(locals())
        var_form = constrainedQAOAVarForm(cost, p, mixer, initial_state)
        super().__init__(cost, var_form, optimizer,
                         operator_mode=operator_mode, initial_point=initial_point)

    @classmethod
    def init_params(cls, params, algo_input):
        """
        Initialize via parameters dictionary and algorithm input instance
        Args:
            params (dict): parameters dictionary
            algo_input (EnergyInput): EnergyInput instance
        """
        if algo_input is None:
            raise AquaError("EnergyInput instance is required.")

        operator = algo_input.qubit_op

        qaoa_params = params.get(QuantumAlgorithm.SECTION_KEY_ALGORITHM)
        operator_mode = qaoa_params.get('operator_mode')
        p = qaoa_params.get('p')
        initial_point = qaoa_params.get('initial_point')
        batch_mode = qaoa_params.get('batch_mode')

        # Set up optimizer
        opt_params = params.get(QuantumAlgorithm.SECTION_KEY_OPTIMIZER)
        optimizer = get_pluggable_class(PluggableType.OPTIMIZER,
                                        opt_params['name']).init_params(opt_params)

        return cls(operator, optimizer, p=p, operator_mode=operator_mode,
                   initial_point=initial_point, batch_mode=batch_mode,
                   aux_operators=algo_input.aux_ops)

class constrainedQAOAVarForm(QAOAVarForm):
    def __init__(self, cost_operator, p, mixer_operator=None, initial_state=None):
        super().__init__(cost_operator, p, initial_state)

        if mixer_operator is None:
            v = np.zeros(self._cost_operator.num_qubits)
            ws = np.eye(self._cost_operator.num_qubits)
            self._mixer_operator = reduce(
                lambda x, y: x + y,
                [
                    Operator([[1, Pauli(v, ws[i, :])]])
                    for i in range(self._cost_operator.num_qubits)
                ]
            )
        else:
            self._mixer_operator = mixer_operator

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, 0)]
vertices = 3
colors = 3
n_qubits = vertices * colors

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

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

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 = 1 # steps of QAOA
constrained = constrainedQAOA(cost_operator, cobyla, mixer_operator, p, init_state)

from qiskit import IBMQ
IBMQ.load_accounts()


# Please input the 20 qubit architecture!
backend = IBMQ.get_backend('ibmq_20_tokyo')

quantum_instance = QuantumInstance(backend=backend_melbourne)

result = constrained.run(quantum_instance)

result['eigvals']

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

result['eval_count']

np.savetxt('3_triangle_3colors_p1_20q_qpu', result['eigvecs'][0])