import copy
import numpy as np
import opt_einsum
from acqdp.tensor_network import ContractionTask, defaultOrderResolver, ContractionCost


class Compiler:
    """
    Compile a :class:`ContractionScheme` indicating contraction scheme and sliced edges into a :class:`ContractionTask`, a hardware-aware format that can be readily used for tensor network contraction.

    Typically, a :class:`ContractionScheme` found by one of the :class:`OrderFinder` only aims to minimize the theoretical floating number operations and / or largest size of intermediate tensors.
    Although those are typically accurate indicators of the overall contracion cost, they fail to accomodate inefficiencies from elsewhere.
    The compiler takes care of machine-related inefficiencies by slightly modifying the order.

    :ivar do_patch: If set to false, the patching and reorganization of the contraction order below will be skipped.
    :ivar patch_size: Whenever two adjacent branches both have size less than the patch size, the branches are merged together. This is to avoid inefficiency called by skewed tensor shapes in tensor multiplication, with a slight sacrifice in floating number operations.
    :ivar reorg_thres: Threshold for contraction order reorganization. The order is seperated into small tensor multiplications and large tensor multiplications. All pairwise tensor multiplications involving tensors with size less than `reorg_thres` will be put forward whenever possible.
    """
    def __init__(self,
                 reorg_thres=23,
                 patch_size=5,
                 do_patch=False,
                 **kwargs):
        self.reorg_thres = reorg_thres
        self.patch_size = patch_size
        self.do_patch = do_patch

    def compile(self, tn, scheme, **kwargs):
        """
        :param tn: The :class:`TensorNetwork` the input contraction scheme is based on.
        :type tn: :class:`TensorNetwork`
        :param scheme: The contraction scheme from which the task is generated.
        :type scheme: :class:`ContractionScheme`

        :returns: :class:`ContractionTask` -- The compiled contraction task ready for execution.
        """
        tn = tn._expand_and_delta()
        res = self._generate_template(tn, scheme=scheme, **kwargs)
        res.set_data({
            node: tn.network.nodes[(0, node)]['tensor'].contract()
            for node in tn.nodes_by_name
        })
        res.update_fix(tn)
        return res

    def _generate_template(self, tn, scheme, test=0, **kwargs):
        inputs = {node[1]: [None] for node in tn.nodes}
        data_ref = {(0, node): inputs[node] for node in inputs}
        order, split = scheme.order, scheme.slices
        split = [(1, i) for i in split]
        tn_copy = tn.copy().expand(recursive=True)
        tn_copy.open_edges += [i[1] for i in split]
        tn_copy.fix()
        for edge_name in list(tn_copy.edges_by_name):
            if (len(tn_copy.network[(1, edge_name)])
                    == 0) and edge_name not in tn_copy.open_edges:
                tn_copy.network.remove_node((1, edge_name))
        init_order, final_order = self._order_patch(tn_copy.copy(), order, split,
                                                   **kwargs)
        k = self._generate_template_fix(tn, [], data_ref)
        if isinstance(k, ContractionTask):
            return k
        lst, fix_dict = k

        for item in init_order:
            if item[1] == '#':
                break
            stn_name, stn = tn_copy.encapsulate(item[0], stn_name=item[1])
            lst += self._generate_template_basic(stn, data_ref, (0, stn_name))

        cnt = len(lst)

        fix_lst, fix_dict_new = self._generate_template_fix(
            tn_copy, split, data_ref)
        lst += fix_lst
        fix_dict.update(fix_dict_new)

        tn_copy.open_edges = [
            i for i in tn_copy.open_edges if (1, i) not in split
            ]
        for i in split:
            tn_copy.fix_edge(i[1])
        tn_copy.fix()

        lst += self._generate_template_basic(tn_copy, data_ref, '#')
        _, _, shapes = tn_copy.subscripts()
        if len(final_order) > 0:
            path = []
            lst_nodes = list(n[1] for n in tn_copy.nodes)
            for o in final_order:
                path.append(
                    (lst_nodes.index(o[0][0]), lst_nodes.index(o[0][1])))
                lst_nodes.remove(o[0][0])
                lst_nodes.remove(o[0][1])
                lst_nodes.append(o[1])
            expr = opt_einsum.contract_expression(lst[-1][3]['subscripts'],
                                                  *shapes,
                                                  optimize=path)
            lst[-1][3]['expr'] = expr
        res = ContractionTask(commands=lst,
                               fix_dict=fix_dict,
                               inputs=inputs,
                               output=data_ref['#'],
                               shape=tn.shape,
                               open_edges=tn.open_edges,
                               sub_outputs=tn_copy.open_edges,
                               cnt=cnt)
        return res

    def _generate_template_fix(self, tn, split, data_ref):
        lst = []

        fix_dict = self._generate_template_edge_fix(tn, split)
        if not isinstance(fix_dict, dict):
            return fix_dict
        for node in tn.nodes:
            fix_idx = []
            fix_to = []
            for i, edge in enumerate(tn.network.nodes[node]['edges']):
                if (1, edge) in fix_dict:
                    fix_idx.append(i)
                    fix_to.append(fix_dict[(1, edge)][1])
            if len(fix_idx) > 0:
                new_res = [None]
                lst.append(('f', data_ref[node], new_res, {'fix_idx': fix_idx, 'fix_to': fix_to}))
                data_ref[node] = new_res
        res_dict = {(k, tuple(fix_dict[k][0])): fix_dict[k][1] for k in fix_dict}
        return lst, res_dict

    def _generate_template_edge_fix(self, tn, split):
        fix_dict = {}
        for edge in tn.edges:
            if edge in split:
                tn.network.nodes[edge]['fix_to'] = list(range(tn.network.nodes[edge]['dim']))
                fix_dict[edge] = (list(range(tn.network.nodes[edge]['dim'])), [0])
            elif 'fix_to' in tn.network.nodes[edge]:
                ll = len(tn.network.nodes[edge]['fix_to'])
                if ll == 0:
                    return ContractionTask(output=[(0, np.zeros(tn.shape))])
                else:
                    fix_dict[edge] = (list(tn.network.nodes[edge]['fix_to']), list(tn.network.nodes[edge]['fix_to']))
        return fix_dict

    def _generate_template_delta(self, tn, data_ref):
        open_edge_names = {}
        edges = tn.edges_by_name_from_nodes_by_name()
        if len(tn.nodes) == 0:
            id_node = np.array(1.)
            tn.add_node('*', [], id_node)
            data_ref[(0,'*')] = [(0, id_node)]
        for (i, edge) in enumerate(tn.open_edges):
            if edge not in open_edge_names:
                open_edge_names.update({edge: 0})
            else:
                open_edge_names[edge] += 1
                new_node = np.eye(tn.network.nodes[(1, edge)]['dim'])
                data_ref[(0, ("*", edge, open_edge_names[edge]))] = [(0, new_node)]
                tn.add_node(("*", edge, open_edge_names[edge]), [edge, (edge, open_edge_names[edge])], new_node)
                tn.open_edges[i] = (edge, open_edge_names[edge])
        for edge in open_edge_names:
            diag = None
            if edge not in edges and open_edge_names[edge] == 0:
                diag = [1.] * tn.network.nodes[(1, edge)]['dim']
            if 'fix_to' in tn.network.nodes[(1, edge)]:
                diag = [1. if i in tn.network.nodes[(1, edge)]['fix_to'] else 0 for i in range(tn.network.nodes[(1, edge)]['dim'])]
                tn.network.nodes[(1, edge)].pop('fix_to')
            if diag is not None:
                new_node = np.array(diag)
                data_ref[(0, ('*', edge, 0))] = [(0, new_node)]
                tn.add_node(('*', edge, 0), [edge], new_node)

    def _generate_template_basic(self, tn, data_ref, new_name, **kwargs):
        new_res = [None]
        data_ref[new_name] = new_res
        if new_name == '#':
            self._generate_template_delta(tn, data_ref)
        lhs = [data_ref[node] for node in tn.nodes]

        return [('c', lhs, new_res, self._generate_template_einsum(tn, subscripts=tn.subscripts(), **kwargs), [node for node in tn.nodes])]

    def _generate_template_einsum(self, tn, subscripts, **kwargs):
        lhs_str, rhs_str, _ = subscripts
        command = ','.join(lhs_str) + '->' + rhs_str
        return {'subscripts': command,  'dtype': tn.dtype}

    def _order_patch(self, tn, order, split_edges, **kwargs):
        if self.do_patch:
            from acqdp.tensor_network.contraction_tree import ContractionTree
            tn_copy = tn.copy()
            tn_copy.open_edges += [i[1] for i in split_edges]
            tn_copy.fix()
            for edge_name in list(tn_copy.edges_by_name):
                if (len(tn_copy.network[(1, edge_name)])
                        == 0) and edge_name not in tn_copy.open_edges:
                    tn_copy.network.remove_node((1, edge_name))
            assert tn_copy.dtype == tn.dtype
            order = ContractionTree.from_order(tn_copy, order).full_order
            init_order, final_order = self._reorg_order(tn_copy, order)
            tn_copy.open_edges = [
                i for i in tn_copy.open_edges if (1, i) not in split_edges
            ]
            for i in split_edges:
                tn_copy.fix_edge(i[1])
            tn_copy.fix()
            if len(final_order) > 0:
                final_order = self._patch_order(tn_copy, final_order)
            return init_order, final_order
        else:
            return [], order

    def _reorg_order(self, graph, order):
        graph.fix()
        size_dic = {}
        subs, order = defaultOrderResolver.order_to_subscripts_tree(
            graph, order)
        assert len(subs) == len(order), f'{len(subs)}, {len(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])
        init_order = []
        node_set = set([a[1] for a in graph.nodes])
        while True:
            for o in order:
                if size_dic[o[1]] < self.reorg_thres and np.all(
                    [a in node_set for a in o[0]]):
                    init_order.append(o)
                    order.remove(o)
                    for a in o[0]:
                        node_set.remove(a)
                    node_set.add(o[1])
                    graph.encapsulate(nodes=o[0], stn_name=o[1])
                    break
            else:
                break
        return init_order, order

    def _patch_order(self, graph, order):
        from acqdp.tensor_network.contraction_tree import ContractionTree
        graph_copy = graph.copy()
        graph_copy.fix()
        tree = ContractionTree.from_order(graph_copy, order)

        res = tree.patch(self.patch_size)
        return res
