import json
import nltk
import pandas as pd

vocab = open("house_all_vocab.txt").readlines()
vocab_dict = dict()
for i, token in enumerate(vocab):
	vocab_dict[i] = token.strip()

values = []
all_instructions = []
unique_stories = dict()

dataset = dict()

for i in range(1, 6):
	for split in ["train", "dev", "test"]:
		with open("tokenized_house%r_discrete_%s.json" % (i, split), 'r') as f:
			data = json.load(f)
			dataframe = pd.DataFrame(data["Dataset"])
			ids = list(dataframe['ID'])
			stories = set(list(dataframe['Story']))
			instructions = [tuple(v) for v in list(dataframe["Instruction"])]

			for inst_, id_ in zip(instructions, ids):
				dataset["house_%d_%d" % (i, id_)] = " ".join([vocab_dict[l] for l in inst_]).strip()

			# check ids are not repeated
			assert len(ids) == len(list(set(ids))), "IDs are repeating within a split"

			values.extend(ids)
			all_instructions.extend(instructions)
			assert len(ids) == len(instructions), "there should be as many IDs are instructions"
			print("House %d, split %s size is %d" % (i, split, len(ids)))
			unique_stories["house_%d_%s" % (i, split)] = stories

num_stories = 0
for key in unique_stories:
	num_stories += len(unique_stories[key])
print("Num stories ", num_stories)


for key1 in unique_stories:
	for key2 in unique_stories:
		if key1 != key2:
			intersection = unique_stories[key1].intersection(unique_stories[key2])
			assert len(intersection) == 0, "Intersection should be 0"


# Check if values are all unique

assert len(values) == len(list(set(values))), "IDs are repeating within a split"
print("Size of the dataset is ", len(values))

print("Total number of instructions are ", len(all_instructions))
unique_instructions = list(set(all_instructions))

vocab = open("house_all_vocab.txt").readlines()
vocab_dict = {}
for i, tk in enumerate(vocab):
      vocab_dict[i] = tk.strip()

f = open("chai_unique_instructions", "w")
for inst in unique_instructions:
     if len(inst) >= 3:
            inst_str = " ".join([vocab_dict[i] for i in inst])
            f.write(inst_str + "\n")
f.flush()
f.close()

print("Total number of unique instructions are %r which percentage wise %r" % (len(unique_instructions), (len(unique_instructions) * 100 / float(len(all_instructions)))))

num_long_instructions = 0
num_tokens_per_instructions = 0
for inst in all_instructions:
	if len(inst) >= 3:
		num_long_instructions += 1
		num_tokens_per_instructions += len(inst)
print("Number of instructions with atleast 3 words are ", num_long_instructions)
print("Number of tokens per instructions are", num_tokens_per_instructions/float(num_long_instructions))
#
# counter = dict()
# for tup in all_instructions:
# 	if tup in counter:
# 		counter[tup] = counter[tup] + 1
# 	else:
# 		counter[tup] = 1
#
# # Sort the dictionary by value and save the content
# f = open("tmp_instructions.txt", "w")
# for key, value in sorted(counter.iteritems(), key=lambda (k,v): (-v,k)):
# 	instruction_string = " ".join([vocab_dict[i] for i in key]).strip()
# 	f.write("%s: %d\n" % (instruction_string, value))
# f.flush()
# f.close()

# Create the segments data
# lines = open("./dataset_trajectory_segments.txt").readlines()
# index = 0
# curr = ""
# id_ = 0
# json_lengths = []
# json_sum, json_cnt = 0, 0
# for line in lines:
# 	line = line.strip()
# 	if len(line) == 0:
# 		continue
# 	if line.startswith("$$$$$$$$$"):
# 		index += 1
# 		curr = ""
# 		continue
# 	if line.startswith("=========="):
# 		index = 0
# 		curr = ""
# 		# example with inst, house and id_
# 		# example = (id_, house, inst)
# 		# inst1 = inst.lower().replace(" ", "")
# 		# inst2 = dataset["house_%d_%d" % (house, id_)].lower().replace(" ", "")
# 		# if inst1 != inst2:
# 		# 	print("Warning: Instructions should be the same. Found %s and %s " % (inst1, inst2))
# 		# id_ += 1
# 		continue
# 	curr = curr + line
# 	if index == 4:
# 		inst = curr
# 	elif index == 6:
# 		house = int(curr[len("house")])
# 	elif index == 8:
# 		trajectory = curr
# 		jobj = json.loads(trajectory)
# 		dios = jobj["dios"]
# 		#print(len(dios), dios[0], dios[1])
# 		if len(nltk.word_tokenize(inst)) >= 3:
# 			json_lengths.append(len(dios))
# 			json_sum += len(dios)
# 			json_cnt += 1
# 	else:
# 		curr = ""
#
# print("Max is %d and min is %d " %(max(json_lengths), min(json_lengths)))
# avg_length = sum(json_lengths)/len(json_lengths)
# avg_length2 = json_sum / float(json_cnt)
# print("Num segments %d " %  json_cnt)
# print("Average JSON length = {}".format(avg_length))
# print("Average JSON length = {}".format(avg_length2))