Marcell / collector

import collections
import importlib
import os

from natsort import natsorted

from projects.marcell.models import IATETerm, EuroVocTerm


DEFAULT_COLS = ['ID', 'FORM', 'LEMMA', 'UPOS', 'XPOS', 'FEATS', 'HEAD', 'DEPREL', 'DEPS', 'MISC']


def write(document):
    if document.annotated() and document.chunks.exists():
        print('Writing %s in CONLLU format.' % document.id)

        cols = importlib.import_module('projects.%s.mappings' % document.pipeline.project.name).CONLLU_COLS

        text_path = os.path.join(document.path, 'text.conllu')
        if cols != DEFAULT_COLS:
            text_path = os.path.join(document.path, 'text.conllup')

        with open(text_path, 'w', encoding='UTF-8') as text_file:
            _write_header(document, text_file, cols)
            _write_paragraphs(document, text_file, cols)


def write_to_dir(document, export_dir_path):
    if document.annotated() and document.chunks.exists():
        print('Writing %s in CONLLU format.' % document.id)

        cols = importlib.import_module('projects.%s.mappings' % document.pipeline.project.name).CONLLU_COLS

        text_path = os.path.join(export_dir_path, '%s-%s.conllu' % (document.lang, document.id))
        if cols != DEFAULT_COLS:
            text_path = os.path.join(export_dir_path, '%s-%s.conllup' % (document.lang, document.id))

        with open(text_path, 'w', encoding='UTF-8') as text_file:
            _write_header(document, text_file, cols)
            _write_paragraphs(document, text_file, cols)


def _write_header(document, text_file, cols):
    if cols != DEFAULT_COLS:
        text_file.write('# global.columns = %s\n' % ' '.join(cols))

    text_file.write('# newdoc id = %s-%s\n' % (document.lang, document.id))
    _write_metadata(document, text_file)


def _write_metadata(document, text_file):
    project_mappings = importlib.import_module('projects.%s.mappings' % document.pipeline.project.name)

    metadata = {'language': document.lang,
                'date': document.publication_date.strftime('%Y-%m-%d'),
                'title': document.title,
                'type': document.type}

    if document.status:
        metadata['status'] = document.status

    entype = _en(project_mappings.DOC_TYPES, document.type)
    if document.type != entype:
        metadata['entype'] = entype

    if document.keywords.exists():
        metadata['keywords'] = ' | '.join([keyword.label for keyword in document.keywords.all()])

    if document.source_url:
        metadata['url'] = document.source_url

    for meta in document.metadata.order_by('sequence'):
        meta_multivalue, meta_separator = _meta_multivalue(project_mappings.META_TYPES, meta.name)
        translated_name = _en(project_mappings.META_TYPES, meta.name)
        if translated_name is not None:
            if meta_multivalue and not meta_separator:
                if translated_name in metadata:
                    metadata[translated_name].append(meta.value)
                else:
                    metadata[translated_name] = [meta.value]
            else:
                metadata[translated_name] = [meta.value]

    for name, value in metadata.items():
        if type(value) == list:
            text_file.write('# %s = %s\n' % (name, ';'.join(value)))
        else:
            text_file.write('# %s = %s\n' % (name, value))


def _meta_multivalue(meta_mapping, pl_name):
    for meta_type in meta_mapping:
        if meta_type['pl'] == pl_name:
            return meta_type['multivalue'], meta_type['separator']
    return False, None


def _en(translations, pl_name):
    for translation in translations:
        if translation['pl'] == pl_name:
            if translation['en'] is None:
                return pl_name
            else:
                return '%s' % translation['en']
    return None


def _write_paragraphs(document, text_file, cols):
    for di, chunk in enumerate(document.chunks.order_by('sequence'), 1):
        if chunk.utterances.exists():
            for ui, utt in enumerate(chunk.utterances.order_by('sequence')):
                par_id = '%s-%s-u%d.%d' % (document.lang, document.id, di, ui)
                text_file.write('# newpar id = %s\n' % par_id)
                text_file.write('# who = %s\n' % utt.speaker.abbrev)
                _write_sentences(par_id, utt.anno, text_file, cols)
        else:
            par_id = '%s-%s-p%d' % (document.lang, document.id, di)
            text_file.write('# newpar id = %s\n' % par_id)
            _write_sentences(par_id, chunk.anno, text_file, cols)


def _write_sentences(par_id, anno, text_file, cols):
    for si, sent in enumerate(anno['chunks'][0]['sentences'], 1):
        sent_id = '%ss%d' % (par_id, si)
        text_file.write('# sent_id = %s\n' % sent_id)

        text = _sent2text(sent)
        text_file.write('# text = %s\n' % text)

        _write_tokens(sent['tokens'], anno, text_file, cols)

        text_file.write('\n')


def _sent2text(sent):
    text = ''
    for tok in sent['tokens']:
        if tok['ns']:
            text += tok['orth']
        else:
            text += '%s ' % tok['orth']
    return text.strip()


def _write_tokens(tokens, anno, text_file, cols):
    tokens_map = {}
    for ti, tok in enumerate(tokens, 1):
        tokens_map[tok['id']] = str(ti)

    names = []
    if 'names' in anno:
        names = _get_local_nes(tokens, anno['names'])

    nps = _get_local_nps(tokens)

    iate = []
    if 'iate' in anno:
        iate = _get_local_longest_iate_terms(tokens, anno['iate'])

    eurovoc = []
    if 'eurovoc' in anno:
        eurovoc = _get_local_longest_eurovoc_terms(tokens, anno['eurovoc'])

    for tok in tokens:

        row = []
        for col in cols:
            if col == 'ID':
                row.append(tokens_map[tok['id']])
            elif col == 'FORM':
                row.append(tok['orth'])
            elif col == 'LEMMA':
                row.append(tok['base'])
            elif col == 'UPOS':
                row.append(tok['upostag'])
            elif col == 'XPOS':
                row.append(tok['ctag'])
            elif col == 'FEATS':
                row.append(tok['feats'])
            elif col == 'HEAD':
                head = tokens_map[tok['head']] if tok['head'] else '0'
                row.append(head)
            elif col == 'DEPREL':
                row.append(tok['deprel'])
            elif col == 'DEPS':
                row.append(tok['deps'])
            elif col == 'MISC':
                misc = tok['misc']
                if misc == '_' and tok['ns']:
                    misc = 'SpaceAfter=No'
                elif misc != '_' and tok['ns'] and 'SpaceAfter=No' not in misc.split('|'):
                    misc = '|'.join(misc.split('|').append('SpaceAfter=No'))
                row.append(misc)
            elif col.split(':')[1] == 'NE':
                col_val = _get_ne_col_value(tok, names)
                row.append(col_val)
            elif col.split(':')[1] == 'NP':
                col_val = _get_np_col_value(tok, nps)
                row.append(col_val)
            elif col.split(':')[1] == 'IATE':
                col_val = _get_iate_col_value(tok, iate)
                row.append(col_val)
            elif col.split(':')[1] == 'EUROVOC':
                col_val = _get_eurovoc_col_value(tok, eurovoc)
                row.append(col_val)
            elif col.split(':')[1] == 'EUROVOCMT':
                col_val = _get_eurovocmt_col_value(tok, eurovoc)
                row.append(col_val)
            else:
                row.append('_')

        text_file.write('%s\n' % '\t'.join(row))


def _get_local_nes(tokens, nes):
    local_nes = []
    for ne in nes:
        if not _ne_is_continous(ne):
            continue
        if _ne_is_person_subtype(ne):
            continue
        for tok in tokens:
            if tok['id'] in ne['tokens']:
                local_nes.append(ne)
                break

    not_intersecting_nes = _get_not_intersecting_nes(local_nes)

    return not_intersecting_nes


def _ne_is_continous(ne):
    ne_tokens_as_numbers = _tokens_ids_to_numbers(ne['tokens'])
    return sorted(ne_tokens_as_numbers) == list(range(min(ne_tokens_as_numbers), max(ne_tokens_as_numbers) + 1))


def _tokens_ids_to_numbers(tokens_ids):
    return [int(tok_id.lstrip('t')) for tok_id in natsorted(tokens_ids)]


def _ne_is_person_subtype(ne):
    ne_type_parts = ne['type'].lower().split('_')
    if ne_type_parts[0] == 'persname' and len(ne_type_parts) > 1:
        return True
    return False


def _get_not_intersecting_nes(nes):
    not_intersecting_nes = []
    ordered_nes = natsorted(nes, key=lambda ne: natsorted(ne['tokens'])[0])
    for ne1 in ordered_nes:
        intersects = False
        for ne2 in not_intersecting_nes:
            if _nes_intersects(ne1, ne2):
                intersects = True
                break
        if not intersects:
            not_intersecting_nes.append(ne1)
    return not_intersecting_nes


def _nes_intersects(ne1, ne2):
    if any(tok_id in ne1['tokens'] for tok_id in ne2['tokens']):
        return True
    return False


def _get_ne_col_value(tok, nes):
    for ne in nes:
        ne_type = _get_ne_type(ne)
        if tok['id'] in ne['tokens'] and ne_type is not None:
            ordered_ne_tokens = natsorted(ne['tokens'])
            if tok['id'] == ordered_ne_tokens[0]:
                return 'B-%s' % ne_type
            else:
                return 'I-%s' % ne_type
    return 'O'


def _get_ne_type(ne):
    ne_type = None

    ne_type_parts = ne['type'].lower().split('_')
    if ne_type_parts[0] == 'persname' and len(ne_type_parts) == 1:
        ne_type = 'PER'
    elif ne_type_parts[0] == 'orgname':
        ne_type = 'ORG'
    elif ne_type_parts[0] in ['geogname', 'placename']:
        ne_type = 'LOC'
    elif ne_type_parts[0] == 'date':
        ne_type = 'DATE'
    elif ne_type_parts[0] == 'time':
        ne_type = 'TIME'

    return ne_type


def _get_local_nps(tokens):
    nps = []

    id2token_map = _create_id2token_map(tokens)
    head2tokens_map = _create_head2tokens_map(tokens)

    for tok in tokens:
        if _is_np_head(tok):
            nps.append(_tokens_to_ids([tok]))

            np_tokens = []
            _get_subtree(tok, np_tokens, head2tokens_map)
            if np_tokens and _np_is_continous(np_tokens):
                nps.append(_tokens_to_ids(np_tokens))

    widest_nps = []
    if nps:
        for tokens_ids in _remove_subnps(nps):
            np_tokens = _ids_to_tokens(tokens_ids, id2token_map)
            np_tokens = _to_np(np_tokens)
            widest_nps.append({'tokens': _tokens_to_ids(np_tokens)})

    not_intersecting_nps = _get_not_intersecting_nps(widest_nps)

    return not_intersecting_nps


def _create_id2token_map(tokens):
    id2token_map = {}
    for tok in tokens:
        id2token_map[tok['id']] = tok
    return id2token_map


def _create_head2tokens_map(tokens):
    head2tokens_map = {}
    for tok in tokens:
        if tok['head']:
            if tok['head'] in head2tokens_map:
                head2tokens_map[tok['head']].append(tok)
            else:
                head2tokens_map[tok['head']] = [tok]
    return head2tokens_map


def _is_np_head(token):
    return token['upostag'] in ['NOUN', 'PROPN']


def _get_subtree(tok, np_tokens, head2tokens_map):
    np_tokens.append(tok)
    if tok['id'] in head2tokens_map:
        for child_token in head2tokens_map[tok['id']]:
            _get_subtree(child_token, np_tokens, head2tokens_map)


def _np_is_continous(np_tokens):
    np_tokens_as_numbers = _tokens_to_numbers(np_tokens)
    return sorted(np_tokens_as_numbers) == list(range(min(np_tokens_as_numbers), max(np_tokens_as_numbers) + 1))


def _tokens_to_numbers(tokens):
    return [int(tok['id'].lstrip('t')) for tok in natsorted(tokens, key=lambda token: token['id'])]


def _get_not_intersecting_nps(nps):
    not_intersecting_nps = []
    ordered_nps = natsorted(nps, key=lambda np: natsorted(np['tokens'])[0])
    for np1 in ordered_nps:
        intersects = False
        for np2 in not_intersecting_nps:
            if _nps_intersects(np1, np2):
                intersects = True
                break
        if not intersects:
            not_intersecting_nps.append(np1)
    return not_intersecting_nps


def _nps_intersects(np1, np2):
    if any(tok_id in np1['tokens'] for tok_id in np2['tokens']):
        return True
    return False


def _tokens_to_ids(tokens):
    return [tok['id'] for tok in natsorted(tokens, key=lambda token: token['id'])]


def _remove_subnps(nps):
    curr_result = []
    result = []
    for ele in sorted(map(collections.OrderedDict.fromkeys, nps), key=len, reverse=True):
        if not any(ele.keys() <= req.keys() for req in curr_result):
            curr_result.append(ele)
            result.append(list(ele))

    return result


def _ids_to_tokens(tokens_ids, id2token_map):
    return [id2token_map[tok_id] for tok_id in natsorted(tokens_ids)]


def _to_np(np_tokens):
    cleaned_tokens = []
    cleaning = True
    ordered_tokens = [tok for tok in natsorted(np_tokens, key=lambda token: token['id'])]
    for tok in ordered_tokens:
        if not cleaning:
            cleaned_tokens.append(tok)
        elif tok['upostag'] in ['ADP', 'CCONJ', 'PUNCT', 'SCONJ']:
            pass
        else:
            cleaning = False
            cleaned_tokens.append(tok)
    if cleaned_tokens[-1]['base'] == ',':
        cleaned_tokens.pop()
    return cleaned_tokens


def _get_np_col_value(tok, nps):
    for np in nps:
        ordered_np_tokens = natsorted(np['tokens'])
        if tok['id'] == ordered_np_tokens[0]:
            return 'B-NP'
        elif tok['id'] in ordered_np_tokens:
            return 'I-NP'
    return 'O'


def _get_local_iate_terms(tokens, iate):
    local_terms = []
    for term in iate:
        for tok in tokens:
            if tok['id'] in term['tokens']:
                local_terms.append(term)
                break
    return local_terms


def _get_local_longest_iate_terms(tokens, iate):
    local_terms = []
    for term in iate:
        for tok in tokens:
            if tok['id'] in term['tokens']:
                local_terms.append(term)
                break
    return _get_longest_terms(local_terms)


def _get_iate_col_value(tok, iate):
    iate_vals = []

    ordered_iate = natsorted(iate, key=lambda term: natsorted(term['tokens'])[0])
    for ti, term in enumerate(ordered_iate, 1):
        if tok['id'] in term['tokens']:
            ordered_term_tokens = natsorted(term['tokens'])
            if tok['id'] == ordered_term_tokens[0]:
                term_obj = IATETerm.objects.get(tid=term['id'])
                domains = term_obj.eurovoc_terms.order_by('tid')
                if domains.exists():
                    iate_vals.append('%d:%s-%s' % (ti, term['id'], ','.join([domain.tid for domain in domains])))
                else:
                    iate_vals.append('%d:%s' % (ti, term['id']))
            else:
                iate_vals.append(str(ti))

    if iate_vals:
        return ';'.join(iate_vals)

    return '_'


def _get_local_eurovoc_terms(tokens, eurovoc):
    local_terms = []
    for term in eurovoc:
        for tok in tokens:
            if tok['id'] in term['tokens']:
                local_terms.append(term)
                break
    return local_terms


def _get_local_longest_eurovoc_terms(tokens, eurovoc):
    local_terms = []
    for term in eurovoc:
        for tok in tokens:
            if tok['id'] in term['tokens']:
                local_terms.append(term)
                break
    return _get_longest_terms(local_terms)


def _get_longest_terms(terms):
    longest_terms = []
    for t1 in terms:
        t1_longest = True
        for t2 in terms:
            if all(tok1 in t2['tokens'] for tok1 in t1['tokens']) and len(t1['tokens']) < len(t2['tokens']):
                t1_longest = False
                break
        if t1_longest and t1 not in longest_terms:
            longest_terms.append(t1)
    return longest_terms


def _terms_intersects(t1, t2):
    if any(tok_id in t1['tokens'] for tok_id in t2['tokens']):
        return True
    return False


def _get_eurovoc_col_value(tok, eurovoc):
    eurovoc_vals = []

    ordered_eurovoc = natsorted(eurovoc, key=lambda term: natsorted(term['tokens'])[0])
    for ti, term in enumerate(ordered_eurovoc, 1):
        if tok['id'] in term['tokens']:
            ordered_term_tokens = natsorted(term['tokens'])
            if tok['id'] == ordered_term_tokens[0]:
                eurovoc_vals.append('%d:%s' % (ti, term['id']))
            else:
                eurovoc_vals.append(str(ti))

    if eurovoc_vals:
        return ';'.join(eurovoc_vals)

    return '_'


def _get_eurovocmt_col_value(tok, eurovoc):
    eurovocmt_vals = []

    ordered_eurovoc = natsorted(eurovoc, key=lambda term: natsorted(term['tokens'])[0])
    for ti, term in enumerate(ordered_eurovoc, 1):
        if tok['id'] in term['tokens']:
            ordered_term_tokens = natsorted(term['tokens'])
            if tok['id'] == ordered_term_tokens[0]:
                mts = EuroVocTerm.objects.get(tid=term['id'], type='descriptor').get_subdomains()
                mts_ids = natsorted([mt.tid for mt in mts])
                eurovocmt_vals.append('%d:%s' % (ti, ','.join(mts_ids)))
            else:
                eurovocmt_vals.append(str(ti))

    if eurovocmt_vals:
        return ';'.join(eurovocmt_vals)

    return '_'