from acqdp.tensor_network.contraction import OrderCounter, ContractionCost
import opt_einsum
import copy
import numpy as np

ORDER_RESOLVER_KWARGS = {'optimal': 8, 'dp': 20, 'thres': 1e8}


class OrderResolver:
    """

    Interfaces for conveniently converting between different formats of contraction orders. Also serves as a converter from non-pairwise contraction orders to pairwise contraction orders.

    :ivar optimal: The maximum number of tensors in a single step of a non-pairwise order, for which brute-force search is used to reconstruct the pairwise order.
    :ivar dp: The maximum number of tensors in a single step of a non-pairwise order, for which dynamic programming is used to reconstruct the pairwise order.
    :ivar thres: The minimum cost of a non-pairwise contraction order via default path to invoke more advanced search for a pairwise order, such as the `dp` or `optimal` methods.

    """

    def __init__(self,
                 optimal=8,
                 dp=20,
                 thres=1e8,
                 **kwargs):
        self.optimal = optimal
        self.dp = dp
        self.thres = thres

    def order_to_contraction_scheme(self, tn, order):
        subscripts = self.order_to_subscripts(tn, order)
        cost = ContractionCost()
        counter = OrderCounter(order)
        new_order = []
        for o, i in zip(order, subscripts):
            try:
                res = self.subscript_to_path(*i)
                new_order += self.path_to_paired_order(o, res[0], counter)
                cost += res[1]
            except IndexError as e:
                print(o, i, res[0], res[1])
                raise e
        from acqdp.tensor_network import ContractionScheme
        return ContractionScheme(new_order, cost=cost)

    def order_to_path(self, tn, order):
        tn_copy = tn.copy()
        tn_copy.fix()
        edges_list = tn_copy.edges_by_name
        nodes_list = tn_copy.nodes_by_name
        lhs = [
            ''.join(
                opt_einsum.get_symbol(edges_list.index(e))
                for e in tn_copy.network.nodes[(0, n)]['edges'])
            for n in nodes_list
        ]
        rhs = ''.join(
            opt_einsum.get_symbol(edges_list.index(e))
            for e in tn_copy.open_edges)
        path = []
        for o in order:
            i, j = nodes_list.index(o[0][0]), nodes_list.index(o[0][1])
            path.append((i, j))
            nodes_list.pop(max(i, j))
            nodes_list.pop(min(i, j))
            nodes_list.append(o[1])
        return ','.join(lhs) + '->' + rhs, path, {
            opt_einsum.get_symbol(edges_list.index(e)): e for e in edges_list
        }

    def order_to_subscripts_tree(self, tn, order):
        from acqdp.tensor_network.contraction_tree import ContractionTree
        tree = ContractionTree.from_order(tn, order)
        return tree.full_subscripts, tree.full_order

    def order_to_subscripts(self, tn, order):
        tn_copy = tn.copy()
        subs = []
        for o in order:
            if o[1] == '#':
                break
            try:
                _, stn = tn_copy.encapsulate(o[0], stn_name=o[1])
                subs.append(stn.subscripts(o[0]))
            except Exception as e:
                print(e)
                raise e
        if len(order) > 0:
            subs.append(tn_copy.subscripts(order[-1][0]))
        return subs

    def print_order(self, tn, order):
        subs, order = self.order_to_subscripts_tree(tn, order)
        steps = []
        for i, sub in enumerate(subs):
            a = len(set(sub[0][0]).difference(sub[0][1]))
            b = len(set(sub[0][1]).difference(sub[0][0]))
            c = len(set(sub[0][1]).intersection(sub[0][0]).difference(sub[1]))
            d = len(set(sub[0][1]).intersection(sub[0][0]).intersection(sub[1]))
            steps.append(
                f'Step {i}: <{a+b+c+d}> {sub[0][0]}[{a} x {c} x {d}] * {sub[0][1]}[{c} x {b} x {d}] -> {sub[1]}[{a} x {b} x {d}] \n {order[i]}'
            )
        return '\n'.join(steps)

    def subscript_to_path(self, lhs, rhs, shapes):
        optimize = None
        try:
            path, path_info = opt_einsum.contract_path(','.join(lhs) + '->' + rhs,
                                                    *shapes,
                                                    shapes=True,
                                                    optimize='auto')
        except ValueError as e:
            print(','.join(lhs) + '->' + rhs, shapes)
            raise e
        if len(lhs) > 2 and path_info.opt_cost > self.thres:
            if len(lhs) < self.optimal:
                optimize = 'optimal'
            elif len(lhs) < self.dp:
                optimize = 'dp'
        if optimize is not None:
            path, path_info = opt_einsum.contract_path(','.join(lhs) + '->' + rhs,
                                                    *shapes,
                                                    shapes=True,
                                                    optimize=optimize)
        return path, ContractionCost(path_info.opt_cost,
                                     path_info.largest_intermediate)

    def path_to_paired_order(self, order, path, counter):
        lhs, rhs = copy.copy(order[0]), order[1]
        # print(lhs, rhs)
        new_order = []
        for i in path:
            try:
                if len(i) == 1:
                    return [[[lhs.pop(i[0])], rhs]]
                else:
                    new_order.append([[lhs[i[0]], lhs[i[1]]], counter.cnt])
                    lhs.pop(max(i[0], i[1]))
                    lhs.pop(min(i[0], i[1]))
                    lhs.append(new_order[-1][1])
            except IndexError as e:
                print(order[0], i)
                raise e
        new_order[-1][1] = rhs
        return new_order


defaultOrderResolver = OrderResolver()


class LocalOptimizer:
    """

    :class:`LocalOptimizer` takes a pairwise contraction tree and do local optimization on the tree. Note that a connected subgraph of the contraction tree represents a sequence of intermediate steps that take some (potentially intermediate) tensors as input, and outputs a later intermediate tensor. This sequence of steps can be optimized based on solely the shapes of the input and output tensors associated to the subgraph, without looking at the rest of the contraction tree. This allows local reorganization of the tree.

    Each iteration of local optimization consists of two steps: In the first step, the contraction tree is divided into non-overlapping subgraphs each up to a given size, with the internal connections in the subgraphs removed. The internal connections are then reconstructed in the second phase using accelerated optimum contraction tree finding approach. The iterations can be repeated by each time randomly selecting a different patch of subgraph division.

    :ivar size: The maximum number of nodes to be included in one subgraph.
    :ivar resolver: The resolver to help reconstruct the pairwise order from subgraph divisions.

    """

    def __init__(self,
                 size=13,
                 **kwargs):
        self.size = size
        self.resolver = OrderResolver(**kwargs.get('resolver_params', {}))

    def _flatten_order(self, tn, order, offset=0):
        tn_copy = tn.copy()
        tn_copy.fix()
        size_dic = {}
        cnt_offset = 0
        subs = self.resolver.order_to_subscripts(tn_copy, order)
        for i, o in enumerate(order):
            size_dic[o[1]] = len(subs[i][1])
            for j, k in enumerate(o[0]):
                size_dic[k] = len(subs[i][0][j])
        cnt = -1
        while -cnt < len(order):
            rhs_list = {}
            for i, o in enumerate(order[:cnt]):
                rhs_list[o[1]] = (i, o[0])
            o = order[cnt]
            lhs = o[0]
            if len(lhs) >= self.size:
                cnt -= 1
                continue
            # choose the largest tensor
            sizes = {k: size_dic[k] for k in o[0]}
            largest_list = sorted(sizes, key=lambda x: sizes[x])[::-1]
            for l in largest_list:
                if (l not in rhs_list) or (sizes[l] < 7) or (
                        len(rhs_list[l][1]) + len(lhs) > self.size):
                    continue
                if cnt_offset < offset:
                    cnt_offset += 1
                    continue
                else:
                    o[0].remove(l)
                    o[0] += rhs_list[l][1]
                    order.pop(rhs_list[l][0])
                    break
            else:
                cnt -= 1
        return order

    def optimize(self, tn, order, num_iter):
        for _ in range(num_iter):
            new_order = self._flatten_order(tn,
                                           order.order,
                                           offset=np.random.randint(
                                               self.size))
            order = self.resolver.order_to_contraction_scheme(tn, new_order)
        return order
