import logging

from KafNafParserPy import Cspan, Ctarget, Ccoreference


logger = logging.getLogger(None if __name__ == '__main__' else __name__)


def add_coreference_to_naf(nafobj, corefclasses, mentions):

    start_count = get_starting_count(nafobj)
    coref_according_to_index = get_ordered_coreference_chains(
        corefclasses,
        mentions
    )
    if logger.getEffectiveLevel() <= logging.DEBUG:
        from .util import view_coref_classes
        from collections import OrderedDict
        logger.debug("Coreference classes: {}".format(
            view_coref_classes(
                nafobj,
                mentions,
                OrderedDict(sorted(coref_according_to_index.items()))
            )
        ))

    for key in sorted(coref_according_to_index):
        mids = coref_according_to_index.get(key)
        if mids is not None:
            mids = set(mids)
            nafCoref = Ccoreference()
            cid = 'co' + str(start_count)
            start_count += 1
            nafCoref.set_id(cid)
            nafCoref.set_type('entity')
            data = sorted(
                get_terms_from_offsets(
                    nafobj,
                    mention.span,
                    mention.head_offset
                )
                for mention in map(mentions.get, mids)
            )
            for term_id_span, head_id in data:
                coref_span = create_span(term_id_span, head_id)
                nafCoref.add_span_object(coref_span)
            nafobj.add_coreference(nafCoref)


def create_span(term_id_span, head_id):
    '''
    Creates naf span object where head id is set
    :param term_id_span: list of term ids
    :param head_id: identifier for the head id
    :return: naf span object
    '''

    mySpan = Cspan()
    for term in term_id_span:
        if term == head_id:
            myTarget = Ctarget()
            myTarget.set_id(term)
            myTarget.set_as_head()
            mySpan.add_target(myTarget)
        else:
            mySpan.add_target_id(term)
    return mySpan


def get_ordered_coreference_chains(corefclasses, mentions):
    '''
    Function that creates new coreference dictionary with index of first
    mention as key (for ordering)

    :param corefclasses: identified coreference classes
    :return: dictionary of coreference classes with new indeces
    '''

    coref_dict = {}

    for mids in corefclasses.values():
        first_index = 90000000000
        for mid in mids:
            mention = mentions[mid]
            if mention.head_offset < first_index:
                first_index = mention.head_offset
        coref_dict[first_index] = mids

    return coref_dict


def get_starting_count(nafobj):

    coref_counter = 1
    for coref in nafobj.get_corefs():
        coref_counter += 1

    return coref_counter


def get_terms_from_offsets(nafobj, offset_span, head_offset=-1):

    wids = []
    head_wid = ''
    for token in nafobj.get_tokens():
        if int(token.get_offset()) in offset_span:
            wids.append(token.get_id())
            if int(token.get_offset()) == head_offset:
                head_wid = token.get_id()

    tids = []
    head_tid = ''
    for term in nafobj.get_terms():
        if term.get_span().get_span_ids()[0] in wids:
            tids.append(term.get_id())
        if head_wid in term.get_span().get_span_ids():
            head_tid = term.get_id()

    return tids, head_tid

@@ -271,21 +271,22 @@
 ding_to_
-index
+offset
  = get_o
@@ -361,479 +361,51 @@
   )%0A
-    if logger.getEffectiveLevel() %3C= logging.DEBUG:%0A        from .util import view_coref_classes%0A        from collections import OrderedDict%0A        logger.debug(%22Coreference classes: %7B%7D%22.format(%0A            view_coref_classes(%0A                nafobj,%0A                mentions,%0A                OrderedDict(sorted(coref_according_to_index.items()))%0A            )%0A        ))%0A%0A    for key in sorted(coref_according_to_index):%0A        mids = coref_according_to_index.get(key)
+%0A    for mids in coref_according_to_offset:
 %0A   
@@ -2017,19 +2017,20 @@
 x = 
-90000000000
+float('inf')
 %0A   
@@ -2234,24 +2234,25 @@
 %0A    return 
+%5B
 coref_dict%0A%0A
@@ -2249,16 +2249,48 @@
 ref_dict
+%5Bk%5D for k in sorted(coref_dict)%5D
 %0A%0A%0Adef g

Clean code in dump.py