#!/usr/bin/env python

"""10-10-80-splits a file containing item sequences (separated by 1 empty line)
into test, development and training set, leaving sequences intact and basing the
splits on the number of sequences rather than items. The generated files get the
name of the original file, with .test/.dev/.train inserted after the first
.-separated component."""

# FIXME produces one empty sequence

import codecs
import math
import os
import sys

def read_sequence(stream):
    sequence = []
    for line in stream:
        line = line.rstrip()
        if not line:
            return sequence
        else:
            sequence.append(line)
    if sequence:
        return sequence
    return None

splits = ('test', 'dev', 'train')
ratio = {}
ratio['test'] = 0.1
ratio['dev'] = 0.1
ratio['train'] = 0.8

path_all = sys.argv[1]

num_all = 0

with codecs.open(path_all, 'r', encoding='UTF-8') as data:
    sequence = read_sequence(data)
    while sequence != None:
        num_all += 1
        sequence = read_sequence(data)

num = {}
num['test'] = int(math.ceil(ratio['test'] * num_all))
num['dev'] = int(math.ceil(ratio['dev'] * num_all))

(dir_all, file_all) = os.path.split(path_all)
(prefix, suffix) = file_all.split('.', 1)

path = {}
for split in splits:
    file_split = '.'.join((prefix, split, suffix))
    path[split] = os.path.join(dir_all, file_split)

with codecs.open(path_all, 'r', encoding='UTF-8') as data:
    with codecs.open(path['test'], 'w', encoding='UTF-8') as f:
        for i in xrange(0, num['test']):
            for line in read_sequence(data):
                f.write('%s\n' % line)
            f.write('\n')
    with codecs.open(path['dev'], 'w', encoding='UTF-8') as f:
        for i in xrange(num['test'], num['test'] + num['dev']):
            for line in read_sequence(data):
                f.write('%s\n' % line)
            f.write('\n')
    with codecs.open(path['train'], 'w', encoding='UTF-8') as f:
        for i in xrange(num['test'] + num['dev'], num_all):
            for line in read_sequence(data):
                f.write('%s\n' % line)
            f.write('\n')
