#!/usr/bin/python
#
# Bismillahi-r-Rahmani-r-Rahim
# Create an RTE dataset

import os
#import nltk
#from nltk.tokenize import sent_tokenizer
from nltk import tokenize
from entailment.NGramOverlap import NGramOverlap

data_dir = '/home/daoud/Data/rte'

pos_threshold = 0.578
neg_threshold = 0.4
rte = NGramOverlap(3, pos_threshold)
max_pairs = 10

def ent(text, hypothesis):
    global rte
    if len(text) < 20 or len(hypothesis) < 20:
        return False
    if len(text) > 500 or len(hypothesis) > 500:
        return False
    if text.count('"') % 2 != 0 or hypothesis.count('"') %2 != 0:
        return False
    if text.count('\n') > 4 or hypothesis.count('\n') > 4:
        return False
    return rte.ent( (text, hypothesis) )

def output_pairs(positive, negative):
    print "Positive:", len(positive), '\n', positive
    print "Negative:", len(negative), '\n', negative

def process_documents(path):
    positive_pairs = []
    negative_pairs = []
    
    subdirs = os.listdir(path)
    subdirs.sort()
    for s in subdirs:
        new_path = os.path.join(path,s)
        files = os.listdir(new_path)
        files.sort()
        for f in files:
            file_path = os.path.join(new_path,f)
            print file_path
            file_ = open(file_path)
            headline = file_.readline()[:-1]
            body = file_.read()
            sentences = tokenize.sent_tokenize(body)
            if ent(sentences[0], headline) > pos_threshold:
                print sentences[0]
                print headline
                print
                positive_pairs.append( (sentences[0], headline) )
            for i in range(len(sentences) - 2):
                text = sentences[i]
                hypothesis = sentences[i+2]
                val = ent(text, hypothesis)
                if val > neg_threshold:
                    print "-- Non entailing --"
                    print text
                    print hypothesis
                    print
                    negative_pairs.append( (val, text, hypothesis) )
            if len(positive_pairs) > max_pairs and len(negative_pairs) > max_pairs:
                output_pairs(positive_pairs, negative_pairs)
                return

if __name__ == "__main__":
    global data_dir
    process_documents(data_dir)
