#!/usr/bin/env perl
#=======================================================================
# name - mk_Restarts
# purpose - wrapper script to run programs which regrid surface restarts
#=======================================================================
use strict;
use warnings;
use FindBin qw($Bin);
use lib ("$Bin");
use Cwd qw(getcwd);

# global variables
#-----------------
my ($saltwater, $openwater, $seaice, $lake, $landice, $catch, $catchcn, $route);
my ($surflay, $rsttime, $grpID, $numtasks, $walltime, $rescale, $qos, $partition, $yyyymm);
my ($mk_catch_j, $mk_catch_log, $weminIN, $weminOUT, $weminDFLT);
my ($zoom);

# mk_catchcn job and log file names
#----------------------------------
$mk_catch_j   = "mk_catch.j";
$mk_catch_log = "mk_catch.log";

# main program
#-------------
{
    my ($cmd, $line, $pid);

    init();

    #---------------------------
    # catch and catchcn restarts
    #---------------------------
    if ($catch or $catchcn) {
        write_mk_catch_j() unless -e $mk_catch_j;

        # run interactively if already on interactive job nodes
        #------------------------------------------------------
        if (-x $mk_catch_j) {
            $cmd = "./$mk_catch_j";
            system_($cmd);
        }
        else {
            $cmd = "sbatch -W $mk_catch_j";
            print "$cmd\n";
            chomp($line = `$cmd`);
            $pid = (split /\s+/, $line)[-1];
        }
    }

    #------------------
    # saltwater restart
    #------------------
    if ($saltwater) {
        $cmd = "$Bin/mk_LakeLandiceSaltRestarts "
            .       "OutData/\*.til "
            .       "InData/\*.til "
            .       "InData/\*saltwater_internal_rst\* 0 $zoom";
        system_($cmd);
    }

    if ($openwater) {
        $cmd = "$Bin/mk_LakeLandiceSaltRestarts "
            .       "OutData/\*.til "
            .       "InData/\*.til "
            .       "InData/\*openwater_internal_rst\* 0 $zoom";
        system_($cmd);
    }

    if ($seaice) {
        $cmd = "$Bin/mk_LakeLandiceSaltRestarts "
            .       "OutData/\*.til "
            .       "InData/\*.til "
            .       "InData/\*seaicethermo_internal_rst\* 0 $zoom";
        system_($cmd);
    }

    #-------------
    # lake restart
    #-------------
    if ($lake) {
        $cmd = "$Bin/mk_LakeLandiceSaltRestarts "
            .       "OutData/\*.til "
            .       "InData/\*.til "
            .       "InData/\*lake_internal_rst\* 19 $zoom";
        system_($cmd);
    }

    #----------------
    # landice restart
    #----------------
    if ($landice) {
        $cmd = "$Bin/mk_LakeLandiceSaltRestarts "
            .       "OutData/\*.til "
            .       "InData/\*.til "
            .       "InData/\*landice_internal_rst\* 20 $zoom";
        system_($cmd);
    }

    #--------------
    # route restart
    #--------------
    if ($route) {
        $cmd = "$Bin/mk_RouteRestarts OutData/\*.til $yyyymm";
        system_($cmd);
    }

    wait_for_pid($pid) if $pid;
}

#=======================================================================
# name - init
# purpose - get runtime flags to determine which restarts to regrid
#=======================================================================
sub init {
    use Getopt::Long;
    my $help;
    $| = 1;  # flush buffer after each output operation

    GetOptions( "saltwater" => \$saltwater,
                "openwater" => \$openwater,
                "seaice"    => \$seaice,
                "lake"      => \$lake,
                "landice"   => \$landice,
                "catch"     => \$catch,
                "catchcn"   => \$catchcn,
                "wemin=i"   => \$weminIN,
                "wemout=i"  => \$weminOUT,
                "route"     => \$route,

                "surflay=i" => \$surflay,
                "rsttime=i" => \$rsttime,
                "grpID=s"   => \$grpID,

                "ntasks=i"  => \$numtasks,
                "walltime=s"=> \$walltime,
                "rescale"   => \$rescale,
                "qos=s"     => \$qos,
                "partition=s" => \$partition,
                "zoom=i"    => \$zoom,
                "h|help"    => \$help );
    # defaults
    #---------
    $rsttime = 0 unless $rsttime;
    $catch = 0 unless $catch;
    $catchcn = 0 unless $catchcn;
    $rescale = 0 unless $rescale;
    $weminDFLT = 26;
    $weminIN = $weminDFLT unless defined($weminIN);
    $weminOUT = $weminDFLT unless defined($weminOUT);
    $zoom = 8 unless $zoom;

    usage() if $help;

    # error if no restart specified
    #------------------------------
    die "Error. No restart specified;"
        unless $saltwater or $lake or $landice or $catch or $catchcn;

    # rsttime and grpID values are needed for catchcn
    #----------------------------------------------
    if ($catchcn) {
        die "Error. Must specify rsttime for catchcn;" unless $rsttime;
        die "Error. rsttime not in yyyymmddhh format: $rsttime;"
            unless $rsttime =~ m/^\d{10}$/;
    }
    if ($catch or $catchcn) {
        unless ($grpID) {
            $grpID = `$Bin/getsponsor.pl -d`;
            print "Using default grpID = $grpID\n";
        }
        unless ($walltime) { $walltime = "1:00:00" }
        unless ($numtasks) { $numtasks = 84 }
        $qos = "" unless $qos;
        $partition = "" unless $partition;
    }

    # rsttime value is needed for route
    #----------------------------------
    if ($route) {
        die "Error. Must specify rsttime for route;" unless $rsttime;
        die "Error. Cannot extract yyyymm from rsttime: $rsttime"
            unless $rsttime =~ m/^\d{6,}$/;
        $yyyymm = $1 if $rsttime =~ /^(\d{6})/;
    }
}

#=======================================================================
# name - write_mk_catch_j
# purpose - write job file to make catch and/or catchcn restart
#=======================================================================
sub write_mk_catch_j {
    my ($grouplist, $cwd, $QOSline, $PARTline, $FH);

    $grouplist = "";
    $grouplist = "SBATCH --account=$grpID" if $grpID;

    $cwd = getcwd;

    $QOSline = "";
    if ($qos) {
        $QOSline = "SBATCH --qos=$qos";
        if ($qos eq "debug") {
            $QOSline = "" unless $numtasks <= 532 and $walltime le "1:00:00";
        }
    }
    $PARTline = "";
    if ($partition) {
        $PARTline = "SBATCH --partition=$partition";
    }
    print("\nWriting jobscript: $mk_catch_j\n");
    open CNj, ">> $mk_catch_j" or die "Error opening $mk_catch_j: $!";

    $FH = select;
    select CNj;

    print <<"EOF";
#!/bin/csh -f
#$grouplist
#SBATCH --ntasks=$numtasks
#SBATCH --time=$walltime
#SBATCH --job-name=catchcnj
###SBATCH --constraint=hasw
#SBATCH --output=$cwd/$mk_catch_log
#$QOSline
#$PARTline

source $Bin/g5_modules
set echo

#limit stacksize unlimited
unlimit

set catchFLG   = $catch
set catchcnFLG = $catchcn
set weminIN  = $weminIN
set weminOUT = $weminOUT
set rescaleFLG = $rescale

set numtasks   = $numtasks
set rsttime    = $rsttime
set surflay    = $surflay
set zoom       = $zoom

set esma_mpirun_X = ( $Bin/esma_mpirun -np \$numtasks )
set mk_CatchRestarts_X   = ( \$esma_mpirun_X $Bin/mk_CatchRestarts )
set mk_CatchCNRestarts_X = ( \$esma_mpirun_X $Bin/mk_CatchCNRestarts )
set Scale_Catch_X   = $Bin/Scale_Catch
set Scale_CatchCN_X = $Bin/Scale_CatchCN

set OUT_til   = OutData/\*.til
set IN_til    = InData/\*.til

if (\$catchFLG) then
    set catchIN = InData/\*catch_internal_rst\*
    set params = ( \$OUT_til \$IN_til \$catchIN \$surflay \$zoom )
    \$mk_CatchRestarts_X \$params

    if (\$rescaleFLG) then
        set catch_regrid = OutData/\$catchIN:t
        set catch_scaled = \${catch_regrid}.scaled
        set params = ( \$catchIN \$catch_regrid \$catch_scaled \$surflay )
        set params = ( \$params \$weminIN \$weminOUT )
        \$Scale_Catch_X \$params

        mv \$catch_regrid \${catch_regrid}.1
        mv \$catch_scaled \$catch_regrid
    endif
endif

if (\$catchcnFLG) then
    set catchcnIN = InData/\*catchcn_internal_rst\*
    set params = ( \$OUT_til \$IN_til \$catchcnIN \$surflay \$rsttime \$zoom )
    \$mk_CatchCNRestarts_X \$params

    if (\$rescaleFLG) then
        set catchcn_regrid = OutData/\$catchcnIN:t
        set catchcn_scaled = \${catchcn_regrid}.scaled
        set params = ( \$catchcnIN \$catchcn_regrid \$catchcn_scaled \$surflay )
        set params = ( \$params \$weminIN \$weminOUT )
        \$Scale_CatchCN_X \$params

        mv \$catchcn_regrid \${catchcn_regrid}.1
        mv \$catchcn_scaled \$catchcn_regrid
    endif
endif
exit
EOF
;
    close CNj;
    select $FH;
    chmod 0755, $mk_catch_j if $ENV{"SLURM_JOBID"};
}

#=======================================================================
# name - system_
# purpose - wrapper for perl system command
#=======================================================================
sub system_ {
    my $cmd = shift @_;
    print "\n$cmd\n";
    die "Error: $!;" if system($cmd);
}

#=======================================================================
# name - wait_for_pid
# purpose - wait for batch job to finish
#
# input parameter
# => $pid: process ID of batch job to wait for
#=======================================================================
sub wait_for_pid {
    my ($pid, $first, %found, $line, $id);
    $pid = shift @_;
    return unless $pid;

    $first = 1;
    while (1) {
        %found = ();
        #--foreach $line (`qstat | grep $ENV{"USER"}`) {
        foreach $line (`squeue | grep $ENV{"USER"}`) {
            $line =~ s/^\s+//;
            $id = (split /\s+/, $line)[0];
            $found{$id} = 1;
        }
        last unless $found{$pid};
        print "\nWaiting for job $pid to finish\n" if $first;
        $first = 0;
        sleep 10;
    }
    print "Job $pid is DONE\n\n" unless $first;
}

#=======================================================================
# name - usage
# purpose - print usage information
#=======================================================================
sub usage {
    use File::Basename ("basename");
    my $name = basename $0;
    print <<"EOF";

usage $name [-saltwater] [-lake] [-landice] [-catch] [-h]

option flags
   -saltwater       regrid saltwater internal restart
   -lake            regrid lake internal restart
   -landice         regrid landice internal restart
   -catch           regrid catchment internal restart
   -catchcn         regrid catchment CN internal restart
   -wemin weminIN   minimum snow water equivalent threshold for input catch/cn [$weminDFLT]
   -wemout weminOUT minimum snow water equivalent threshold for output catch/cn [$weminDFLT]
   -route           create the route internal restart

   -surflay n       number of surface layers (catch & catchcn)
   -rsttime n10     restart time in format, yyyymmddhh (catchcn) or yyyymm (route)
   -grpID grpID     group ID for batch submittal (catchcn)

   -ntasks nt       number of tasks to assign to catchcn batch job [112]
   -walltime wt     walltime in format \"hh:mm:ss\" for catchcn batch job [1:00:00]
   -rescale
   -qos val         use \"SBATCH --qos=val directive\" for batch jobs;
                    \"-qos debug\" will not work unless these conditions are met
                         -> numtasks <= 532
                         -> walltime le \"1:00:00\"
   -partition val   use \"SBATCH --partition=val directive\" for batch jobs
   -zoom n          zoom value to send to land regridding codes [8]
   -h               print usage information

EOF
exit;
}
