import itertools
import os
from collections import defaultdict
import matplotlib.pyplot as plt
from multiprocessing import Pool, cpu_count
import networkx as nx
import cirq
from Qcover.backends import Backend
from Qcover.utils import get_graph_weights


class CircuitByCirq(Backend):
    """generate a instance of CircuitByCirq"""

    def __init__(self,
                 research: str = "QAOA",
                 is_parallel: bool = None) -> None:
        """initialize a instance of CircuitByCirq"""
        super(CircuitByCirq, self).__init__()
        self._p = None
        self._origin_graph = None
        self._is_parallel = False if is_parallel is None else is_parallel
        self._research = research

        self._nodes_weight = None
        self._edges_weight = None
        self._element_to_graph = None
        self._pargs = None
        self._expectation_path = []
        self._element_expectation = dict()

    @property
    def element_expectation(self):
        return self._element_expectation

    def get_operator(self, element, qubit_num):
        qubits = cirq.LineQubit.range(qubit_num)
        nodes = [element] if isinstance(element, int) else list(element)
        op = 1
        for i in range(qubit_num):
            if i in nodes:
                op *= cirq.Z(qubits[i])
        return op

    def get_QAOA_circuit(self, p, graph, node_to_qubit):
        gamma_list, beta_list = self._pargs[: p], self._pargs[p:]
        circ = cirq.Circuit()
        ql = cirq.LineQubit.range(len(graph))
        circ.append(cirq.H(ql[i]) for i in range(len(graph)))
        for k in range(p):
            for i in graph.nodes:
                u = node_to_qubit[i]
                circ.append(cirq.rz(2 * gamma_list[k] * self._nodes_weight[i]).on(ql[u]))

            for edge in graph.edges:
                u, v = node_to_qubit[edge[0]], node_to_qubit[edge[1]]
                if u == v:
                    continue

                circ.append(cirq.CX(ql[u], ql[v]))
                circ.append(cirq.rz(2 * gamma_list[k] * self._edges_weight[edge[0], edge[1]]).on(ql[v]))
                circ.append(cirq.CX(ql[u], ql[v]))

            for nd in graph.nodes:
                u = node_to_qubit[nd]
                circ.append(cirq.Moment(cirq.rx(2 * beta_list[k]).on(ql[u])))

        return circ

    def get_GHZ_circuit(self, p, graph, node_to_qubit):
        gamma_list, beta_list = self._pargs[: p], self._pargs[p:]
        circ = cirq.Circuit()
        ql = cirq.LineQubit.range(len(graph))
        circ.append(cirq.H(ql[i]) for i in range(len(graph)))

        pivot = len(self._origin_graph) // 2
        for k in range(p):
            for edge in graph.edges:
                u, v = node_to_qubit[edge[0]], node_to_qubit[edge[1]]
                if u == v:
                    continue

                circ.append(cirq.CX(ql[u], ql[v]))
                circ.append(cirq.rz(2 * gamma_list[k] * self._edges_weight[edge[0], edge[1]]).on(ql[v]))
                circ.append(cirq.CX(ql[u], ql[v]))

            for nd in graph.nodes:
                if nd != pivot:
                    u = node_to_qubit[nd]
                    circ.append(cirq.Moment(cirq.rx(2 * beta_list[k]).on(ql[u])))

        return circ

    def get_expectation(self, element_graph, p=None):
        """
        transform the graph to circuit according to the computing_framework
        Args:
            graph (nx.Graph): graph to be transformed to circuit
            params (np.array): Optimal parameters
            original_e (Optional[None, int, tuple])
        Return:
            if original_e=None, then the graph is the whole original graph generated by
            generate_weighted_graph(), so just return the circuit transformed by it

            if original_e is a int, then the subgraph is generated by node(idx = original_e
            in whole graph), so return the it's idx mapped by node_to_qubit[], and the circuit

            if original_e is a tuple, then the subgraph is generated by edge(node idx = original_e
            in whole graph), so return the it's idx mapped by node_to_qubit[] as
            tuple(mapped node_id1, mapped node_id2), and the circuit
        """
        if self._is_parallel is False:
            p = self._p if p is None else p
            original_e, graph = element_graph
        else:
            p = self._p if len(element_graph) == 1 else element_graph[1]
            original_e, graph = element_graph[0]

        node_to_qubit = defaultdict(int)
        node_list = list(graph.nodes)
        for i in range(len(node_list)):
            node_to_qubit[node_list[i]] = i

        if self._research == "QAOA":
            circ = self.get_QAOA_circuit(p, graph, node_to_qubit)
        elif self._research == "GHZ":
            pivot = len(self._origin_graph) // 2
            if pivot in graph:
                circ = self.get_GHZ_circuit(p, graph, node_to_qubit)
            else:
                circ = self.get_QAOA_circuit(p, graph, node_to_qubit)

        qubits = cirq.LineQubit.range(len(node_list))
        qubit_map = dict(zip(qubits, range(len(node_list))))

        if isinstance(original_e, int):
            weight = self._nodes_weight[original_e]
            op = cirq.Z(qubits[node_to_qubit[original_e]])
        else:
            weight = self._edges_weight[original_e]
            op = cirq.Z(qubits[node_to_qubit[original_e[0]]]) * cirq.Z(qubits[node_to_qubit[original_e[1]]])

        state = cirq.final_state_vector(circ)
        exp_res = op.expectation_from_state_vector(state, qubit_map=qubit_map)

        return weight, exp_res.real

    def expectation_calculation(self, p=None):
        if self._nodes_weight is None or self._edges_weight is None:
            nodes_weight, edges_weight = get_graph_weights(self._origin_graph)
            self._nodes_weight, self._edges_weight = nodes_weight, edges_weight

        self._element_expectation = {}
        if self._is_parallel:
            return self.expectation_calculation_parallel(p)
        else:
            return self.expectation_calculation_serial(p)

    def expectation_calculation_serial(self, p=None):
        cpu_num = cpu_count()
        os.environ['OMP_NUM_THREADS'] = str(cpu_num)
        os.environ['OPENBLAS_NUM_THREADS'] = str(cpu_num)
        os.environ['MKL_NUM_THREADS'] = str(cpu_num)
        os.environ['VECLIB_MAXIMUM_THREADS'] = str(cpu_num)
        os.environ['NUMEXPR_NUM_THREADS'] = str(cpu_num)

        res = 0
        for item in self._element_to_graph.items():
            w_i, exp_i = self.get_expectation(item, p)
            if isinstance(item[0], tuple):
                self._element_expectation[item[0]] = exp_i
            res += w_i * exp_i

        print("Total expectation of original graph is: ", res)
        self._expectation_path.append(res)
        return res

    def expectation_calculation_parallel(self, p=None):
        cpu_num = 1
        os.environ['OMP_NUM_THREADS'] = str(cpu_num)
        os.environ['OPENBLAS_NUM_THREADS'] = str(cpu_num)
        os.environ['MKL_NUM_THREADS'] = str(cpu_num)
        os.environ['VECLIB_MAXIMUM_THREADS'] = str(cpu_num)
        os.environ['NUMEXPR_NUM_THREADS'] = str(cpu_num)

        circ_res = []
        args = list(itertools.product(self._element_to_graph.items(), [p]))
        pool = Pool(os.cpu_count())
        circ_res.append(pool.map(self.get_expectation, args))

        pool.terminate()  # pool.close()
        pool.join()
        res = 0
        for it in circ_res[0]:
            res += it[0] * it[1]

        # res = sum(circ_res[0])
        print("Total expectation of original graph is: ", res)
        self._expectation_path.append(res)
        return res

    def optimization_visualization(self):
        plt.figure()
        plt.plot(range(1, len(self._expectation_path) + 1), self._expectation_path, "ob-", label="cirq")
        plt.ylabel('Expectation value')
        plt.xlabel('Number of iterations')
        plt.legend()
        plt.show()
