import json
import nltk
import os

class HouseVocab:

    def __init__(self, house_names):
        self.house_names = house_names
        self.data_filenames = [house_name + '_discrete.json' for house_name in house_names]
        self.vocab = {}
        self.vocab_list = []

    def create_vocab(self, vocab_filename):
        if False:
        # if os.path.isfile(vocab_filename):
            with open(vocab_filename, 'r') as f:
                self.vocab = json.load(open('dict_'+vocab_filename))
                with open(vocab_filename) as f:
                    self.vocab_list = f.read().splitlines()

        else:
            vocab = {}
            idx_cnt = 0
            for hi in range(len(self.house_names)):

                data = json.load(open(self.data_filenames[hi]))
                datapoints_json = data["Dataset"]

                for dp in datapoints_json:
                    instruction = dp['Instruction']
                    words = set([x.lower() for x in nltk.word_tokenize(instruction)])
                    for word in words:
                        if word not in vocab:
                            vocab[word] = idx_cnt
                            idx_cnt += 1

                # save vocab
                with open('dict_' + vocab_filename, 'w') as f:
                    json.dump(vocab, f)
                self.vocab = vocab

                # save vocab list
                vocab_list = list(vocab.keys())
                f = open(vocab_filename, 'w')
                for word in vocab_list:
                    f.write("%s\n" % word)
                f.close()
                self.vocab_list = vocab_list

    def tokenize_instructions(self, hidx):
        vocab = self.vocab
        data = json.load(open(self.data_filenames[hidx]))
        datapoints_json = data["Dataset"]

        for dp in datapoints_json:
            instruction = dp['Instruction']
            instruction_idxs = [vocab[token.lower()] for token in nltk.tokenize.word_tokenize(instruction)]
            dp['Instruction'] = instruction_idxs

        with open('tokenized_' + self.data_filenames[hidx], 'w') as f:
            json.dump(data, f)


if __name__ == '__main__':
    pass