import kahypar
import opt_einsum as oe
import numpy
from numpy import random
import itertools
import copy
from os.path import join, abspath, dirname
import cma
import os
import tempfile
import pickle
from multiprocessing import Pool
import cProfile
from scipy import sparse
from acqdp.utility import opt_einsum_ext

from acqdp.tensor_network.local_optimizer import OrderResolver
from acqdp.tensor_network.contraction import ContractionScheme, ContractionCost, OrderCounter
from acqdp.tensor_network.order_finder import OrderFinder

CMA_TEMP_DIR = tempfile.TemporaryDirectory()
CMA_FILE_PREFIX = CMA_TEMP_DIR.name + os.sep

KAHYPAR_PROFILE_DIR = join(abspath(dirname(__file__)), 'kahypar_profiles')

modes = ['direct', 'recursive']
objectives = ['cut', 'km1']
queryOrderResolver = OrderResolver(optimal=9, dp=9)


def query(order_finder,
          s,
          eps,
          graph_copy,
          seed=numpy.random.randint(2**31 - 1),
          counter=None,
          objective=0,
          cutoff=25):
    K = s // 2 + 3
    mode = s % 2
    tg, nodes = graph_copy.tanner_graph
    order = order_finder._order_by_params(numpy.array(tg.todense()),
                                         copy.copy(nodes),
                                         K=K,
                                         eps=eps,
                                         mode=mode,
                                         objective=0,
                                         cutoff=cutoff,
                                         seed=seed,
                                         counter=counter)
    res = queryOrderResolver.order_to_contraction_scheme(graph_copy, order)
    return res


class KHPOrderFinder(OrderFinder):
    """
    :class:`KHPOrderFinder` utilizes hypergraph decomposition and cma-es opmization scheme for finding contraction orders for intermediate-size tensor networks.

    :ivar num_threads: Number of threads to concurrently search for the best order.
    :ivar num_iters: Number of iterations for cma optimization.
    :ivar num_cmas: Number of times cma-es optimization scheme is called for exploring better parameter combinations.
    :ivar cma_args: Arguments for cma-es optimization scheme. See the documentation of cma for more details
    :ivar eps: In case a good parameter combination is already known, inputting a eps would skip the cma optimization scheme and directly yield orders found by hypergraph decomposition.
    """

    def __init__(
            self,
            num_threads=28,
            num_iters=50,
            num_cmas=1,
            cma_args={
                'bounds': [0.5, 0.5, 0.5],
                'std_dev': 0.2,
                'kwargs': {
                    'seed': 1175,
                    'bounds': [[0.01, 0, 0], [1, 1, 1]],
                    'tolfun': 1e-2,
                    'verb_filenameprefix': CMA_FILE_PREFIX
                },
            },
            eps=None,
            **kwargs):
        self.num_threads = num_threads
        self.num_iters = num_iters
        self.num_cmas = num_cmas
        self.eps = eps
        self.cma_args = cma_args

    def _simplify(self, graph):
        order = []
        counter = OrderCounter()
        flag = True
        while flag:
            flag = False
            for edge in graph.edges:
                found = False
                for n0, n1 in itertools.combinations(graph.network[edge], 2):
                    s0, s1 = set(graph.network[n0]), set(graph.network[n1])
                    fix = len([
                        e for e in s0.intersection(s1)
                        if len(graph.network[e]) == 2
                    ])
                    sz = max(len(s0), len(s1))
                    new_sz = len(s0.union(s1)) - fix
                    if new_sz <= sz:
                        flag = True
                        new_name = counter.cnt
                        new_name, _ = graph.encapsulate(nodes=[n0[1], n1[1]],
                                                        stn_name=new_name)
                        order.append([[n0[1], n1[1]], new_name])
                        found = True
                        break
                if found:
                    break
        return order, counter

    def find_order(self, graph):
        graph_copy = graph._expand_and_delta()
        graph_copy.fix()
        init_order, counter = self._simplify(graph_copy)
        if self.eps is None:
            seed_f = numpy.random.uniform(0, 1)

            def probe(x):
                eps = [x[0], 0.95 + 0.049 * x[1]]
                with Pool(self.num_threads) as p:
                    reses = p.starmap(
                        query,
                        [(self, s, eps, graph_copy, int(
                            (2**31 - 1) * seed), copy.copy(counter))
                         for s in range(self.num_threads)
                         for seed in x[2:]])
                res = min(reses, key=lambda x: (x.cost, reses.index(x)))
                return res

            curr_y = None
            for _ in range(self.num_cmas):
                seed_f = numpy.random.uniform(0, 1)
                self.cma_args['kwargs']['seed'] = numpy.random.randint(2**31 -
                                                                       1)

                def probe_value(x):
                    res = probe([x[0], x[1], seed_f])
                    return numpy.log10(float(res.cost.s))

                es = cma.CMAEvolutionStrategy(
                    self.cma_args['bounds'], self.cma_args['std_dev'],
                    self.cma_args['kwargs']).optimize(probe_value,
                                                      iterations=self.num_iters)
                x = es.result[0]
                y = es.result[1]
                if curr_y is None or curr_y > y:
                    self.eps = [x[0], 0.95 + 0.049 * x[1]]
                    curr_y = y
            res = probe(x)
            yield ContractionScheme(init_order + res.order, cost=res.cost)
        while True:
            reses = [
                query(self,
                      s,
                      self.eps,
                      graph_copy,
                      seed=numpy.random.randint(2**31 - 1),
                      counter=copy.copy(counter)) for s in range(10)
            ]
            res = min(reses, key=lambda x: (x.cost, reses.index(x)))
            yield ContractionScheme(init_order + res.order, cost=res.cost)

    def _order_by_params(self, graph, nodes, **kwargs):
        graph = self._refresh(graph)
        counter = kwargs.get('counter')

        K = int(kwargs.get('K'))
        eps = kwargs.get('eps', None)
        top = kwargs.get('top', True)
        kwargs_c = copy.copy(kwargs)
        kwargs_c['eps'] = eps[0]
        if not top:
            kwargs_c['K'] = 2
            kwargs_c['eps'] = eps[1]
        context = self._set_context(**kwargs_c)
        cutoff = kwargs.get('cutoff')
        hypergraph, _, edges = self._init_hypergraph_tanner(graph, kwargs_c['K'])
        if not edges:
            return [[nodes, '#']]
        kahypar.partition(hypergraph, context)
        partitions_names = [[] for _ in range(K)]
        for i, n in list(enumerate(nodes)):
            partitions_names[hypergraph.blockID(i)].append(n)
        if any(len(part) == len(nodes) for part in partitions_names):
            return [[nodes, '#']]
        order = []
        all_names = []
        for (i, partite_nodes) in enumerate(
                sorted(partitions_names, key=lambda x: -len(x))):
            if len(partite_nodes) == 0:
                continue
            elif len(partite_nodes) > 1:
                new_stn = counter.cnt
                all_names.append(new_stn)
                new_g, graph, nodes = self._sub(graph, nodes, partite_nodes,
                                               new_stn)
                assert graph.shape[0] == len(nodes) + 1, (graph.shape[0],
                                                          len(nodes))
                kwargs['top'] = False
                if len(partite_nodes) > cutoff:
                    new_order = self._order_by_params(new_g, partite_nodes,
                                                     **kwargs)
                    new_order[-1][1] = new_stn
                else:
                    new_order = [[partite_nodes, new_stn]]
                order += new_order
            else:
                all_names.append(partite_nodes[0])
        order.append([all_names, '#'])
        return order

    def _set_context(self, **kwargs):
        mode = modes[int(kwargs.get('mode'))]
        objective = objectives[int(kwargs.get('objective'))]
        K = int(kwargs.get('K'))
        eps = kwargs.get('eps')
        seed = kwargs.get('seed')
        profile_mode = {'direct': 'k', 'recursive': 'r'}[mode]
        profile = f"{objective}_{profile_mode}KaHyPar_sea20.ini"
        context = kahypar.Context()
        context.loadINIconfiguration(join(KAHYPAR_PROFILE_DIR, profile))
        context.setK(K)
        context.setSeed(seed)
        context.setEpsilon(kwargs.get('epsilon', eps * (K - 1)))
        context.suppressOutput(kwargs.get('quiet', True))
        return context

    def _sub(self, graph, nodes, sub_nodes, new_name):
        inds = [False] + [n in sub_nodes for n in nodes]
        out_nodes = [n for n in nodes if n not in sub_nodes] + [new_name]
        subb = graph[inds]
        rest = graph[numpy.invert(inds)]
        open_edges = numpy.any(subb, axis=0) & numpy.any(rest, axis=0)
        return numpy.append(open_edges[None, :], subb,
                            axis=0), numpy.append(rest,
                                                  open_edges[None, :],
                                                  axis=0), out_nodes

    def _init_hypergraph_tanner(self, graph, k=2):
        nodes = list(range(1, graph.shape[0]))
        edges = list(range(graph.shape[1]))
        hyperedge_indices = [0]
        hyperedges = []
        edge_weights = []
        for edge in range(graph.shape[1]):
            hyperedges += list(numpy.where(graph[1:, edge])[0])
            hyperedge_indices.append(len(hyperedges))
            edge_weights.append(1)
        hp = kahypar.Hypergraph(len(nodes), len(edges), hyperedge_indices,
                                hyperedges, k, edge_weights, []), nodes, edges
        return hp

    def _refresh(self, graph):
        return graph[:, numpy.sum(graph, axis=0) > 0]
