<filename>tests/RQAOA.py
import time

from typing import Optional
from collections import defaultdict
import numpy as np
import networkx as nx
from Qcover.optimizers import Optimizer, COBYLA
from Qcover.backends import Backend, CircuitByQiskit
import warnings
warnings.filterwarnings("ignore")


class RQcover:
    """
    Qcover is a QAOA solver
    """
    # pylint: disable=invalid-name
    def __init__(self,
                 graph: nx.Graph = None,
                 p: int = 1,
                 optimizer: Optional[Optimizer] = COBYLA(),
                 backend: Optional[Backend] = CircuitByQiskit()
                 ) -> None:

        assert graph is not None
        self._simple_graph = graph
        self._p = p
        self._backend = backend
        self._backend._p = p
        self._optimizer = optimizer
        self._optimizer._p = p

        self._nodes_weight = []
        self._edges_weight = []
        # self._path = dict()

    @classmethod
    def generate_weighted_graph(cls, nodes, edges, weight_range=10):
        """
        generate graph from nodes list and edges list which identify the nodes and edges
        that should be add in graph, and the random weight range of every node and edge.
        Args:
            nodes (list/set): list of node idex / node-weight map, element form is tuple(nid, weight)
            edges (list/set): list of edge: (e_idex1, e_idex2) / edge-weight map, element form is tuple(nid1, nid2, edge_weight)
            weight_range (int): random weight range of every node and edge
        Returns:
            g (nx.Graph): graph generated by args
        """
        g = nx.Graph()
        if isinstance(nodes, list) and isinstance(edges, list):
            for v in nodes:
                w = np.random.choice(range(weight_range))
                g.add_node(v, weight=w)

            for e in edges:
                w = np.random.choice(range(weight_range))
                g.add_edge(e[0], e[1], weight=w)
        else:
            for item in nodes:
                g.add_node(item[0], weight=item[1])

            for item in edges:
                g.add_edge(item[0], item[1], weight=item[2])
        return g

    @staticmethod
    def get_graph_weights(graph):
        """
        get the weights of nodes and edges in graph
        Args:
            graph (nx.Graph): graph to get weight of nodes and edges
        Return:
            node weights form is dict{nid1: node_weight}, edges weights form is dict{(nid1, nid2): edge_weight}
        """
        nodew = nx.get_node_attributes(graph, 'weight')
        edw = nx.get_edge_attributes(graph, 'weight')
        edgew = edw.copy()
        for key, val in edw.items():
            edgew[(key[1], key[0])] = val

        return nodew, edgew

    def generate_subgraph(self, graph, dtype: str, p):
        """
        according to the arguments of dtype and p to generate subgraphs from graph
        Args:
            graph (nx.Graph): graph to be decomposed
            dtype (string): set "node" or "edge", the ways according to which to decompose the graph
            p (int): the p of subgraphs
        Return:
            subg_dict (dict) form as {node_id : subg, ..., (node_id1, node_id2) : subg, ...}
        """
        if dtype not in ["node", "edge"]:
            print("Error: wrong dtype, dtype should be node or edge")
            return None

        nodes_weight, edges_weight = self.get_graph_weights(graph)

        subg_dict = defaultdict(list)
        if dtype == 'node':
            for node in graph.nodes:
                node_set = {(node, nodes_weight[node])}
                edge_set = set()
                for i in range(p):
                    new_nodes = { (nd2, nodes_weight[nd2]) for nd1 in node_set for nd2 in graph[nd1[0]] }
                    new_edges = {(nd1[0], nd2, edges_weight[nd1[0], nd2]) for nd1 in node_set for nd2 in graph[nd1[0]]}
                    node_set |= new_nodes
                    edge_set |= new_edges

                subg = self.generate_weighted_graph(node_set, edge_set)
                subg_dict[node] = subg
        else:
            for edge in graph.edges:
                node_set = {(edge[0], nodes_weight[edge[0]]), (edge[1], nodes_weight[edge[1]])}
                edge_set = {(edge[0], edge[1], edges_weight[edge[0], edge[1]])}

                for i in range(p):
                    new_nodes = {(nd2, nodes_weight[nd2]) for nd1 in node_set for nd2 in graph[nd1[0]]}
                    new_edges = {(nd1[0], nd2, edges_weight[nd1[0], nd2]) for nd1 in node_set for nd2 in
                                 graph.adj[nd1[0]]}
                    node_set |= new_nodes
                    edge_set |= new_edges

                subg = self.generate_weighted_graph(node_set, edge_set)
                subg_dict[edge] = subg
        return subg_dict

    def graph_decomposition(self, graph, p):
        """
        according to dtype to decompose graph
        Args:
            graph (nx.Graph): graph to be composed
            p (int): the p of subgraphs
        """
        if p <= 0:
            warnings.warn(" the argument of p should be >= 1 in qaoa problem, "
                          "so p would be set to the default value at 1")
            p = 1

        subg_node = self.generate_subgraph(graph, "node", p)
        subg_edge = self.generate_subgraph(graph, "edge", p)
        element_to_graph = {}
        for k, v in subg_node.items():
            element_to_graph[k] = v

        for k, v in subg_edge.items():
            element_to_graph[k] = v
        return element_to_graph

    @staticmethod
    def solve_basic_graph(graph: nx.Graph):
        nodes = graph.nodes
        min_H = np.inf
        basic_sol = dict()
        for i in range(len(nodes)):
            cur_H = 0
            tmp_sol = dict()
            binary_list = [int(x) for x in bin(i)[2:]]
            for i, nd in enumerate(nodes):
                tmp_sol[nd] = binary_list[i]

            for nd in graph.nodes:
                sigma_nd = 1 if tmp_sol[nd] == 1 else -1
                cur_H += graph.nodes[nd]["weight"] * sigma_nd

            for ed in graph.edges:
                u, v = ed
                sigma_u = 1 if tmp_sol[u] == 1 else -1
                sigma_v = 1 if tmp_sol[v] == 1 else -1
                cur_H += graph.adj[u][v]["weight"] * sigma_u * sigma_v

            if cur_H < min_H:
                min_H = cur_H
                basic_sol = tmp_sol
        return basic_sol

    @staticmethod
    def get_solution(graph, basic_sol: dict = None):
        sol = defaultdict(lambda: -1)
        node_num = len(graph.nodes)
        queue = list(basic_sol.keys()) if basic_sol is not None else []
        for i in range(node_num):
            if sol[i] == -1:  #len(graph.neighbors(i)) != 0 and
                queue.append(i)
                while len(queue) > 0:
                    u = queue[0]
                    queue.pop(0)
                    if sol[u] == -1:
                        sol[u] = 0
                    neighbor_u = list(graph.neighbors(u))
                    for v in neighbor_u:
                        if sol[v] == -1:
                            sol[v] = sol[u] if graph.adj[u][v]["weight"] > 0 else 1 - sol[u]
                            queue.append(v)
        return sol

    def calculate(self, pargs, p=None):
        """
        The framework function which use the backend to calculate the value of expectation,
        and be used as the object function in the optimization function of the optimizer
        Args:
            pargs: the value of the parameter alpha and beta in the circuit
            p: the integer used to define the number of layers the current circuit needs to be superimposed
        Returns:
            the value of expectation calculated by backends
        """
        p = self._p if p is None else p
        element_to_graph = self.graph_decomposition(graph=self._simple_graph, p=p)

        self._backend._pargs = pargs
        self._backend._element_to_graph = element_to_graph
        return self._backend.expectation_calculation(p)

    def run_rqaoa(self, node_threshold=1):
        current_g = self._simple_graph
        node_num = len(current_g.nodes)
        pathg = nx.Graph()
        while node_num > node_threshold:
            nodes_weight, edges_weight = self.get_graph_weights(current_g)
            self._backend._nodes_weight = nodes_weight
            self._backend._edges_weight = edges_weight

            self._optimizer.optimize(objective_function=self.calculate)   #x, fun, nfev =
            exp_sorted = sorted(self._backend._element_expectation.items(), key=lambda item: abs(item[1]), reverse=True)
            u, v = exp_sorted[0][0]
            correlation = 1 if exp_sorted[0][1] - 1e-8 > 0 else -1
            pathg.add_edge(u, v, weight=correlation)
            if u > v:
                u, v = v, u

            for nd in current_g.neighbors(v):
                if nd == u:
                    continue
                if nd not in current_g.neighbors(u):
                    current_g.add_edge(u, nd, weight=correlation * current_g.adj[v][nd]["weight"])
                else:
                    current_g.adj[u][nd]["weight"] += correlation * current_g.adj[v][nd]["weight"]
                    # current_g.adj[nd][u]["weight"] = current_g.adj[u][nd]["weight"]
            current_g.remove_node(v)
            node_num = len(current_g.nodes)

        basic_sol = None
        if node_threshold > 1:
            basic_sol = self.solve_basic_graph(current_g)

        solution = self.get_solution(pathg, basic_sol)
        return solution

    def run(self, is_parallel=False, mode="RQAOA"):
        """
        run Qcover code to solve the given problem
        Args:
            node_num: nodes number in the graphical representation of the problem
            edge_num: edges number in the graphical representation of the problem
            is_parallel: run programs in parallel or not

        Returns:
            results of the problem, which including the optimal parameters of the circuit model,
            the optimal expectation value and the number of times the optimizer iterates
        """
        self._backend._is_parallel = is_parallel
        if mode != "QAOA":
            res = self.run_rqaoa()

        return res


# usage example
if __name__ == '__main__':
    p = 1

    g = nx.Graph()
    nodes = [(0, 1), (1, 1), (2, 1)]
    edges = [(0, 1, -1), (1, 2, -1)]
    # edges = [(0, 1, 3), (1, 2, 2), (0, 2, 1)]
    for nd in nodes:
        u, w = nd[0], nd[1]
        g.add_node(int(u), weight=int(w))
    for ed in edges:
        u, v, w = ed[0], ed[1], ed[2]
        g.add_edge(int(u), int(v), weight=int(w))

    from Qcover.applications import MaxCut
    mxt = MaxCut(g)
    ising_g, shift = mxt.run()

    optc = COBYLA(options={'tol': 1e-3, 'disp': True})
    qiskit_bc = CircuitByQiskit(expectation_calc_method="statevector")

    rqc = RQcover(ising_g, p,
                  optimizer=optc,
                  backend=qiskit_bc)

    st = time.time()
    sol = rqc.run()
    # rqc.run_rqaoa()
    ed = time.time()
    print("time cost is:", ed - st)
    print("solution is:", sol)