# coding: utf-8

from allennlp.data.dataset_readers.dataset_utils.span_utils import bioul_tags_to_spans
import copy


def find_indices(lst, condition):
    """
    Find the indices of elements in a list `lst` that match a condition.
    """
    return [i for i, elem in enumerate(lst) if condition(elem)]


def expand_arg(relation, arg_nb, n_c, formatted_relation):
    """
    Update formatted_relation to take into account the possible expansion of relation arguments.
    """
    arg_str = "arg{}".format(arg_nb)
    arg = (relation[arg_str + "_begin_char"], relation[arg_str + "_end_char"], relation[arg_str + "_text"])
    if arg in n_c.keys():
        formatted_relation[arg_str + "_begin_char"] = n_c[arg].start_char
        formatted_relation[arg_str + "_end_char"] = n_c[arg].end_char

        formatted_relation[arg_str + "_begin_token"] = n_c[arg].start
        formatted_relation[arg_str + "_end_token"] = n_c[arg].end - 1

        formatted_relation[arg_str + "_text"] = n_c[arg].text
    else:
        formatted_relation[arg_str + "_begin_token"] = relation.get(arg_str + "_index")
        formatted_relation[arg_str + "_end_token"] = relation.get(arg_str + "_index")


def check_overlapping(relation, formatted_relation):
    """
    Check if there is no overlapping between the two expanded arguments of a relation.
    If there is an overlap, we drop the expansion for the relation.
    """
    arg1_b, arg1_e = formatted_relation["arg1_begin_char"], formatted_relation["arg1_end_char"]
    arg2_b, arg2_e = formatted_relation["arg2_begin_char"], formatted_relation["arg2_end_char"]

    overlap = False or (arg1_b < arg2_e and arg1_e >= arg2_b) or (arg2_b < arg1_e and arg2_e >= arg1_b)

    if overlap:
        arg1_i, arg2_i = relation["arg1_index"], relation["arg2_index"]
        relation["arg1_begin_token"], relation["arg1_end_token"] = arg1_i, arg1_i
        relation["arg2_begin_token"], relation["arg2_end_token"] = arg2_i, arg2_i
        return relation
    else:
        return formatted_relation


class predictionFormatter:
    """
    A class that format the prediction returned by HMTL model.
    If necessary, it also expands
    """

    def __init__(self):
        pass

    def format(self, predictions, sent_char_offset, input_text):
        tokenized_text = predictions["tokenized_text"]
        predicted_tasks = predictions.keys()

        formatted_predictions = {}
        formatted_predictions["tokenized_text"] = tokenized_text

        ### Format NER and EMD ###
        for task_name in ["ner", "emd"]:
            if task_name in predicted_tasks:
                decoded_bioul = []
                assert len(predictions[task_name]) == 1

                spans = bioul_tags_to_spans(predictions[task_name][0])
                for tag, (begin, end) in spans:
                    entity = {
                        "type": tag,
                        "begin_token": begin,
                        "end_token": end,
                        "begin_char": sent_char_offset[begin],
                        "end_char": sent_char_offset[end] + len(tokenized_text[end]),
                        "tokenized_text": tokenized_text[begin : (end + 1)],
                        "text": input_text[
                            sent_char_offset[begin] : (sent_char_offset[end] + len(tokenized_text[end]))
                        ],
                    }
                    decoded_bioul.append(entity)

                formatted_predictions[task_name] = decoded_bioul

        ### Format Relation ###
        if "relation" in predicted_tasks:
            decoded_relation_arcs = []
            assert len(predictions["relation"]) == 1

            for i, relation in enumerate(predictions["relation"][0]):
                indices = find_indices(relation, lambda x: x != "*")
                for ind in indices:
                    tag = relation[ind]
                    if tag[:4] == "ARG1":
                        arg1_index, arg1_text = ind, tokenized_text[ind]
                    if tag[:4] == "ARG2":
                        arg2_index, arg2_text = ind, tokenized_text[ind]
                rel = {
                    "type": tag[5:],
                    "arg1_index": arg1_index,
                    "arg1_text": arg1_text,
                    "arg1_begin_char": sent_char_offset[arg1_index],
                    "arg1_end_char": sent_char_offset[arg1_index] + len(arg1_text),
                    "arg2_index": arg2_index,
                    "arg2_text": arg2_text,
                    "arg2_begin_char": sent_char_offset[arg2_index],
                    "arg2_end_char": sent_char_offset[arg2_index] + len(arg2_text),
                }
                decoded_relation_arcs.append(rel)

            formatted_predictions["relation_arcs"] = decoded_relation_arcs

        ### Format Coreference ###
        if "coref" in predicted_tasks:
            decoded_coref_arcs = []
            decoded_coref_clusters = []
            assert len(predictions["coref"]) == 1

            for cluster in predictions["coref"][0]:
                ## Format the clusters
                decoded_cluster = []
                for mention in cluster:
                    begin, end = mention
                    m = {
                        "begin": begin,
                        "end": end,
                        "begin_char": sent_char_offset[begin],
                        "end_char": sent_char_offset[end] + len(tokenized_text[end]),
                        "tokenized_text": tokenized_text[begin : (end + 1)],
                        "text": input_text[
                            sent_char_offset[begin] : (sent_char_offset[end] + len(tokenized_text[end]))
                        ],
                    }
                    decoded_cluster.append(m)
                decoded_coref_clusters.append(decoded_cluster)

                ## Format the arcs
                for i in range(len(cluster) - 1):
                    mention1_begin, mention1_end = cluster[i]
                    mention2_begin, mention2_end = cluster[i + 1]
                    coref_arc = {
                        "mention1_begin": mention1_begin,
                        "mention1_end": mention1_end,
                        "mention1_begin_char": sent_char_offset[mention1_begin],
                        "mention1_end_char": sent_char_offset[mention1_end] + len(tokenized_text[mention1_end]),
                        "tokenized_text1": tokenized_text[mention1_begin : (mention1_end + 1)],
                        "text1": input_text[
                            sent_char_offset[mention1_begin] : (
                                sent_char_offset[mention1_end] + len(tokenized_text[mention1_end])
                            )
                        ],
                        "mention2_begin": mention2_begin,
                        "mention2_end": mention2_end,
                        "mention2_begin_char": sent_char_offset[mention2_begin],
                        "mention2_end_char": sent_char_offset[mention2_end] + len(tokenized_text[mention2_end]),
                        "tokenized_text2": tokenized_text[mention2_begin : (mention2_end + 1)],
                        "text2": input_text[
                            sent_char_offset[mention2_begin] : (
                                sent_char_offset[mention2_end] + len(tokenized_text[mention2_end])
                            )
                        ],
                    }
                    decoded_coref_arcs.append(coref_arc)

            formatted_predictions["coref_arcs"] = decoded_coref_arcs
            formatted_predictions["coref_clusters"] = decoded_coref_clusters

        return formatted_predictions

    def expand_relations(self, predictions, doc):
        """
        HMTL predicts the relation between the last head tokens.
        This is a simple heuristic to expand relations using a dependecy tree.
        """
        if "relation_arcs" in predictions:
            predictions["relation_arcs_expanded"] = []
            noun_chunks = {}
            for chunk in doc.noun_chunks:
                noun_chunks[(chunk.root.idx, chunk.root.idx + len(chunk.root.text), chunk.root.text)] = chunk

            for relation in predictions["relation_arcs"]:
                formatted_relation = copy.deepcopy(relation)

                expand_arg(relation, 1, noun_chunks, formatted_relation)
                expand_arg(relation, 2, noun_chunks, formatted_relation)
                formatted_relation = check_overlapping(relation, formatted_relation)

                del formatted_relation["arg1_index"], formatted_relation["arg2_index"]

                predictions["relation_arcs_expanded"].append(formatted_relation)

        return predictions

    def expand_emd(self, predictions, doc):
        """
        HMTL predicts the heads of a mention.
        Simple heuristic to expand entity mentions using a dependecy tree.
        """
        if "emd" in predictions:
            noun_chunks = {}
            for chunk in doc.noun_chunks:
                noun_chunks[(chunk.root.idx, chunk.root.idx + len(chunk.root.text), chunk.root.text)] = chunk

            predictions["emd_expanded"] = []
            for emd in predictions["emd"]:
                expanded_emd = copy.deepcopy(emd)
                id_ = (emd["begin_char"], emd["end_char"], emd["text"])

                if id_ in noun_chunks.keys():
                    expanded_emd["begin_char"] = noun_chunks[id_].start_char
                    expanded_emd["end_char"] = noun_chunks[id_].end_char

                    expanded_emd["begin_token"] = noun_chunks[id_].start
                    expanded_emd["end_token"] = noun_chunks[id_].end - 1

                    expanded_emd["text"] = noun_chunks[id_].text
                    expanded_emd["tokenized_text"] = [token.text for token in noun_chunks[id_]]

                predictions["emd_expanded"].append(expanded_emd)

        return predictions

    def expand(self, predictions, doc):
        """
        Perform both EMD and Relation expansion
        """
        predictions = self.expand_relations(predictions, doc)
        predictions = self.expand_emd(predictions, doc)

        return predictions
