from collections import OrderedDict
import xml.etree.ElementTree as ET
import sys


class Token:
    orth = None
    base = None
    ann = 'O'
    rel = 'O'

    def __str__(self):
        return str((self.orth, self.base, self.ann))


#ccl_file = '/home/kotu/projects/rel2iob/2018-05-15-kpwr-event_all/blogi/00102007.xml'
ccl_file = sys.argv[1]
tree = ET.parse(ccl_file)
chunks = list(tree.getroot())
text_struct = OrderedDict()
sent_ann_tokens = dict()
for chunk in chunks:
    for sentence in list(chunk):
        sentence_id = sentence.get('id')
        text_struct[sentence_id] = list()
        for token in list(sentence):
            if token.tag == 'tok':
                t = Token()
                t.orth = token.find('orth').text
                t.base = token.find('lex[@disamb="1"]/base').text
                for ann in token.findall('ann'):
                    ann_value = ann.text
                    if ann_value != '0':
                        ann_chan = ann.get('chan')
                        t.ann = ann_chan
                        sent_ann_tokens[(sentence_id, ann_chan, ann_value)] = t
                text_struct[sentence_id].append(t)

rel_file = ccl_file.replace('.xml', '.rel.xml')
tree = ET.parse(rel_file)
relations = list(tree.getroot())
for relation in relations:
    rel_name = relation.get('name')
    rel_from = relation.find('from')
    sent_from = rel_from.get('sent')
    chan_from = rel_from.get('chan')
    val_from = rel_from.text
    rel_to = relation.find('to')
    sent_to = rel_to.get('sent')
    chan_to = rel_to.get('chan')
    val_to = rel_to.text
    if sent_from == sent_to and rel_name != 'mod':
        token_from = sent_ann_tokens[(sent_from, chan_from, val_from)]
        token_to = sent_ann_tokens[(sent_to, chan_to, val_to)]
        token_from.rel = rel_name + "_h"
        token_to.rel = rel_name + "_t"


print('-DOCSTART- -X- -X- O\n')
for sentence_id, tokens in text_struct.items():
    for token in tokens:
        '''print(' '.join([token.orth, token.base,
                        'B-{}'.format(token.ann.upper()) if token.ann.upper() != 'O' else 'O',
                        'B-{}'.format(token.rel.upper()) if token.rel.upper() != 'O' else 'O']))'''
        print(' '.join(['{}/{}'.format(token.orth, token.ann.upper()) if token.ann.upper() != 'O' else token.orth, token.base,
                        'B-{}'.format(token.ann.upper()) if token.ann.upper() != 'O' else 'O',
                        'B-{}'.format(token.rel.upper()) if token.rel.upper() != 'O' else 'O']))
    print()
