import random
import json
import math

class DataSplitter:

    def __init__(self, parent_filename):
        self.filename = parent_filename
        self.dataset = []

    def read_data(self):

        ## read in entire dataset
        with open(self.filename) as json_data:
            d = json.load(json_data)
        self.dataset = d["Dataset"]

    def split_data(self, goal_dist):
        dataset, n = self.dataset, len(self.dataset)
        train_prob, dev_prob, test_prob = goal_dist
        train_num= int(train_prob * n)
        dev_num, test_num = train_num + int(dev_prob * n), n

        ## get story count
        story_count = {}
        for idx, dp in enumerate(dataset):
            story = dp['Story']
            if story not in story_count:
                story_count[story] = []
            story_count[story].append(idx)

        ## split into train, dev, test
        lowest_cross_entropy, best_distribution, best_dataset_split = None, None, None

        for attempt in range(0, 200):

            count = 0
            train_idxs, dev_idxs, test_idxs = [], [], []
            curr_list = train_idxs
            story_list = list(story_count.keys())

            random.shuffle(story_list)

            for story in story_list:
                story_idxs = story_count[story]
                new_count = count + len(story_idxs)
                if (count < train_num and new_count > train_num):
                    curr_list = dev_idxs
                if (count < dev_num and new_count > dev_num):
                    curr_list = test_idxs
                curr_list.extend(story_idxs)
                count = new_count

            train_dataset = [dataset[idx] for idx in train_idxs]
            dev_dataset = [dataset[idx] for idx in dev_idxs]
            test_dataset = [dataset[idx] for idx in test_idxs]

            counts = [ len(train_dataset), len(dev_dataset), len(test_dataset)]
            norm_ = sum(counts)
            probs = [count_/float(norm_) for count_ in counts]
            cross_entropy = -sum([ w * math.log(v + 0.001) for (w, v) in zip(goal_dist, probs)])
            # print("Cross Entropy is %r and Distribution is %r " % (cross_entropy, probs))
            if lowest_cross_entropy is None or cross_entropy < lowest_cross_entropy:
                lowest_cross_entropy, best_distribution, best_dataset_split = \
                    cross_entropy, probs, (train_dataset, dev_dataset, test_dataset)

        print("Lowest cross entropy is %r and best split is %r" % (lowest_cross_entropy, best_distribution))
        train_dataset, dev_dataset, test_dataset = best_dataset_split

        ## write to file
        fp = self.filename.split(".")

        train_json = {"Dataset":train_dataset}
        train_filename = fp[0] + "_train." + fp[1]
        with open(train_filename, 'w') as f:
            json.dump(train_json, f)

        dev_json = {"Dataset":dev_dataset}
        dev_filename = fp[0] + "_dev." + fp[1]
        with open(dev_filename, 'w') as f:
            json.dump(dev_json, f)

        test_json = {"Dataset":test_dataset}
        test_filename = fp[0] + "_test." + fp[1]
        with open(test_filename,'w') as f:
            json.dump(test_json, f)

        print('train set size = {}/{} = {}%'.format(len(train_dataset), n, len(train_dataset)/n))
        print('dev set size = {}/{} = {}%'.format(len(dev_dataset), n, len(dev_dataset)/n))
        print('test set size = {}/{} = {}%'.format(len(test_dataset), n, len(test_dataset)/n))

if __name__ == '__main__':
    pass

## JUST TO DOUBLE CHECK
# train_stories, dev_stories, test_stories = set(), set(), set()
# for i in train_dataset:
#     train_stories.add(i['Story'])
# for i in dev_dataset:
#     dev_stories.add(i['Story'])
# for i in test_dataset:
#     test_stories.add(i['Story'])
#
# print('Train Stories:----')
# for st in train_stories:
#     print(st + '\n')
# print('Dev Stories:----')
# for st in dev_stories:
#     print(st+ '\n')
# print('Test Stories:----')
# for st in test_stories:
#     print(st+ '\n')