#!/usr/bin/perl
# dqsub submits jobs using qsub with better options
# and is released under the terms of the GNU GPL version 3, or any
# later version, at your option. See the file README and COPYING for
# more information.
# Copyright 2014 by Don Armstrong <don@donarmstrong.com>.


use warnings;
use strict;

use Getopt::Long;
# use Pod::Usage;

=head1 NAME

dqsub - submits jobs using qsub with better options

=head1 SYNOPSIS

dqsub [options]

 Options:
   --queue, -q Queue to use
   --interactive, -I call qsub interactively
   --nodes nodes to use
   --array array mode (one of 'chdir' or 'xargs' or '')
   --array-from file to read arrays from (default STDIN)
   --array-slot-limit
   --ppn processors per node to use
   --mem memory to request
   --dir Directory to run the script in (default current directory)
   --debug, -d debugging level (Default 0)
   --help, -h display this help
   --man, -m display manual

=head1 OPTIONS

=over

=item B<--array>

This describes how dqsub will generate array jobs.

If no B<--array> is given, then the command and any additional
arguments given will be run using qsub.

If B<--array> is C<chdir>, then each line of the input given in
B<--array-from> will be used as a directory and the command and any
additional arguments given will run in each directory.

IF B<--array> is C<xargs>, then each line of the input given will be
considered to be an additional argument which will be given to the
command run in the current directory.

=item B<--array-from>

File to read array arguments from. If not provided, and B<--array> is
given, arguments will be read from STDIN.

=item B<--debug, -d>

Debug verbosity. (Default 0)

=item B<--help, -h>

Display brief usage information.

=item B<--man, -m>

Display this manual.

=back

=head1 EXAMPLES

dqsub

=cut

use IO::File;
use Cwd qw(getcwd abs_path);
use vars qw($DEBUG);

my %options = (nodes           => 1,
               ppn             => 2,
               mem             => '2G',
               debug           => 0,
               help            => 0,
               man             => 0,
               interactive     => 0,
              );

GetOptions(\%options,
           'queue|q=s',
           'interactive|I!',
           'nodes=i',
           'array=s',
           'array_from|array-from=s',
           'array_slot_limit|array-slot-limit=i',
           'ppn|processors-per-node=i',
           'mem|memory=s',
           'dir=s',
           'debug|d+','help|h|?','man|m');

# pod2usage() if $options{help};
# pod2usage({verbose=>2}) if $options{man};

$DEBUG = $options{debug};

my @USAGE_ERRORS;
if (not @ARGV and not $options{interactive}) {
    push @USAGE_ERRORS,"You must provide a command to run";
}
if (defined $options{array} and $options{array} !~ /^(?:|chdir|xargs)$/i) {
    push @USAGE_ERRORS,"--array must be one of chdir, xargs or '' if provided";
    $options{array} = lc($options{array});
    if ($options{array} eq '') {
        $options{array} = undef;
    }
}
if ($options{interactive} and @ARGV) {
    push @USAGE_ERRORS,"Don't provide commands when you're asking for an interactive shell";
}

# pod2usage(join("\n",@USAGE_ERRORS)) if @USAGE_ERRORS;
print STDERR join("\n",@USAGE_ERRORS) and exit 1 if @USAGE_ERRORS;

# OK. Generate the options to qsub which we'll be using
my @qsub_options = generate_qsub_options(\%options);

if ($options{interactive}) {
    print STDERR 'running: qsub '.join(' ',@qsub_options) if $DEBUG;
    exec('qsub',@qsub_options);
} else {
    my @array = ();
    if ($options{array}) {
        @array = read_array_options(\%options) if $options{array};
        # the -t option gives the range of elements for an array job
        push @qsub_options,'-t','1-'. scalar @array;
        if ($options{array_slot_limit}) {
            $qsub_options[$#qsub_options] .= '%'.$options{array_slot_limit};
        }
    }
    call_qsub(\@qsub_options,write_qsub_script(\%options,\@ARGV,\@array));
}

sub generate_qsub_options{
    my ($options) = @_;
    my @qo;
    if (defined $options->{queue} and length $options->{queue}) {
        push @qo,'-q',$options->{queue};
    }
    if (defined $options->{dir}) {
        push @qo,'-d',abs_path($options->{dir});
    } else {
        push @qo,'-d',getcwd;
    }
    ## handle the -l options
    my @l;
    push @l, 'nodes='.$options->{nodes};
    if (defined $options->{ppn}) {
        $l[$#l] .= ':ppn='.$options->{ppn};
    }
    if ($options->{mem}) {
        push @l,'mem='.$options->{mem};
    }
    push @qo,'-l',join(',',@l) if @l;
    if ($options->{interactive}) {
        push @qo,'-I';
    }
    return @qo;
}

sub read_array_options{
    my ($options) = @_;
    my $fh = \*STDIN;
    if (defined $options->{array_from}) {
        $fh = IO::File->new(defined $options->{array_from}) or
            die "Unable to open $options->{array_from} for reading: $!";
    }
    my @array;
    for (<$fh>) {
        chomp;
        push @array,$_;
    }
    return @array;
}

sub call_qsub {
    my ($qsub_options,$script) = @_;
    my $qsub_fh;
    open $qsub_fh,'|-','qsub',@{$qsub_options},'-' or
        die "Unable to start qsub: $!";
    print {$qsub_fh} $script or
        die "Unable to print to qsub: $!";
    close($qsub_fh) or
        die "Unable to close qsub filehandle: $!";
}

sub write_qsub_script {
    my ($opt,$arg,$array) = @_;

    my $script = "#!/bin/bash\n";
    my $command = join(' ',map {qq('$_')} @{$arg});
        $script .= <<EOF;
# this script was written by dqsub
EOF
    if (defined $opt->{array}) {
        my $array_opt = join("\n",@{$array});
        $script .= <<EOF;
OPT=\$(sed -n -e "\$PBS_ARRAYID p"<<'_HERE_DOC_END_'
$array_opt
_HERE_DOC_END_
)
EOF
        if ($opt->{array} eq 'chdir') {
            $script .= <<EOF;
cd "\$OPT";
exec $command;
EOF
        } else {
            $script .= <<EOF;
exec $command "\$OPT";
EOF
        }
    } else {
        $script .= <<EOF;
# there's no array, so just executing the command with arguments
exec $command;
EOF
    }
    return $script;
}


__END__