#!/usr/bin/env python

import codecs
import logging
import nltk
import sys

logging.basicConfig(level=logging.INFO)

try:
    sys.argv.remove('--force-newline')
    force_newline = True
except ValueError:
    force_newline = False

lang = sys.argv[1]

out = codecs.getwriter('UTF-8')(sys.stdout)

def new_sentence():
    global current_sentence, current_doc
    if current_sentence:
        current_doc.append(current_sentence)
        current_sentence = ''

def new_doc():
    global current_doc
    new_sentence()
    if current_doc:
        docs.append(current_doc)
        current_doc = []

# Part 1: read in test corpus in IOB format

current_doc = []
current_sentence = ''
docs = []
new_doc()

for line in sys.stdin:
    line = line.strip()
    if line:
        char, label = line.split(' ')
        char = int(char)
        if label == 'S':
            new_sentence()
        current_sentence += unichr(char)
    else:
        new_doc()

# Part 2: sentence-segment each document using Punkt:

punkt = nltk.data.load('tokenizers/punkt/%s.pickle' % lang)
total_sentences = 0

for doc in docs:
    doc = ''.join(doc)
    output = 0
    sentences0 = punkt.tokenize(doc, realign_boundaries=True)
    if force_newline:
        # Force splitting at newlines (very important for PAISA):
        sentences = []
        for sentence in sentences0:
            sentences.extend(sentence.split('\n'))
    else:
        sentences = sentences0
    total_sentences += len(sentences)
    logging.debug('<SENT>'.join(sentences))
    # Align sentences with doc, collecting sentence-starting indices:
    starts = []
    sentindex = 0
    prevlen = 0
    for sentence in sentences:
        try:
            sentindex = doc.index(sentence[:-1], sentindex + prevlen)
            prevlen = len(sentence)
        except ValueError:
            sys.stderr.write('Sentence not found:\n%s\n\nIn document:\n%s\n\n' % (sentence, doc))
        starts.append(sentindex)
    for i in range(len(doc)):
        if i in starts:
            label = 'S'
        else:
            label = 'I'
        print '%d %s' % (ord(doc[i]), label)
    print

logging.info('%d sentences in total', total_sentences)
