# -*- coding: utf-8 -*-
from data_utils import get_file
import string
import random
import cPickle

def make_reuters_dataset(path='datasets/temp/reuters21578/', min_samples_per_topic=15):
    import os
    import re
    from preprocessing.text import Tokenizer

    wire_topics = []
    topic_counts = {}
    wire_bodies = []

    for fname in os.listdir(path):
        if 'sgm' in fname:
            s = open(path + fname).read()
            tag = '<TOPICS>'
            while tag in s:
                s = s[s.find(tag)+len(tag):]
                topics = s[:s.find('</')]
                
                if topics and not '</D><D>' in topics:
                    topic = topics.replace('<D>', '').replace('</D>', '')
                    wire_topics.append(topic)
                    topic_counts[topic] = topic_counts.get(topic, 0) + 1
                else:
                    continue

                bodytag = '<BODY>'
                body = s[s.find(bodytag)+len(bodytag):]
                body = body[:body.find('</')]
                wire_bodies.append(body)

    # only keep most common topics
    items = topic_counts.items()
    items.sort(key = lambda x: x[1])
    kept_topics = set()
    for x in items:
        print x[0] + ': ' + str(x[1])
        if x[1] >= min_samples_per_topic:
            kept_topics.add(x[0])
    print '-'
    print 'Kept topics:', len(kept_topics)

    # filter wires with rare topics
    kept_wires = []
    labels = []
    topic_indexes = {}
    for t, b in zip(wire_topics, wire_bodies):
        if t in kept_topics:
            if t not in topic_indexes:
                topic_index = len(topic_indexes)
                topic_indexes[t] = topic_index
            else:
                topic_index = topic_indexes[t]

            labels.append(topic_index)
            kept_wires.append(b)

    # vectorize wires
    tokenizer = Tokenizer()
    tokenizer.fit(kept_wires)
    X = tokenizer.transform(kept_wires)

    print 'Sanity check:'
    for w in ["banana", "oil", "chocolate", "the", "dsft"]:
        print '...index of', w, ':', tokenizer.word_index.get(w)
    cPickle.dump(dataset, open('datasets/data/reuters.pkl', 'w'))



def load_data(path="reuters.pkl", nb_words=100000, maxlen=None, test_split=0.2, seed=113):
    path = get_file(path, origin="https://s3.amazonaws.com/text-datasets/reuters.pkl")
    f = open(path, 'rb')

    X, labels = cPickle.load(f)
    f.close()
    random.seed(seed)
    random.shuffle(X)
    random.seed(seed)
    random.shuffle(labels)

    if maxlen:
        new_X = []
        new_labels = []
        for x, y in zip(X, labels):
            if len(x) < maxlen:
                new_X.append(x)
                new_labels.append(y)
        X = new_X
        labels = new_labels

    X = [[1 if w >= nb_words else w for w in x] for x in X]
    X_train = X[:int(len(X)*(1-test_split))]
    y_train = labels[:int(len(X)*(1-test_split))]

    X_test = X[int(len(X)*(1-test_split)):]
    y_test = labels[int(len(X)*(1-test_split)):]

    return (X_train, y_train), (X_test, y_test)


if __name__ == "__main__":
    make_reuters_dataset()
    (X_train, y_train), (X_test, y_test) = load_data()

    dataset = (X, labels) 
    print '-'
    print 'Saving...'