# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/01_optimization.ipynb (unless otherwise specified).

__all__ = ['QuantumOptimizer', 'SUPPORTED_TASKS', 'DEFAULT_LOCAL']

# Cell
import pennylane as qml
import numpy as np
from matplotlib import pyplot as plt
import networkx as nx
import warnings

SUPPORTED_TASKS = {'maximum_clique': qml.qaoa.max_clique,
                   'minimum_vertex_cover': qml.qaoa.min_vertex_cover}
DEFAULT_LOCAL = 'braket.local.qubit'

class QuantumOptimizer:
    """quantum-based combinatorial optimization with Amazon Braket


    Usage:

    ```python
    >>> qo = QuantumOptimizer(g,
                              task='maximum_clique',
                              device_name='braket.local.qubit', # default braket local simulator
                              device_arn='device_arn',
                              s3_folder='s3_folder')
        qo.fit()
        results = qo.results()
    ```

    **Parameters:**

    * **g** : a networkx Graph object
    * **task** : one of {'maximum_clique'}
    * **device_name** : one of {'braket.local.qubit', 'braket.aws.qubit'}
    * **device_arn** : Device ARN. Only required if name != 'braket.local.simulator'.
    * **s3_folder** : S3 folder. Only required if name != 'braket.local.simulator'.
    * **n_layers**  : number of layers. default:4
    * **constrained**  : default:False
    * **params**  : user-supplied initialization of params.  If None, a default random initialization is used.


    """
    def __init__(self, g, task=None, n_layers=4, constrained=False,
                 device_name='braket.local.qubit', device_arn=None, s3_folder=None,
                 params=None):
        """
        constructor
        """
        # error checks
        if task not in SUPPORTED_TASKS: raise ValueError(f'task {task} is not supported')
        if device_name != DEFAULT_LOCAL and (device_arn is None or s3_folder is None):
            raise ValueError('device_arn and s3_folder are required if using managed AWS device')
        if not isinstance(g, nx.Graph): raise ValueError('g must be instance of networkx.Graph')

        # input vars
        self.g = g
        self.task = task
        self.device_name = device_name
        self.device_arn = device_arn
        self.s3_folder = s3_folder
        self.n_layers = n_layers
        self.constrained=constrained

        # computed vars
        self.wires = self.g.number_of_nodes()
        np.random.seed(1967)
        self.params = self._default_params() if params is None else params
        self.circuit, self.cost_function = self._prepare()
        self._fit_called = False

    def _default_params(self):
        #return np.random.uniform(size=[2, n_layers])
        return np.full((2,self.n_layers), 0.5)


    def _device(self, shots=None):
        kwargs = {}
        if self.device_name != DEFAULT_LOCAL:
            kwargs['device_arn'] = self.device_arn
            kwargs['s3_destination_folder'] = self.s3_folder
        if shots is not None:
            kwargs['shots'] = shots
        dev = qml.device(self.device_name, wires=self.wires, **kwargs)
        return dev

    def _prepare(self):

        # qaoa layer
        cost_h, mixer_h = SUPPORTED_TASKS[self.task](self.g, constrained=self.constrained)
        def qaoa_layer(gamma, alpha):
            qml.qaoa.cost_layer(gamma, cost_h)
            qml.qaoa.mixer_layer(alpha, mixer_h)

        # build circuit
        def circuit(params, **kwargs):
            for i in range(self.wires):  # Prepare an equal superposition over all qubits
                qml.Hadamard(wires=i)
            qml.layer(qaoa_layer, self.n_layers, params[0], params[1])

        # build cost function
        cost_function = qml.ExpvalCost(circuit, cost_h, self._device(), optimize=True)

        return circuit, cost_function


    def fit(self, iterations=10, verbose=1):
        """
        approximate a solution to given task
        """
        # progress bar
        from collections import OrderedDict
        from .utils import add_metrics_to_log, log_to_message, ProgressBar
        if verbose:
            print(f"Optimize for {iterations} iterations...")
            pb = ProgressBar(iterations)
        log = OrderedDict()

        # optimize
        optimizer = qml.GradientDescentOptimizer()
        for i in range(iterations):
            self.params = optimizer.step(self.cost_function, self.params)
            cost_eval = self.cost_function(self.params)
            #if verbose:
            #    print(f"Completed iteration {i + 1}, cost function:", cost_eval)
            log['cost'] = cost_eval
            log['iteration'] = i+1
            if verbose:
                pb.bar(i, log_to_message(log))
        self._fit_called = True
        return

    def results(self, shots = 100000):
        """
        Sample and return results.
        Returns a tuple of the form: `(nodes, probs)`, where `nodes` are the list
        of vertices comprising the most likely solution and `probs` are the raw solution probabilities.
        """

        if not self._fit_called: raise Exception('The fit method must be called first.')

        dev = self._device(shots=shots)

        @qml.qnode(dev)
        def samples(params):
            self.circuit(params)
            return [qml.sample(qml.PauliZ(i)) for i in range(self.wires)]

        # compute probabilities
        from collections import Counter
        s = samples(self.params).T
        s = (1 - s.numpy()) / 2
        s = map(tuple, s)
        counts = Counter(s)
        indx = np.ndindex(*[2] * self.wires)
        probs = {p: counts.get(p, 0) / shots for p in indx}

        # extract nodes
        import operator
        result_bool = max(probs.items(), key=operator.itemgetter(1))[0]
        nodes = [i for i, b in enumerate(result_bool) if b]

        return nodes, probs


    def plot_samples(self, probs):
        """
        Plot sample for toy problems for testing purposes.
        """

        import matplotlib.pyplot as plt

        plt.style.use("seaborn")
        labels = ["{0:{fill}8b}".format(i, fill='0') for i in range(len(probs))]

        plt.bar(range(2 ** self.wires), probs.values())
        plt.xticks([i for i in range(len(probs))], labels, rotation='vertical', size=12)
        plt.yticks(size=12)

        plt.xlabel("Sample", size=20)
        plt.ylabel("Probability", size=20)

        fig = plt.gcf()
        fig.set_size_inches(16, 8)
        plt.show()