# -*- coding: utf-8 -*-
import random

read_file = 'enwiki_conv_2.txt'
write_file = 'enwiki_conv_3.txt'

def get_context_word(num):
    fr = open(read_file, 'r')
    fw = open(write_file, 'w')
    cword_dict = {}
    for r in fr.readlines():
        words = r.split()
        for i, word in enumerate(words):
            l = []
            for n in range(num):
                if i+(n+1) > len(words)-1:
                    continue
                l.append(words[i+(n+1)])
            for n in reversed(range(num)):
                if i-(n+1) < 0:
                    continue
                l.append(words[i-(n+1)])
            if word not in cword_dict:
                cword_dict[word] = l
            else:
                cword_dict[word] += l
    for key, value in cword_dict:
        for v in value:
            fw.write(key + '\t' + v + '\n')

def main():
    #num = random.randint(1,5)
    num = 1
    get_context_word(num)

if __name__ == '__main__':
    main()
