#!/usr/bin/python3

"""This script contains some utility functions used throughout my WMR
related code.

Requires the env var PYTHONPATH to be set to the root directory
containin the code for WMR.
"""

import random
import pickle

def store_pickle(obj, pickle_file):
    with open(pickle_file, "wb") as f:
        pickle.dump(obj, f)
        
def load_pickle(pickle_file):
    with open(pickle_file, "rb") as f:
        return pickle.load(f)

def split_dataset(dataset, split):
    """
    Splits the dataset with where split is the percentage of the
    dataset that will form the train set, and the remaining 1 - split
    is the percentage that will form the test set.
    """
    max_row = len(dataset)
    yet_to_pick =  max_row * split
    train = []
    rows_picked = set()

    # NOTE: this has to be fixed to allow different executions to have
    # the same train/test split.
    random.seed(1000)

    # -- construct train
    while yet_to_pick:
        r = random.randint(0, max_row - 1)

        if r not in rows_picked:
            rows_picked.add(r)
            train.append(dataset[r])
            yet_to_pick -= 1

    # -- remaining is test set
    test = []
    for i in range(0, max_row):
        if i not in rows_picked:
            test.append(dataset[i])

    return train, test