from pydmrs.core import SortDictDmrs, span_pred_key, abstractSortDictDmrs
from pydmrs.matching.common import are_equal_links


# ------------------------------------------------------------------------
def match_nodes(nodes1, nodes2, excluded=[]):
    """
    :param nodes1: A list of Nodes from the DMRS to be matched, sorted by span_pred_key.
    :param nodes2: A list of Nodes from the DMRS against which we match, sorted by span_pred_key.
    :param excluded: A list of nodeids which should not be used for matching.

    :return: A list of lists of nodeid pairs. The first element in the pair is from small DMRS, the second from the
    larger one. The pairs are listed in reverse span_pred_key order of the corresponding nodes.  Returns [] if no
    match found.
    """
    if not nodes1 or not nodes2:
        return []
    matches = []
    earliest = len(nodes1)
    longest = 0
    for i, node2 in enumerate(nodes2):
        if len(nodes2) - i < longest:  # Not enough nodes left to beat the current longest match.
            break
        if excluded and node2.nodeid in excluded:
            continue
        for j, node1 in enumerate(nodes1):
            if j > earliest:  # To avoid repetition.
                break
            if node1 == node2:
                best_matches = match_nodes(nodes1[j + 1:], nodes2[i + 1:], excluded=excluded)
                if best_matches:
                    for match in best_matches:
                        match.append((node1.nodeid, node2.nodeid))
                else:
                    best_matches = [[(node1.nodeid, node2.nodeid)]]
                earliest = j
                longest = max(longest, len(best_matches[0]))
                matches.extend(best_matches)
    if matches:
        max_len = len(max(matches, key=len))
        return [m for m in matches if len(m) == max_len]
    else:
        return []


def add_quantifier_matches(dmrs1, dmrs2, longest_matches):
    for m in longest_matches:
        q_pairs = []
        for nodeid1, nodeid2 in m:
            try:
                q_link1 = dmrs1.get_in(nodeid1, rargname='RSTR', post='H').pop()
                q_link2 = dmrs2.get_in(nodeid2, rargname='RSTR', post='H').pop()
            except KeyError:
                continue
            if dmrs1[q_link1.start] == dmrs2[q_link2.start]:
                q_pairs.append((q_link1.start, q_link2.start))
        m.extend(q_pairs)


def get_compounds(dmrs, compound_preds):
    compounds = []
    for node in dmrs.iter_nodes():
        if str(node.pred) in compound_preds:
            arg1 = dmrs.get_out_nodes(node.nodeid, rargname='ARG1').pop().nodeid
            arg2 = dmrs.get_out_nodes(node.nodeid, rargname='ARG2').pop().nodeid
            compounds.append({"node": node, "args": (arg1, arg2)})
    return compounds


def add_compound_matches(small_dmrs, large_dmrs, longest_matches, compound_preds):
    small_compounds = get_compounds(small_dmrs, compound_preds)
    large_compounds = get_compounds(large_dmrs, compound_preds)

    for m in longest_matches:
        cmpd_pairs = []
        for small_cmpd in small_compounds:
            query_arg1 = None
            query_arg2 = None
            for small, large in m:
                if small == small_cmpd['args'][0]:
                    query_arg1 = large
                elif small == small_cmpd['args'][1]:
                    query_arg2 = large
                if query_arg1 and query_arg2:
                    break
            else:
                continue
            for large_cmpd in large_compounds:
                if (query_arg1, query_arg2) == large_cmpd['args']:
                    if small_cmpd['node'] == large_cmpd['node']:
                        cmpd_pairs.append((small_cmpd['node'].nodeid, large_cmpd['node'].nodeid))
        m.extend(cmpd_pairs)


def find_extra_surface_nodeids(nodeids, large_dmrs):
    """ Finds nodeids present in the aligned matched region of the large DMRS,
        but which have no equivalents in the small DMRS.

        :param nodeids Nodeids from the large DMRS which have equivalents in the small one, sorted by span_pred_key of
        their nodes.
        :param large_dmrs The large DMRS.

        :return A list of additional nodeids sharing the span with nodeids but without equivalents in the small DMRS.
    """
    max_cto = large_dmrs[nodeids[-1]].cto
    extra_nodeids = []
    reached_start = False
    reached_end = False
    for i, node in enumerate(large_dmrs.nodes):
        if node.nodeid == nodeids[0]:
            first_overlap_orderid = i
            min_cfrom = node.cfrom
            max_cto = max(max_cto, node.cto)
            while True and first_overlap_orderid > 0:
                prev_node = large_dmrs.nodes[first_overlap_orderid - 1]
                prev_cfrom = prev_node.cfrom
                if prev_cfrom == min_cfrom and prev_node.cto <= max_cto:
                    first_overlap_orderid -= 1
                    extra_nodeids.append(prev_node.nodeid)
                    max_cto = max(max_cto, prev_node.cto)
                else:
                    break
            reached_start = True
        elif not reached_start:
            continue
        elif reached_end and node.cfrom >= max_cto:
            break
        else:
            max_cto = max(max_cto, node.cto)
            if node.nodeid not in nodeids and node.nodeid not in extra_nodeids:
                extra_nodeids.append(node.nodeid)
        if node.nodeid == nodeids[-1]:
            reached_end = True

    return extra_nodeids


def get_links(dmrs, nodeids):
    """
    :param dmrs: A Dmrs object.
    :param nodeids: A list of nodeids.
    :return: A list of all links starting and ending on a node from nodeids.
    """
    links = []
    eq_links = set()
    for nodeid in nodeids:
        node_links = dmrs.get_out(nodeid)
        for link in node_links:
            if link.end in nodeids:
                links.append(link)
        node_links = dmrs.get_eq(nodeid)
        for link in node_links:
            if link not in eq_links:
                eq_links.add(link)
    links.extend(eq_links)
    return links


def get_subgraph(dmrs, subgraph_nodeids):
    """ Returns a subgraph of dmrs containing only nodes with subgraph_nodeids and all the links between them.
    :param dmrs: A Dmrs object.
    :param subgraph_nodeids: A list of nodeids.
    :return A SortDictDmrs containing only nodes with subgraph_nodeids and links between them.