#!/usr/bin/perl
# create_hostlist.pl - Generate a host list suitable for input to MPI
# author: James Dinan <dinan@cse.ohio-state.edu>

# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA

### DEFAULT OPTIONS ###
$prefix = "node";
$suffix = "";
$append = "";
$separator = "\n";

### SOME GLOBAL THINGS ###
$version = "0.3";

### PARSE CMD LINE ###
if ($#ARGV < 0) {
	print "Try -h for help.\n";
	exit 0;
}

while ($#ARGV >= 0) {
	$_ = $ARGV[0];
	shift;

	if (/^-(-)?h(elp)?$|^-(-)?u(sage)?$/i) {
		usage();
		exit 0;
	}
	elsif (/^-p$/) {
		$prefix = $ARGV[0];
		shift or die "Not enough arguments";
	}
	elsif (/^-s$/) {
		$suffix = $ARGV[0];
		shift or die "Not enough arguments";
	}
	elsif (/^-i$/) {
		$separator = " ";
	}
	elsif (/^-a$/) {
		$append = $ARGV[0];
		shift or die "Not enough arguments";
	}
	elsif (/^(-?)([0-9]+)$/) {
		# Here and in the next one, $1 is used as a flag to tell whether
		# the user gave a minus sign on this number/range
		if ($1) {
			delete $nodes{$2};
		} else {
			$nodes{$2} = 1;
		}
			
	}
	elsif (/^(-?)([0-9]+)-([0-9]+)$/) {
		if ($3 <= $2) {
			print "Error: $2-$3 is not a valid range\n";
			exit 1;
		}

		for (my $c = $2; $c <= $3; $c++) {
			if ($1) {
				delete $nodes{$c};
			} else {
				$nodes{$c} = 1;
			}
		}
	}
	else {
		print "Unrecognized parameter: $_.  Try -h for help.\n";
		exit 1;
	}
}

# This call to sort is magic.  Thank god for google.
# Output the nodes in a numerically sorted order.
$append_str = " $append" unless ($append eq "");

foreach my $node (sort { $a <=> $b } keys %nodes) {
	print "$prefix$node$suffix$append_str$separator";
}

### SUBROUTINES ###
sub usage {
	my $suffix_str = ($suffix eq "") ? "None" : $suffix;
	my $append_str = ($append eq "") ? "None" : $append;

	print<<EndOfUsage;
HostList Writer (v$version): by James Dinan <dinan\@cse.ohio-state.edu>
                Generate a host list suitable for input to MPI

Usage: create_hostlist.pl [flags] [node numbers]

Flags:
       -h      This Message
       -p str  Specify a hostname prefix (default: $prefix)
       -s str  Specify a hostname suffix. eg: ".osu.edu" (default: $suffix_str)
       -a str  Append the string str to each line (default: $append_str)
       -i      Inline host list (eg "h1 h2..."), all on one line

Node Numbers:
       XX      Add node number XX to the list (eg: 1 2 3...)
      -XX      Remove node number XX to the list (eg: -2 -3...)
       XX-YY   Add nodes XX up to YY to the list (eg: 1-10)
      -XX-YY   Remove nodes XX up to YY to the list (eg: -1-10)

Hostlist Example:
       ./create_hostlist.pl -p node -s .osu.edu 1-10 -5 -2
          node1.osu.edu
          node3.osu.edu
          node4.osu.edu
          node6.osu.edu
          node7.osu.edu
          node8.osu.edu
          node9.osu.edu
          node10.osu.edu

Inline Example (eg: mpirun -np 16 \$(create_hostlist...) ./par_app):
       ./create_hostlist.pl -i -p bm 1-10 -2
          bm1 bm3 bm4 bm5 bm6 bm7 bm8 bm9 bm10
EndOfUsage
}