#!/usr/bin/perl

use POSIX qw(ceil floor);
use Cwd;
use strict;

if (@ARGV < 3) { # print a warning if we didn't get the right number of features
	print "Correct syntax is:\n\tbuild_sets.pl corpus_name train_ratio input_file_list\n";
	exit;
}

my $corpus_name = $ARGV[0]; # testing and training files are named $corpus_name.testing and $corpus_name.training, respectively
my $train_ratio = $ARGV[1]; # floor(author_doc_count * train_ratio) docs from each author will be randomly assigned to the training set

# pop the top two entried off the list, leaving us with a file list
shift(@ARGV);
shift(@ARGV);

my $cwd = getcwd(); # get the current directly
print "Output Directory: $cwd\n";

# we'll use these to keep track of how many files/docs we looked at
my $files = 0; 
my $doccount = 0;

my @docs = (); # initalize docs array

foreach my $name (@ARGV) { # work though our list of files
	if (-e $name) { # make sure the files exist
		print "Reading from file $name\n"; # let user know what we're up to
	
		open (INFILE, "$name"); # open up the input file
		my @input = <INFILE>; # slurp the input file into memory
		close (INFILE); # close out the input file
		
		my $fullinput = join(' ', @input); # link together all of the lines in the input into a single scalar
		
		while ($fullinput =~ /(<DOC>.+?<\/DOC>)/s) { 
			#there's still at least one doc in the string, so extract it
			push @docs, $1; # push the doc we found into our array of docs
			$fullinput =~ s/<DOC>.+?<\/DOC>//s; # remove that doc from the string
		}
		
		$files++; # increment file count
	}
}

$doccount = @docs; # figure out how many documents we have

print "Loaded $doccount documents from $files files\n\n"; # print a little info

my %authors = (); # initialize the author hash

for( my $i = 0; $i < @docs; $i++) {
	# loop through each of the docs, extract the author id and load it into a hash
	$docs[$i] =~ /<AUTHOR>(.+)<\/AUTHOR>/; # extract author name
	my $author = $1; # stuff the author name in a conveniently named variable
	$authors{$author} = $authors{$author} . "$i:"; # load the index of the doc into the appropriate hash entry
}

# open handles to the testing and training files
open (TRAINING, ">" . $corpus_name . ".training");
open (TESTING, ">" . $corpus_name . ".testing");

foreach my $key (sort( keys %authors)) { 
	# work on each author in the hash
	chop $authors{$key}; # remove the trailing colon from the entry
	my @indices = split(/:/, $authors{$key}); # get a list of doc indexes
	
	$doccount = @indices; # figure out how many docs we have by the current author
	if ($doccount == 1) { # author has only one article, so skip them
		print "1 doc in set from author $key. Author will not be included\n\n";
		next;
	}
	
	print "$doccount docs in set from author $key\n"; # print a little status
	
	my $choose = floor($doccount * $train_ratio);  # figure out how many docs to choose for training
	if ($choose == 0) { $choose = 1; } # make sure we always pull at least one doc for training
	
	print "Choosing $choose docs from this author for training\n"; # print a little more status
	
	my $chosen = 0; # number of docs already chosen for the training set
	
	foreach my $index (@indices) {
		# work on each of the author's docs
		
		my $set; # we'll use this to note if the file went to testing or training
		
		if(rand(1) < (($choose - $chosen) / $doccount)) {
			# this doc should go into the training set
			$set = "TRAINING "; # note that it's going into training
			print TRAINING "$docs[$index]\n\n"; # write it into the training file
			$chosen++; # increment number of docs chosen
		} else {
			# this doc should go in the testing set
			$set = "TESTING  "; # note that it's going into testing
			print TESTING "$docs[$index]\n\n"; # write it into the testing file
		}
		$doccount--; # remove this doc from the number remaining
		
		$docs[$index] =~ /<DOCID>(.+)<\/DOCID>/; # extract document ID
		print "\t$set$1\n"; # print out the docid of the author's work and where it went
	}
	
	print "\n"; # a padding line between authors
}

# close handles to training and testing files
close TRAINING;
close TESTING;