#!/usr/bin/env perl

#-----------------------------------------------------------------------------------------------
# case.setup - create the $caseroot/case.run script and user_nl_xxx component namelist mod files
#-----------------------------------------------------------------------------------------------

use strict;
use Cwd;
use English;
use Getopt::Long;
use IO::File;
use IO::Handle;
use File::Basename;
use File::Copy;


#-----------------------------------------------------------------------------------------------
sub usage {
    die <<EOF;

SYNOPSIS
     Creates Macros file for target machine if it does not exist
     Creates user_nl_xxx files for target components (and number of instances) if
        they do not exist
     Creates batch run script (case.run) for target machine

USAGE
     case.setup [options]
OPTIONS
     -help [or -h]        Print usage to STDOUT.
     -clean               Removes the batch run script for target machine
                          If the testmode argument is present then keep the test
			  test script if it is present - otherwise remove it
                          The user_nl_xxx and Macros files are never removed by case.setup -
                          you must remove them manually
     -testmode            Keeps the test script when the -clean argument is used
EOF
}

#-----------------------------------------------------------------------------------------------
# Set the directory that contains this script.
my $caseroot = getcwd();   # current working directory
$ENV{CASEROOT}=$caseroot;  # put this in environment
my $eol = "\n";

# Setting autoflush (an IO::Handle method) on STDOUT helps in debugging.  It forces any
# output be printed to STDOUT before the error messages start.
*STDOUT->autoflush();

# Parse command-line options.
my %opts = (loglevel=>"INFO");

GetOptions(
    "h|help"     => \$opts{'help'},
    "clean"      => \$opts{'clean'},
    "testmode"   => \$opts{'testmode'},
    "loglevel=s"     => \$opts{'loglevel'},
)  or usage();

# Give usage message.
usage() if $opts{'help'};

# Check for unparsed argumentss
if (@ARGV) {
    print "ERROR: unrecognized arguments: @ARGV\n";
    usage();
}

my $clean = 0;
if ($opts{'clean'}) {
    $clean = 1;
}

my $testmode = 0;
if ($opts{'testmode'}) {
    $testmode = 1;
}

chdir "$caseroot";
my $cimeroot  = `./xmlquery CIMEROOT  -value`;
my @dirs = ("$cimeroot/utils/perl5lib" );

unshift @INC, @dirs;
require Config::SetupTools;
require Task::TaskMaker;
require Batch::BatchMaker;
require Log::Log4perl;

my $level = Log::Log4perl::Level::to_priority($opts{loglevel});
Log::Log4perl->easy_init({level=>$level,
			  layout=>'%m%n'});

my $logger = Log::Log4perl::get_logger();

my %xmlvars=();
SetupTools::getxmlvars($caseroot, \%xmlvars);

# Check that $DIN_LOC_ROOT exists - and abort if not a namelist compare tests
my $din_loc_root = $xmlvars{'DIN_LOC_ROOT'};
my $testcase = $xmlvars{'TESTCASE'};
if (! -d $din_loc_root) {
    if ($testcase ne 'SBN') {
	$logger->logdie( "ERROR case.setup: inputdata root is not a directory: \"$din_loc_root\" ");
    }
}

# Check that userdefine settings are specified before expanding variable
my $fail_setup = 0;
my $strmacros;
my $strbuild;
my $strrun;
foreach my $attr (keys %xmlvars) {
    if ( $xmlvars{$attr} =~ m/USERDEFINED_required_macros/ ) {
	$strmacros.="ERROR: must set xml variable $attr to generate Macros file \n";
	$fail_setup = 1;
    }
    if ( $xmlvars{$attr} =~ m/USERDEFINED_required_build/ ) {
	$strbuild.="ERROR: must set xml variable $attr to build the model \n";
	$fail_setup = 1;
    }
    if ( $xmlvars{$attr} =~ m/USERDEFINED_required_run/ ) {
	$strrun.="ERROR: must set xml variable $attr to run the model \n";
	$fail_setup = 1;
    }
    $xmlvars{$attr} = SetupTools::expand_xml_var($xmlvars{$attr}, \%xmlvars);
}
if ($fail_setup) {
    my $outstr;
    $outstr = "$strmacros\n" if(defined($strmacros));
    $outstr .= "$strbuild\n" if(defined($strbuild));
    $outstr .= "$strrun\n" if(defined($strrun));
    $outstr .= "Correct above and issue case.setup again ";
    $logger->logdie($outstr);
}

#-----------------------------------------------------------------------------------------------
# Create batch script
#-----------------------------------------------------------------------------------------------

if (! $clean ) {

    my $mach = $xmlvars{'MACH'};
    if (! defined $mach ) {
	$logger->logdie("ERROR case.setup: xml variable MACH is not set ");
    }
    my $case = $xmlvars{'CASE'};
    if (! defined $case ) {
	$logger->logdie("ERROR case.setup: xml variable CASE is not set ");
    }

    #--------------------------------------------------------------------------
    # Create Macros file only if it does not exist
    #--------------------------------------------------------------------------
    if (!-f "./Macros") {
	$logger->info( "Creating Macros file for $mach");
	SetupTools::set_compiler("$xmlvars{OS}",
				 "$xmlvars{MACHDIR}/config_compilers.xml",
				 "$xmlvars{COMPILER}",
				 "$xmlvars{MACH}",
				 "$xmlvars{MPILIB}", 'Macros'  );
    } else {
	$logger->info( "Macros script already created ...skipping");
    }
    # Set tasks to 1 if mpi-serial library
    if($xmlvars{MPILIB} eq "mpi-serial"){
	foreach my $attr (keys %xmlvars){
	    if($attr =~ /NTASKS_/){
		$xmlvars{$attr}=1;
		my $sysmod = "./xmlchange -noecho -file env_mach_pes.xml -id $attr -val 1";
		system($sysmod) == 0 or $logger->logdie("ERROR case.setup: $sysmod failed: $?");

	    }
	}
    }

    my $ninst_fail = 0;
    my $build_threaded = $xmlvars{BUILD_THREADED};
    foreach my $comp (qw(ATM LND ROF OCN ICE GLC WAV ESP)){
	my $ninst  = "NINST_".$comp;
	my $ntasks = "NTASKS_".$comp;
	my $nthrds = "NTHRDS_".$comp;
	if ($xmlvars{$ninst} > $xmlvars{$ntasks}) {
	    if($xmlvars{$ntasks}==1){
		system("./xmlchange -noecho $ntasks=$xmlvars{$ninst}");
		$xmlvars{$ntasks}=$xmlvars{$ninst};
	    }else{
		$logger->logdie("ERROR case.setup: $comp NINST value greater than $comp NTASKS");
	    }
	}
	if($xmlvars{$nthrds} > 1){
	    $build_threaded = "TRUE";
	}
    }

    if ($build_threaded eq "TRUE" and $xmlvars{"COMPILER"} eq "nag") {
        $logger->logdie("ERROR cesm_setup: it is not possible to run with OpenMP if using the NAG Fortran compiler");
    }

    my $sysmod = "./xmlchange -noecho BUILD_THREADED=$build_threaded";
    system($sysmod) == 0 or $logger->logdie("ERROR case.setup: $sysmod failed: $?");

    if (-f "$caseroot/case.run") {
	$logger->info( "Machine/Decomp/Pes configuration has already been done ...skipping ");

    } else {
	_check_pelayouts_require_rebuild();
	unlink("LockedFiles/env_build.xml") if(-e "LockedFiles/env_build.xml");
	my $sysmod = "./Tools/check_lockedfiles -cimeroot $cimeroot";
	system($sysmod) == 0 or $logger->logdie("ERROR case.setup: $sysmod failed: $?");

	my $pestot = `$caseroot/Tools/taskmaker.pl -sumonly`;
	$sysmod = "./xmlchange -noecho -file env_mach_pes.xml -id TOTALPES -val $pestot";
	system($sysmod) == 0 or $logger->logdie("ERROR case.setup: $sysmod failed: $?");
	wait;

        # Compute cost based on PE count
	my $pval = 1;
	my $pcnt = 0;
	while ($pval < $pestot) {
	    $pval = $pval * 2;
	    #$pcnt = $pcnt + 10;      # (perfect scaling)
	    $pcnt = $pcnt + 6;       # (scaling like sqrt(6/10))
	}
	my $pcost = 3 - int($pcnt / 10); # (3 is 64 with 6)

        # Compute cost based on DEBUG
	my $dcost = 0;
	if ($xmlvars{'DEBUG'} == "TRUE") {$dcost = 3;}

        # Compute cost based on run length
        # For simplicity, we use a heuristic just based on STOP_OPTION (not considering
        # STOP_N), and only deal with options longer than ndays
        my $lcost = 0;
        if ($xmlvars{'STOP_OPTION'} =~ /nmonth/) {
           # N months costs 30x as much as N days; since cost is based on log-base-2, add 5
           $lcost = 5;
        }
        elsif ($xmlvars{'STOP_OPTION'} =~ /nyear/) {
           # N years costs 365x as much as N days; since cost is based on log-base-2, add 9
           $lcost = 9;
        }

	my $CCSM_CCOST = $xmlvars{'CCSM_CCOST'};
	my $CCSM_GCOST = $xmlvars{'CCSM_GCOST'};
	my $CCSM_TCOST = $xmlvars{'CCSM_TCOST'};
	my $CCSM_MCOST = $xmlvars{'CCSM_MCOST'};
	my $estcost = $CCSM_CCOST + $CCSM_GCOST + $CCSM_MCOST + $CCSM_TCOST + $pcost + $dcost + $lcost;

	my $sysmod = "./xmlchange -noecho -file env_mach_pes.xml -id TOTALPES -val $pestot";
	system($sysmod) == 0 or $logger->logdie("ERROR case.setup: $sysmod failed: $?");

	$sysmod = "./xmlchange -noecho -file env_mach_pes.xml -id CCSM_PCOST -val $pcost";
	system($sysmod) == 0 or $logger->logdie("ERROR case.setup: $sysmod failed: $?");

	$sysmod = "./xmlchange -noecho -file env_mach_pes.xml -id CCSM_ESTCOST -val $estcost";
	system($sysmod) == 0 or $logger->logdie("ERROR case.setup: $sysmod failed: $?");

	#--------------------------------------------------------------
	# create batch file
	#--------------------------------------------------------------
	$logger->info( "Creating batch script case.run");
	# Use BatchFactory to get the appropriate instance of a BatchMaker,
	# use it to create our batch scripts
	my $batchmaker = Batch::BatchFactory::getBatchMaker( caseroot => $xmlvars{'CASEROOT'},
							     cimeroot => $cimeroot,
							     case     => $xmlvars{'CASE'},
							     mpilib   => $xmlvars{'MPILIB'},
							     machroot => $xmlvars{'MACHDIR'},
							     machine  => $xmlvars{'MACH'},
							     compiler => $xmlvars{'COMPILER'},
							     threaded => $build_threaded,
							     job => 'run'
	                                                     );



	my $inputbatchscript  = "$xmlvars{'MACHDIR'}/template.cesmrun";
	my $outputbatchscript = "$xmlvars{'CASEROOT'}/case.run";

	$batchmaker->makeBatchScript($inputbatchscript, $outputbatchscript);
	$logger->info("Creating batch script case.st_archive");
	$batchmaker->overrideNodeCount(1);
	$batchmaker->set({job=>'st_archive'});
	$batchmaker->makeBatchScript("$xmlvars{'MACHDIR'}/template.st_archive", "$xmlvars{'CASEROOT'}/case.st_archive");

	$logger->info("Creating batch script case.lt_archive");
	$batchmaker->overrideNodeCount(1);
	$batchmaker->set({job=>'lt_archive'});
	$batchmaker->makeBatchScript("$xmlvars{'MACHDIR'}/template.lt_archive", "$xmlvars{'CASEROOT'}/case.lt_archive");

	# Make a copy of env_mach_pes.xml in order to be able
	# to check that it does not change once case.setup is invoked
	if (! -e "LockedFiles/env_mach_pes.xml") {
	    $sysmod = "cp env_mach_pes.xml  LockedFiles/env_mach_pes.xml";
	    system($sysmod) == 0 or $logger->logdie("ERROR case.setup: $sysmod failed: $?");
	    $logger->info("Locking file env_mach_pes.xml ");
	}
    }

    #--------------------------------------------------------------------------
    # Create user_nl files for the required number of instances
    #--------------------------------------------------------------------------
    if (!-f "./user_nl_cpl") {
	$logger->info("Creating user_nl_xxx files for components and cpl");
    }
    my @models = qw( ATM LND ICE OCN GLC ROF WAV ESP );
    # loop over models
    foreach my $model (@models) {
	# what is actual component associated with each model
	my $comp = $xmlvars{"COMP_${model}"};
	_build_usernl_files($caseroot, $cimeroot, $model, $comp);
	if ($comp eq "cism")  {
	    my $sysmod = "$cimeroot/../components/cism/cime_config/cism.template $caseroot";
	    system($sysmod) == 0 or $logger->logdie("ERROR case.setup: $sysmod failed: $?");
	}
    }
    _build_usernl_files($caseroot, $cimeroot, "drv", "cpl");

    #--------------------------------------------------------------------------
    # Run preview namelists for scripts
    #--------------------------------------------------------------------------
    $logger->info("Running preview_namelist script with loglevel $opts{loglevel}");
    system(" $caseroot/preview_namelists -loglevel $opts{loglevel} ");
    if($?){
	$logger->logdie("ERROR: $caseroot/preview_namelists failed: $?")
    }
    $logger->info("See ./CaseDoc for component namelists ");
    $logger->info("If an old case build already exists, might want to run case.clean_build before building ");

    #--------------------------------------------------------------------------
    # Create test script if appropriate
    #--------------------------------------------------------------------------
    if( -e "$caseroot/env_test.xml") {
	if ( -e "$caseroot/case.test") {
	    # do nothing
	} else {
	    $logger->info("Starting testcase.setup ");
	    my $sysmod = "./testcase.setup -caseroot $caseroot";
	    system($sysmod) == 0 or $logger->logdie("ERROR: $sysmod failed: $?");
	    $logger->info("Finished testcase.setup $caseroot");
	}
    }

    my $fh = new IO::File;
    $fh->open(">>CaseStatus") or $logger->logdie("can't open file: CaseStatus");
    my $sdate = `date +"%Y-%m-%d %H:%M:%S"`;
    print $fh "case.setup $sdate $eol";
    $fh->close();

}

#-----------------------------------------------------------------------------------------------
# Clean batch script
#-----------------------------------------------------------------------------------------------

if ($clean) {

    my $case = $xmlvars{'CASE'};
    if (!-f "$caseroot/case.run" ) {
	$logger->info("clean option has already been invoked ...skipping ");
	exit;
    }

    my $sysmod;
    my $id = `date +%y%m%d-%H%M%S`;
    my $backupdir = "PESetupHist/b.${id}";
    if (!-d ${backupdir}) {
	$sysmod = "mkdir -p ${backupdir}";
	system($sysmod) == 0 or $logger->logdie("ERROR case.setup: $sysmod failed: $?");
    }
    # back up relevant files
    $sysmod = "cp case.run env_build.xml env_mach_pes.xml Macros* ${backupdir}";
    system($sysmod) == 0 or $logger->logdie("ERROR case.setup: $sysmod failed: $?");

    # remove relevant files from $caseroot
    $sysmod = "rm ./case.run";
    system($sysmod) == 0 or $logger->logdie("ERROR case.setup: $sysmod failed: $?");

    # only do the following if are NOT in testmode
    if (! $testmode)  {
	# rebuild the models (even on restart)
	$sysmod = "./xmlchange -noecho -file env_build.xml -id BUILD_COMPLETE -val FALSE";
	system($sysmod) == 0 or $logger->logdie("ERROR case.setup: $sysmod failed: $?");

	# backup and then clean test script
	if( -e "$caseroot/case.test") {
	    $sysmod = "cp case.test ${backupdir}";
	    system($sysmod) == 0 or $logger->logdie("ERROR case.setup: $sysmod failed: $?");

	    $sysmod = "rm ./case.test";
	    system($sysmod) == 0 or $logger->logdie("ERROR case.setup: $sysmod failed: $?");
	    $logger->info("Successfully cleaned test script case.test ");
	}
    }

    $logger->info("Successfully cleaned batch script case.run ");
    $logger->info("Some files have been saved to ${backupdir}");

    my $fh = new IO::File;
    $fh->open(">>CaseStatus") or $logger->logdie("can't open file: CaseStatus");
    my $sdate = `date +"%Y-%m-%d %H:%M:%S"`;
    print $fh "case.setup -clean $sdate $eol";
    $fh->close();

}

sub _check_pelayouts_require_rebuild {
    my ($xmlvars) = @_;
    if ( -e "LockedFiles/env_mach_pes.xml") {
	# Look to see if $comp_PE_CHANGE_REQUIRES_REBUILD is defined for any component
	foreach my $comp (qw(ATM LND ROF OCN ICE GLC WAV)){
	    if(defined $xmlvars{$comp."_PE_CHANGE_REQUIRES_REBUILD"} &&
	       $xmlvars{$comp."_PE_CHANGE_REQUIRES_REBUILD"} == "true"){
		# Changing these values in env_mach_pes.xml will force you to clean the corresponding component
		my $old_tasks = `./xmlquery NTASKS_$comp -file LockedFiles/env_mach_pes.xml -value`;
		my $old_threads = `./xmlquery NTHRDS_$comp -file LockedFiles/env_mach_pes.xml -value`;
		my $old_inst = `./xmlquery NINST_$comp -file LockedFiles/env_mach_pes.xml -value`;
		my $new_tasks = `./xmlquery NTASKS_$comp -file env_mach_pes.xml -value`;
		my $new_threads = `./xmlquery NTHRDS_$comp -file env_mach_pes.xml -value`;
		my $new_inst = `./xmlquery NINST_$comp -file env_mach_pes.xml -value`;
		if($old_tasks != $new_tasks || $old_threads != $new_threads || $old_inst != $new_inst){
		    $logger->warn("$comp pe change requires clean build");
		    my $cleanflag = lc substr($comp,0,1);
		    system("./case.clean_build -$cleanflag");
		}
	    }
	}
	unlink("LockedFiles/env_mach_pes.xml");
    }

}



#-----------------------------------------------------------------------------------------------
sub _build_usernl_files {
    # Create user_nl_xxx files
    my ($caseroot, $cimeroot, $model, $comp) = @_;

    my $model = uc($model);
    my $file = `./xmlquery CONFIG_${model}_FILE -value`;
    my $dir  = dirname($file);

    (-d $dir) or $logger->logdie("ERROR _build_user_nl_files: cannot find cime_config directory $dir for component $comp");

    if ($comp eq 'cpl') {
	if ( ! -f "$caseroot/user_nl_cpl" ) {
	    my $sysmod = "cp $dir/user_nl_cpl $caseroot/user_nl_cpl";
	    system($sysmod) == 0 or $logger->logdie("ERROR build_usernl_files: $sysmod failed: $?");
	}
    } else {
	my $NINST = `./xmlquery NINST_${model} -value`;
	if ( -f "${dir}/user_nl_${comp}") {
	    if ($NINST > 1) {
		my $inst_string;
		my $inst_counter = 1;
		while ($inst_counter <= $NINST) {
		    $inst_string = `printf _%04d $inst_counter`;
		    if ( ! -f "$caseroot/user_nl_${comp}${inst_string}" ) {
			my $sysmod = "cp $dir/user_nl_${comp} $caseroot/user_nl_${comp}${inst_string}";
			system($sysmod) == 0 or $logger->logdie("ERROR build_usernl_files: $sysmod failed: $?");
		    }
		    $inst_counter = $inst_counter + 1;
		}
	    } else {
		if ( ! -f "$caseroot/user_nl_${comp}" ) {
		    (-f "${dir}/user_nl_${comp}") or $logger->logdie("ERROR _build_usernl_files: $dir/user_nl_${comp} does not exist");
		    my $sysmod = "cp $dir/user_nl_${comp} $caseroot/user_nl_${comp}";
		    system($sysmod) == 0 or $logger->logdie("ERROR _build_usernl_files: $sysmod failed: $?");
		}
	    }
	}
    }
}