#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Add LIWC words as entities in FoLiA XML file.
Usage: python add_liwc_entities.py <file in>
"""
from bs4 import BeautifulSoup
from datetime import datetime
import argparse
import codecs
import os


def add_entity(soup, sentence, cls, words):
    if sentence.find('entities'):
        entities = sentence.entities
    else:
        entities = soup.new_tag('entities')
        sentence.append(entities)

    entity = soup.new_tag('entity')
    entity['class'] = cls
    for w in words:
        entity.append(soup.new_tag('wref', id=w['xml:id'], t=w.t.string))

    entities.append(entity)


def write_folia_file(soup, folia_in, ext):
    output_xml = soup.prettify()
    head, tail = os.path.split(folia_in)
    p = tail.split('.')   
    file_out = '{h}{s}{n}-{e}.xml'.format(n=p[0], h=head, s=os.sep, e=ext)
    with codecs.open(file_out, 'wb', 'utf8') as f:
        f.write(output_xml)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('file_in', help='the name of the FoLiA XML file add '
                        'LIWC entities to')
    args = parser.parse_args()

    file_name = args.file_in

    # Load liwc dict
    with codecs.open('historic_Dutch_LIWC.dic', 'rb', 'utf8') as f:
        lines = f.readlines()

    liwc_categories = {}
    liwc_dict = {}

    for line in lines:
        # LIWC category
        if line[0].isdigit():
            entry = line.split()
            # remove 0 from strings like 01
            c = str(int(entry[0]))
            liwc_categories[c] = entry[1]
        # word
        elif line[0].isalpha():
            entry = line.split()
            term = entry[0]
            categories = entry[1:]
            liwc_dict[term] = categories

    # Load document
    # doc = folia.Document(file='medea-folia-no_events.xml')
    with open(file_name, 'r') as f:
        soup = BeautifulSoup(f, 'xml')

    words = soup.find_all('w')
    for word in words:
        w = word.t.string
        if w in liwc_dict.keys():
            for cat in liwc_dict[w]:
                cat_label = 'liwc-{}'.format(liwc_categories[cat])
                add_entity(soup, word.parent, cat_label, [word])

    annotation_tag = soup.new_tag('entity-annotation')
    annotation_tag['annotator'] = 'liwc'
    annotation_tag['annotatortype'] = 'auto'
    annotation_tag['datetime'] = datetime.now()
    annotation_tag['set'] = 'liwc-set'

    soup.annotations.append(annotation_tag)

    write_folia_file(soup, file_name, 'liwc')

@@ -137,16 +137,70 @@
 in%3E%0A%22%22%22%0A
+from lxml import etree%0Afrom lxml.etree import Element%0A
 from bs4
@@ -705,20 +705,23 @@
 ia_file(
-soup
+context
 , folia_
@@ -734,41 +734,8 @@
 t):%0A
-    output_xml = soup.prettify()%0A
     
@@ -794,19 +794,16 @@
 lit('.')
-   
 %0A    fil
@@ -875,31 +875,24 @@
 t)%0A    with 
-codecs.
 open(file_ou
@@ -896,25 +896,16 @@
 _out, 'w
-b', 'utf8
 ') as f:
@@ -925,20 +925,188 @@
 ite(
-output_xml)%0A
+etree.tostring(context.root,%0A                               encoding='utf8',%0A                               xml_declaration=True,%0A                               pretty_print=True))
 %0A%0Aif
@@ -1962,145 +1962,114 @@
     
-# doc = folia.Document(file='medea-folia-no_events.xml'
+context = etree.iterparse(file_name, events=('end',)
 )%0A
+%0A
     
-with open(file_name, 'r') as f:%0A        soup = BeautifulSoup(f, 'xml')
+for event, elem in context:%0A        pass
 %0A%0A    
+#
 word
@@ -2091,24 +2091,25 @@
 ll('w')%0A    
+#
 for word in 
@@ -2115,24 +2115,25 @@
  words:%0A    
+#
     w = word
@@ -2142,24 +2142,25 @@
 .string%0A    
+#
     if w in 
@@ -2173,32 +2173,33 @@
 ict.keys():%0A    
+#
         for cat 
@@ -2215,24 +2215,25 @@
 ict%5Bw%5D:%0A    
+#
             
@@ -2283,24 +2283,25 @@
 s%5Bcat%5D)%0A    
+#
             
@@ -2354,282 +2354,8 @@
 %5D)%0A%0A
-    annotation_tag = soup.new_tag('entity-annotation')%0A    annotation_tag%5B'annotator'%5D = 'liwc'%0A    annotation_tag%5B'annotatortype'%5D = 'auto'%0A    annotation_tag%5B'datetime'%5D = datetime.now()%0A    annotation_tag%5B'set'%5D = 'liwc-set'%0A%0A    soup.annotations.append(annotation_tag)%0A%0A
     
@@ -2371,20 +2371,23 @@
 ia_file(
-soup
+context
 , file_n

Read and write FoLiA XML file using iterparse