#!/usr/bin/env python

import os, sys, math, random

from collections import defaultdict

if sys.version_info[0] >= 3:

    xrange = range

def exit_with_help(argv):

    print("""\

Usage: {0} [options] dataset subset_size [output1] [output2]

This script randomly selects a subset of the dataset.

options:

-s method : method of selection (default 0)

     0 -- stratified selection (classification only)

     1 -- random selection

output1 : the subset (optional)

output2 : rest of the data (optional)

If output1 is omitted, the subset will be printed on the screen.""".format(argv[0]))

    exit(1)

def process_options(argv):

    argc = len(argv)

    if argc < 3:

        exit_with_help(argv)

    # default method is stratified selection

    method = 0

    subset_file = sys.stdout

    rest_file = None

    i = 1

    while i < argc:

        if argv[i][0] != "-":

            break

        if argv[i] == "-s":

            i = i + 1

            method = int(argv[i])

            if method not in [0,1]:

                print("Unknown selection method {0}".format(method))

                exit_with_help(argv)

        i = i + 1

    dataset = argv[i]

    subset_size = int(argv[i+1])

    if i+2 < argc:

        subset_file = open(argv[i+2],'w')

    if i+3 < argc:

        rest_file = open(argv[i+3],'w')

    return dataset, subset_size, method, subset_file, rest_file

def random_selection(dataset, subset_size):

    l = sum(1 for line in open(dataset,'r'))

    return sorted(random.sample(xrange(l), subset_size))

def stratified_selection(dataset, subset_size):

    labels = [line.split(None,1)[0] for line in open(dataset)]

    label_linenums = defaultdict(list)

    for i, label in enumerate(labels):

        label_linenums[label] += [i]

    l = len(labels)

    remaining = subset_size

    ret = []

    # classes with fewer data are sampled first; otherwise

    # some rare classes may not be selected

    for label in sorted(label_linenums, key=lambda x: len(label_linenums[x])):

        linenums = label_linenums[label]

        label_size = len(linenums)

        # at least one instance per class

        s = int(min(remaining, max(1, math.ceil(label_size*(float(subset_size)/l)))))

        if s == 0:

            sys.stderr.write('''\

Error: failed to have at least one instance per class

    1. You may have regression data.

    2. Your classification data is unbalanced or too small.

Please use -s 1.

''')

            sys.exit(-1)

        remaining -= s

        ret += [linenums[i] for i in random.sample(xrange(label_size), s)]

    return sorted(ret)

def main(argv=sys.argv):

    dataset, subset_size, method, subset_file, rest_file = process_options(argv)

    #uncomment the following line to fix the random seed

    #random.seed(0)

    selected_lines = []

    if method == 0:

        selected_lines = stratified_selection(dataset, subset_size)

    elif method == 1:

        selected_lines = random_selection(dataset, subset_size)

    #select instances based on selected_lines

    dataset = open(dataset,'r')

    prev_selected_linenum = -1

    for i in xrange(len(selected_lines)):

        for cnt in xrange(selected_lines[i]-prev_selected_linenum-1):

            line = dataset.readline()

            if rest_file:

                rest_file.write(line)

        subset_file.write(dataset.readline())

        prev_selected_linenum = selected_lines[i]

    subset_file.close()

    if rest_file:

        for line in dataset:

            rest_file.write(line)

        rest_file.close()

    dataset.close()

if __name__ == '__main__':

    main(sys.argv)