"""
Developer    : Michael O'Donnell, U.S. Geological Survey
Date         : 09/15/2014
Python ver.  : Developed with Python 2.6.x and 2.7.x (not tested with 3.x)

Objective: This is a Python script (source), which is compiled into an executable,
  that runs in a HTCondor (grid computing middleware). This script executes the
  individual Monte Carlo replications of Apex ST-STM command line software. The
  executable generated from this script (does not require Python on the HTCondor
  client) is called from the HTCondor submit file that is generated with
  ST_Sim_HTCondor_setup.py.

################################################################################
SOFTWARE DISCLAIMER AND LICENSING INFORMATION

Although this program has been used by the U.S. Geological Survey (USGS), no
warranty, expressed or implied, is made by the USGS or the U.S. Government as
to the accuracy and functioning of the program and related program material nor
shall the fact of distribution constitute any such warranty, and no responsibility
is assumed by the USGS in connection therewith.


U.S. GEOLOGICAL SURVEY OPEN SOURCE AGREEMENT VERSION 1.0

THIS OPEN SOURCE AGREEMENT ("AGREEMENT") DEFINES THE RIGHTS OF USE, REPRODUCTION,
DISTRIBUTION, MODIFICATION AND REDISTRIBUTION OF CERTAIN COMPUTER SOFTWARE ORIGINALLY
RELEASED BY THE UNITED STATES GOVERNMENT AS REPRESENTED BY THE GOVERNMENT AGENCY
LISTED BELOW ("GOVERNMENT AGENCY"). THE UNITED STATES GOVERNMENT, AS REPRESENTED
BY GOVERNMENT AGENCY, IS AN INTENDED THIRD-PARTY BENEFICIARY OF ALL SUBSEQUENT
DISTRIBUTIONS OR REDISTRIBUTIONS OF THE SUBJECT SOFTWARE. ANYONE WHO USES, REPRODUCES,
DISTRIBUTES, MODIFIES OR REDISTRIBUTES THE SUBJECT SOFTWARE, AS DEFINED HEREIN,
OR ANY PART THEREOF, IS, BY THAT ACTION, ACCEPTING IN FULL THE RESPONSIBILITIES
AND OBLIGATIONS CONTAINED IN THIS AGREEMENT.
################################################################################
"""

# Standard python modules
import sys, os, subprocess, time, zipfile, traceback
import threading

# Passed arguments to compiled script -------------------------------------------
# Path to ST-SIM software executables
STSim_Lib = sys.argv[1]

# Path to .ssim database created for each Monte Carlo by the ST-STM Splie.exe
#   commandline software
ScenarioLib = sys.argv[2]

# This is a Boolean signifying whether data is transfered to the client as a zip
#   with HTCondor software. The user specifies this need when using
#   ST_Sim_HTCondor_setup.py
TransferZip = eval(sys.argv[3]) # Boolean

# Set this to true to investigate folders and files transfered to HTCondor node
QAQC = False

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
class ReaderThread(threading.Thread):
    """
    Class used to capture stdout and stderr after running a subprocess.popen()
      and wrapping in a thread. This method is required to get a live stream of the
      output, which is useful for debugging.
    """
    def __init__(self, stream):
        threading.Thread.__init__(self)
        self.stream = stream

    def run(self):
        while True:
            line = self.stream.readline()
            if len(line) == 0:
                break
            print "\t", line,


def main():
    """
    This is the main routine that runs the ST-SIM models on the HTCondor client.
    """
    # Get start info
    Get_Start_Info()

    print "\nPassed arguments:"
    print "ST-Sim Library: " + STSim_Lib
    print "Scenario Folder: " + ScenarioLib
    print "Unzipping..." + str(TransferZip)
    print "Executing ApexRms.SyncroSim.Run.exe..."

    # Workspace setup
    HTcondorWS = os.getcwd()
    print "\nWorking directory: " + HTcondorWS

    if TransferZip == True:
        print "\nUnzip data..."
        os.chdir(HTcondorWS)
        for iZip in [os.path.basename(STSim_Lib), os.path.basename(ScenarioLib)]:
            if os.path.exists(iZip):
                print "\tUnzipping: " + iZip
                zf = zipfile.ZipFile(iZip, "r")
                # Extract Monte Carlo data to root
                if iZip == os.path.basename(ScenarioLib):
                    ZipFld = os.path.splitext(iZip)[0]
                    zf.extractall(os.path.dirname(ZipFld))
                # Extract software to zip
                else:
                    zf.extractall(os.path.splitext(iZip)[0])
            else:
                print "Cannot locate file to unzip..." + iZip
                sys.exit(1)
        time.sleep(2)


    print "\nRun the ST-SIM job..."
    # Determine where the database is located
    # Determine where executable is located for STM library
    if TransferZip == True:
        STSim_Lib2 = os.path.basename(STSim_Lib)
        ScenarioLib2 = os.path.basename(ScenarioLib)
        ExeDir = os.path.join(HTcondorWS, STSim_Lib2.replace(".zip", ""))
        SimFld = os.path.join(HTcondorWS, ScenarioLib2.replace(".zip", ""))
        SIMDB_base = os.path.basename(SimFld)
        SimDB = os.path.join(HTcondorWS, SIMDB_base.replace(".input", ""))
    else:
        ExeDir = STSim_Lib
        SimFld = ScenarioLib
        SimDB = SimFld.replace(".input", "")

    # Print a list of directories for QAQC
    if QAQC:
        print "\tExecuatbale directory: " + ExeDir
        print "\tFiles:"
        for root, dirs, files in os.walk(ExeDir):
            for idir in dirs:
                print "\t\t" + idir
                for iFile in os.listdir(os.path.join(root, idir)):
                    print "\t\t\t" + iFile
        for iFile in os.listdir(ExeDir):
            print "\t\t" + iFile
        #
        print "\n\tSTM data directory: " + SimFld
        print "\tFiles:"
        for root, dirs, files in os.walk(SimFld):
            for idir in dirs:
                print "\t\t" + idir
                for iFile in os.listdir(os.path.join(root, idir)):
                    print "\t\t\t" + iFile
        for iFile in os.listdir(SimFld):
            print "\t\t" + iFile

    # Check if executable exists
    if not os.path.exists(os.path.join(ExeDir, "ApexRms.SyncroSim.Run.exe")):
        print "\tCannot locate exe: " + os.path.join(ExeDir, "ApexRms.SyncroSim.Run.exe")
        sys.exit(1)
    else:
        print "\tLocated exe: " + os.path.join(ExeDir, "ApexRms.SyncroSim.Run.exe")

    # Check if library exists
    if not os.path.exists(SimDB):
        print "\tCannot locate .ssim library: " + SimDB
        sys.exit(1)
    else:
        print "\n\tScenario DB Exists: " + SimDB

    # Check if scenario directory exists
    if not os.path.exists(SimFld):
        print "\n\tCannot locate Scenario Folder: " + SimFld
        sys.exit(1)
    else:
        print "\n\tScenario Folder Exists: " + SimFld

    # Run job
    print "\n\tExecuting STM job..."
    os.chdir(ExeDir)
    Arg = "ApexRms.SyncroSim.Run.exe --lib=\"" + SimDB + "\""
    print "\t" + str(Arg)


    '''
    Stream the stdout and stderr:

    The communicate() method will Interact with process: Send data to stdin.
    Read data from stdout and stderr, until end-of-file is reached. Wait for
    process to terminate.

    Python has no asynchronous subprocess module, so reading the stdout
      causes a block and the output is not received until after the execution completes

    To develop an asynchronous subprocess call, I use threads.
    '''
    ret = subprocess.Popen(Arg, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    reader = ReaderThread(ret.stdout)
    reader.start() # Start process
    ret.wait() # Wait until subprocess is done
    reader.join() # Wait until we've processed all output

    # Wait a bit to ensure completion before analyzing output
    time.sleep(2)

    # Check for messages
    e = ret.communicate()
    ErrBool = False
    e2 = e[0].split("\r\n")
    if len(e2) > 1:
        print "\\tStandard Output:"
        for i in e2:
            if len(i.strip()) > 0:
                print "\t\t" + i
    e3 = e[1].split("\r\n")
    if len(e3) > 1:
        print "\tStandard Error:"
        for i in e3:
            if len(i.strip()) > 0:
                print "\t\t" + i
                ErrBool = True
    if ErrBool:
        if sys.exc_info() != (None, None, None):
            tb2 = sys.exc_info()[2]
            tbinfo = traceback.format_tb(tb)[0]
            pymsg = "\nPYTHON ERRORS:\nTraceback info:\n" + \
                tbinfo + "\nError Info:\n" + str(sys.exc_info()[1])
            print pymsg
        sys.exit(1)
    time.sleep(1)

    # Get competion info
    Get_End_Info()


def Get_Start_Info():
    """
    Get start information before running the job. StdOut will end up in the HTCondor
      log files
    """
    try:
        print "Start Time: " + time.strftime("%Y-%m-%d %H:%M:%S")
        print "User: " + os.getenv("USERNAME")
        print "Computer: " + os.getenv("COMPUTERNAME")
    except:
        print "\tCould not retrieve start information."


def Get_End_Info():
    """
    Get the end information after running the job. StdOut will end up in the HTCondor
      log files
    """
    try:
        print "\nEnd Time: " + time.strftime("%Y-%m-%d %H:%M:%S")
    except:
        print "Could not retrieve start information."


if __name__ == '__main__':
    main()
