#! /usr/bin/env perl

use Carp;
use warnings;
use strict;

use IPC::Cmd 0.36 qw( can_run run );
use PS::IPP::Metadata::Config;
use PS::IPP::Metadata::List qw( parse_md_list );
use PS::IPP::Config 1.01 qw( :standard );
use DateTime;
use Data::Dumper;
use File::Basename;

use Getopt::Long qw( GetOptions :config auto_help auto_version gnu_getopt );
use Pod::Usage qw( pod2usage );


# Hard coded values
#my $remote_root   = '/scratch3/watersc1/';  # Far side destination base location
my $remote_root   = '/lustre/scratch1/turquoise/watersc1/ps1/'; # Far side destination base location
$remote_root = '/scratch3/watersc1/';
my $threads       = 2;                      # How many threads are we going to use?
my $job_cost      = 1500 / 60 / 60;           # Estimate of how long a job runs, in hours.
my $proc_per_node = 24;                     # processors per node
my $min_nodes     = 1;                      # smallest allocation to ask for
my $max_nodes     = 1000;                   # largest allocation to ask for
my $min_time      = 1;                      # shortest allocation to ask for
my $max_time      = 8;                      # longest allocation to ask for
my $remote_hostname   = "LANL/Mustang";         # Name of the remote node.
my @return_component_list = ("DBINFO.EXP", "PPSTACK.UNCONV","PPSTACK.UNCONV.MASK","PPSTACK.UNCONV.VARIANCE",
			     "PPSTACK.UNCONV.EXP","PPSTACK.UNCONV.EXPWT","PPSTACK.UNCONV.EXPNUM","PSPHOT.OUTPUT",
			     "LOG.EXP", "PSASTRO.STATS",
			     "PPSTACK.CONFIG","PPSTACK.TARGET.PSF","PSPHOT.STACK.PSF.SAVE");
#,"PPSTACK.CONV.KERNEL");
# Look for programs we need
my $missing_tools;

my $remotetool = can_run('remotetool') or (warn "Can't find remotetool" and $missing_tools = 1);
my $stacktool    = can_run('stacktool') or (warn "Can't find warptool" and $missing_tools = 1);

if ($missing_tools) {
    warn("Can't find required tools.");
    exit($PS_EXIT_CONFIG_ERROR);
}

my ($remote_id,$stack_id,$camera,$dbname,$verbose,$path_base,$no_update);
GetOptions(
    'remote_id=s'    => \$remote_id,
    'camera|c=s'     => \$camera,
    'dbname|d=s'     => \$dbname,
    'path_base=s'    => \$path_base,
    'no_update'      => \$no_update,
    'verbose'        => \$verbose,
    ) or pod2usage( 2 );

pod2usage( -msg => "Unknown option: @ARGV", -exitval => 2 ) if @ARGV;
pod2usage( -msg => "Required options: --remote_id --camera --dbname --path_base", -exitval => 3) unless
    defined($remote_id) and
    defined($camera) and
    defined($path_base) and
    defined($dbname);
my $ipprc = PS::IPP::Config->new( $camera ) or my_die( "Unable to set up", $remote_id);

my $mdcParser = PS::IPP::Metadata::Config->new; # Parser for metadata config files

#
# Step 1: Get a list of the components that make up this remoteRun

# SHould this call listrun to ensure we're in state new?
my $rt_cmd = "$remotetool -listcomponent -remote_id $remote_id";
$rt_cmd   .= " -dbname $dbname " if defined($dbname);

my ( $success, $error_code, $full_buf, $stdout_buf, $stderr_buf ) =
    run(command => $rt_cmd, verbose => $verbose);
unless ($success) {
    $error_code = (($error_code >> 8) or $PS_EXIT_PROG_ERROR);
    &my_die("Unable to run remotetool to determine stage parameters.",
            $remote_id,$error_code);
}
my $compData = $mdcParser->parse(join "", @$stdout_buf) or
    &my_die("Unable to determine component information.",
            $remote_id,$PS_EXIT_PROG_ERROR);
my $compData2= parse_md_list($compData);

#
# Step 1b: Open output files
my $uri_command = $path_base . ".cmd";
my $uri_transfer= $path_base . ".transfer";
my $uri_check   = $path_base . ".check";
my $uri_config  = $path_base . ".config";
my $uri_generate= $path_base . ".generate";
my $uri_return  = $path_base . ".return";

my $disk_command = $ipprc->file_resolve($uri_command,1);
my $disk_transfer= $ipprc->file_resolve($uri_transfer,1);
my $disk_check   = $ipprc->file_resolve($uri_check,1);
my $disk_config  = $ipprc->file_resolve($uri_config,1);
my $disk_generate= $ipprc->file_resolve($uri_generate,1);
my $disk_return  = $ipprc->file_resolve($uri_return,1);

my (undef, $remote_config) = uri_convert($uri_config); # Needs to be done after we've created it.

open(TRANSFER, ">$disk_transfer")  || &my_die("Couldn't open file? $disk_transfer",$remote_id,$PS_EXIT_SYS_ERROR);
open(CHECK,    ">$disk_check")  || &my_die("Couldn't open file? $disk_check",$remote_id,$PS_EXIT_SYS_ERROR);
open(CONFIG,   ">$disk_config")  || &my_die("Couldn't open file? $disk_config",$remote_id,$PS_EXIT_SYS_ERROR);
open(GENERATE, ">$disk_generate") || &my_die("Couldn't open file? $disk_generate",$remote_id,$PS_EXIT_SYS_ERROR);
open(RETURN,   ">$disk_return") || &my_die("Couldn't open file? $disk_return", $remote_id, $PS_EXIT_SYS_ERROR);


#
# Step 2: Iterate over all componenets in this remote run.
my $job_index = 0;
my @pre_commands = ();
my @main_commands = ();
my @post_commands = ();
my %file_filter = ();


foreach my $compEntry (@$compData2) {
    my $stack_id = $compEntry->{stage_id};

# Get exposure level information from the stackRun we're working from.
    my $run_command = "$stacktool -tosum -stack_id $stack_id";
    $run_command .= " -dbname $dbname " if defined($dbname);
    ( $success, $error_code, $full_buf, $stdout_buf, $stderr_buf ) =
	run(command => $run_command, verbose => 0);
    unless ($success) {
	$error_code = (($error_code >> 8) or $PS_EXIT_PROG_ERROR);
	&my_die("Unable to run stacktool -tosum ",
		$stack_id,$error_code);
    }

    my $stackData = $mdcParser->parse(join "", @$stdout_buf) or
	&my_die("Unable to determine stack component information.",
		$stack_id,$PS_EXIT_PROG_ERROR);
    my $stackData2= parse_md_list($stackData);

    my $stack = ${ $stackData2 }[0];
    my $reduction = $stack->{reduction};
    $reduction = 'DEFAULT' unless defined($reduction);
    my $workdir   = $stack->{workdir};
    my $tess_id   = $stack->{tess_id};
    my $skycell_id= $stack->{skycell_id};

    my $ipp_outroot    = "${workdir}/${tess_id}/${skycell_id}/${tess_id}.${skycell_id}.stk.${stack_id}";    
    my $remote_outroot = uri_local_to_remote($ipp_outroot);    
    my $remote_outdir  = uri_local_to_remote("${workdir}/${tess_id}/${skycell_id}");
    # This actually returns all the individual stack/skyfiles that comprise this run.  Because consistency.
    my $command = "$stacktool -inputskyfile -stack_id $stack_id ";
    $command .= " -dbname $dbname " if defined($dbname);
    ( $success, $error_code, $full_buf, $stdout_buf, $stderr_buf ) =
	run(command => $command, verbose => 0);
    unless ($success) {
	$error_code = (($error_code >> 8) or $PS_EXIT_PROG_ERROR);
	&my_die("Unable to run stacktool -inputskyfile ",
		$stack_id,$error_code);
    }

    my $warpData = $mdcParser->parse(join "", @$stdout_buf) or
	&my_die("Unable to determine stack component information.",
		$stack_id,$PS_EXIT_PROG_ERROR);
    my $warpData2= parse_md_list($warpData);

#
# Step 3: Iterate over the sub-components
    my $warp_path_base_string = "";
    foreach my $warpEntry ( @{ $warpData2 } ) {
	my $warp_path_base = $warpEntry->{path_base};
	my $remote_path_base = uri_local_to_remote($warp_path_base);
	$warp_path_base_string .= " $remote_path_base ";

	# Append file names to transfer lists so we can recover if warps are deleted.
	uri_to_outputs($warp_path_base . ".fits");
	uri_to_outputs($warp_path_base . ".wt.fits");
	uri_to_outputs($warp_path_base . ".mask.fits");
	uri_to_outputs($warp_path_base . ".cmf");
#	uri_to_outputs($warp_path_base . ".psf");
    }
    my $mk_mdc_command = "mkdir -p $remote_outdir && sc_mk_stack_mdc.pl $warp_path_base_string > ${remote_outroot}.in.mdc";
    my $recipe_ppstack = $ipprc->reduction($reduction, 'STACK_PPSTACK'); # Recipe to use                                                                          
    my $ppstack_command  = " ppStack -input ${remote_outroot}.in.mdc ";
    $ppstack_command    .= " ${remote_outroot} -stats ${remote_outroot}.stats ";
    $ppstack_command    .= " -recipe PPSTACK $recipe_ppstack ";
    $ppstack_command    .= " -recipe PPSUB STACK_1DG -recipe PSPHOT STACK -recipe PPSTATS STACKSTATS ";
    $ppstack_command    .= " -stack-type DEEP_STACK -F PSPHOT.PSF.SAVE PSPHOT.PSF.SKY.SAVE ";
    $ppstack_command    .= " -F PSPHOT.OUTPUT PSPHOT.OUT.CMF.MEF -F PSPHOT.BACKMDL PSPHOT.BACKMDL.MEF ";
    $ppstack_command    .= " -F SOURCE.PLOT.MOMENTS SOURCE.PLOT.SKY.MOMENTS -F SOURCE.PLOT.PSFMODEL SOURCE.PLOT.SKY.PSFMODEL ";
    $ppstack_command    .= " -F SOURCE.PLOT.APRESID SOURCE.PLOT.SKY.APRESID -R PPSTACK.OUTPUT FITS.TYPE COMP_STACK ";
    $ppstack_command    .= " -R PPSTACK.OUTPUT.VARIANCE FITS.TYPE COMP_STACK  -R PPSTACK.UNCONV FITS.TYPE COMP_STACK ";
    $ppstack_command    .= " -R PPSTACK.UNCONV.VARIANCE FITS.TYPE COMP_STACK";
    $ppstack_command    .= " -tracedest ${remote_outroot}.trace -log ${remote_outroot}.log ";
    $ppstack_command    .= " -threads 4 ";
    $ppstack_command    .= " -dumpconfig ${remote_outroot}.mdc ";
    $ppstack_command    .= " -stack_id $stack_id -skycell_id $skycell_id -tess_id $tess_id ";

    my $post_cmd_echo = " echo -n \"stacktool  -addsumskyfile -stack_id $stack_id ";
    $post_cmd_echo   .= " -path_base $ipp_outroot -uri UNKNOWN ";
    $post_cmd_echo   .=  " -dbname $dbname " if defined $dbname;
    $post_cmd_echo   .= " -hostname $remote_hostname -dtime_script 0 \" > ${remote_outroot}.dbinfo ";
	
    my $post_cmd_SfM  = " ppStatsFromMetadata ${remote_outroot}.stats - STACK_SKYCELL >> ${remote_outroot}.dbinfo ";

    print CONFIG "${mk_mdc_command} && ${ppstack_command} && ${post_cmd_echo} && ${post_cmd_SfM} ";
    $job_index++;

    # Determine which output files need to be returned
    foreach my $component(@return_component_list) {
	my $filename = $ipprc->filename($component,$ipp_outroot);
	my ($ipp_disk, $remote_disk) = uri_to_outputs_for_return( $filename);
	my $remote_outroot_dir = dirname($ipp_disk);
	print CONFIG " && mkdir -p ${remote_root}/tmp/${remote_outroot_dir} && ln -sf $remote_disk ${remote_root}/tmp/${ipp_disk} && touch $remote_disk ";
    }

#     my $cleanup_cmd   = " (rm -f /scratch3/watersc1/sys_temp/${tess_id}.${skycell_id}.stk.${stack_id}.*conv.*.fits || true) ";    
#     print CONFIG " && ${cleanup_cmd} \n";

    print CONFIG "\n";
}
close(CONFIG);
close(TRANSFER);
close(CHECK);
close(RETURN);
close(GENERATE);

#
# Construct the moab command last, so we can use the job_index counter to estimate resources.  Somehow.
my $proc_need = $job_index * $threads;       # how many total processors do we need?
my $node_need = $proc_need / $proc_per_node; # this equals how many nodes?
my $time_need = $job_index * $job_cost;      # How many seconds will this take?

my $fill_factor = 0.8;  # This is the factor of how much of the time allocation we'd like to fill
my ($time_req,$node_req);
if ($node_need * $job_cost < $fill_factor * $min_nodes * $min_time) {
    $time_req = $min_time;
    $node_req = $min_nodes;
}
elsif ($node_need * $job_cost > $fill_factor * $max_nodes * $max_time) {
    $time_req = $max_time;
    $node_req = $max_nodes;
    print STDERR "You've requested the construction of a bundle that appears to need $node_need nodes and $job_cost time per job.  This exceeds the max limits ($max_nodes, $max_time).  Using those max values instead.  Good luck.\n";
}
else {
    $time_req = int(($node_need * $job_cost) / ($fill_factor * $max_nodes)) + 1;
    $node_req = int(($node_need * $job_cost) / ($fill_factor * $time_req)) + 1;
}

$time_req += 2;
$node_req *= 3;
open(COMMAND,  ">$disk_command") || &my_die("Couldn't open file? $disk_command",$remote_id,$PS_EXIT_SYS_ERROR);
print COMMAND "#!/bin/tcsh\n";
print COMMAND "##### Moab controll lines\n";
print COMMAND "#MSUB -l nodes=${node_req}:ppn=${proc_per_node},walltime=${time_req}:00:00\n"; ## CHECK RESOURCES
print COMMAND "#MSUB -j oe\n";
print COMMAND "#MSUB -V\n";
print COMMAND "#MSUB -o ${remote_root}stack.${remote_id}.out\n";
print COMMAND "date\n";
print COMMAND 'srun -n $SLURM_JOB_NUM_NODES -m cyclic -l /bin/hostname | sort -n | awk \'{printf "%s\n", $2}\' > hosts.${SLURM_JOB_ID}' . "\n";
print COMMAND "${remote_root}/stask_chip.py $remote_config " . 'hosts.${SLURM_JOB_ID} 3' . "\n";
print COMMAND "date\n";
close(COMMAND);

unless($no_update) {
    my $command = "remotetool -updaterun -remote_id $remote_id ";
    $command .= " -set_state pending ";
    $command .= " -dbname $dbname " if defined $dbname;

    system($command);
}


## Common SC routines

sub uri_convert { # (ipp_disk,remote_disk) = uri_convert(neb_uri);
    my $neb_uri = shift;
    my $ipp_disk= $ipprc->file_resolve( $neb_uri );
    my $remote_disk = $ipp_disk;
    
    unless(defined($ipp_disk)) {
        my_die();
    }

    $remote_disk =~ s%^.*/%%;   # Remove nebulous path
    $remote_disk =~ s%^\d+\.%%; # Remove ins_id
    $remote_disk =~ s%:%/%g;    # Replace colons with directories
    $remote_disk = "${remote_root}/${remote_disk}";
    return($ipp_disk,$remote_disk);
}

sub uri_convert_and_create { # (ipp_disk,remote_disk) = uri_convert_and_create(neb_uri); ipp_disk is created if it doesn't exist
    my $neb_uri = shift;
    my $ipp_disk= $ipprc->file_resolve( $neb_uri , 1);
    my $remote_disk = $ipp_disk;
    
    unless(defined($ipp_disk)) {
        my_die();
    }

    $remote_disk =~ s%^.*/%%;   # Remove nebulous path
    $remote_disk =~ s%^\d+\.%%; # Remove ins_id
    $remote_disk =~ s%:%/%g;    # Replace colons with directories
    $remote_disk = "${remote_root}/${remote_disk}";
    return($ipp_disk,$remote_disk);
}

sub uri_to_outputs { # (ipp_disk,remote_disk) = uri_to_output(neb_uri); Appends to TRANSFER and CHECK filehandles
    my $neb_uri = shift;
    my ($ipp_disk, $remote_disk) = uri_convert( $neb_uri );
    
    unless (exists($file_filter{$neb_uri})) {
        $file_filter{$neb_uri} = 1;
        print TRANSFER "$ipp_disk\n";
        print CHECK    "$remote_disk\n";
    }
    return($ipp_disk,$remote_disk);
}

sub uri_to_outputs_for_return { # (ipp_disk,remote_disk) = uri_to_outputs_for_return(neb_uri); create ipp_disk, append to RETURN and GENERATE
    my $neb_uri = shift;
    my ($ipp_disk, $remote_disk) = uri_convert_and_create( $neb_uri );
    unless (exists($file_filter{$neb_uri})) {
        $file_filter{$neb_uri} = 1;
        print RETURN "$ipp_disk\n";
        print GENERATE "$remote_disk\n";
    }
    return($ipp_disk,$remote_disk);
}

sub uri_local_to_remote { #(remote_uri) = uri_local_to_remote(local_neb_uri);
    # This needs to replace the nebulous tag with the remote root.
    my $local_uri = shift;
    $local_uri =~ s%^.*?/%%; # neb:/
    $local_uri =~ s%^.*?/%%; # /
    $local_uri =~ s%^.*?/%%; # @HOST@.0/
    my $remote_uri = "${remote_root}/" . $local_uri;

    return($remote_uri);
}
 
sub uri_remote_to_local { #(local_neb_uri) = uri_remote_to_local(remote_uri);
    # This needs to replace the remote root directory with the nebulous tag.
    my $remote_uri = shift;
    $remote_uri =~ s%${remote_root}%%;
    my $local_uri  = "neb:///" . $remote_uri;
    
    return($local_uri);
}


sub my_die { # exit with status; my_die(message,stage_id,exit_code,exit_status);
    my $msg = shift;
    my $id  = shift;
    my $exit_code = shift;
    my $exit_state = shift;

    $exit_code = $PS_EXIT_PROG_ERROR unless defined $exit_code;

    carp($msg);
    
    if (defined $id and not $no_update) {
        my $command = "remotetool -updaterun -remote_id $id";
        $command .= " -fault $exit_code " if defined $exit_code;
        $command .= " -set_state $exit_state " if defined $exit_state;
        $command .= " -dbname $dbname " if defined $dbname;

        system($command);
    }

    exit($exit_code);
}
