"""
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);
               Requires win32 Python module if this script is not compiled
OS           : Tested and designed for Windows
Dependencies : STM_runScenario.exe (Python script compiled from STM_runScenario.py)
IMPORTANT    : Change argments in "Hardcoded Arguments", found below.


Environment:
    Since, ST-STM software can only run on a Windows operating system,
      this script assumes this by requiring that jobs can only run on
      Windows OS. However, if your pool includes Linux machines and
      Mono software is installed on such machines, HTCondor could be used
      in such a scenario but this script would require modification.
    If files (STM software or STM databases) reside on a local drive
      the user must transfer these files to each HTCondor node.

General Workflow:
    Run ST-SIM partial tasks (Monte Carlo simulations of ST-STM Scenarios) in
      a HTCondor environment

    (1.) On the submit machine, split (decompose) the database using the split
           utility (local)
    (2.) ST-STM replication data setup:
         Copy each "Job-{n}" folder to a different node in the HTCondor pool
         --OR--
         Copy a zip file created from "Job-{n}" folder to a different node
            in the HTCondor pool
         --OR--
         Use each "Job-{n}" folder on a networked storage while running jobs
    (3.) Each replication runs on a HTCondor node
    (4.) Merge the data using the "merge" utility and the SSimJobs.xml on the
           submit machine.

HTCondor Requirements:
    (1.) HTCondor middleware must exist or the user must have access to the
         HTCondor pool for submitting jobs.
    (2.) Depending on access permissions, user may need to strore their password or
         modify the submit template that accompanies this script.

HTCondor File Requirements:
    (1.) SubmitTemplate.sub (Requires editing by the end user to match the HTCondor
         system they are using, because requirements will vary depending on how the
         HTCondor system is managed/configured
    (2.) STM_runScenario.exe (Executable created from its source Python script.
         This executable will run on each HTCondor clinet and does not require that
         Python is installed on the client.

ST-STM File Requirements:
    (1.) Apex ST-STM command line executables
    (2.) ST-STM database that the user set up with the ST-STM plugin for SyncroSim.
         The user is required to create their database and configure the model runs
         in this environment prior to splitting (decomposing) the Monte Carlo
         replications and running in parallel (embarassingly parallel) within HTCondor.

Transfering data to HTCondor client:
    If the user wants to transfer data to the HTCondor node, they need to make sure
      there is enough disk space on the node. The easiest method for doing this is to
      figure out how much disk space is required for the input and output of a single
      Monte Carlo and then define this in the Condor submit template (e.g., request_disk=).
    Although HTCondor supports transfering folders, the output transfered back to the
      submit machine cannot exist. Because of how ST-STM software creates output
      directories and files, we zip the data, delete the original folders, process
      the data on the client, then transfer the directories (.input, .output) and
      file (.ssim)

Executing Merge within HTCondor:
    We are not running the STM merge in HTCondor because transfering all the data
      adds significant overhead and the user could not have the option to submit
      the job to a pool if their data resides on their local machine.

Arguments:
    This script can be converted to an executable so users do not require python
      on their machine


################################################################################
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 libraries
import sys, os, subprocess, traceback, time, shutil, stat, zipfile, getopt, _winreg

# Requires win32 module, third party
try:
    import win32com.client
except ImportError:
    print "Error importing win32 module. Not installed or corrupted."
    sys.exit(1)


# Hardcoded Arguments ----------------------------------------------------------
# Set the workspace of your state-transition ssim database
STM_DB = r"<PATH HERE>\ST-Sim-Spatial-Sample-V1-4-2.ssim"

# Stratum ID
sid = "210"

# How many Monte Carlos will you run
NumJobs = "20"

# Set the workspace to the folder containing the command line STM-Sim software
STM_Lib = r"<PATH HERE>\STSim-Zips\STSim-Windows"

# Do you want to zip and transfer your data to the remote client
TransferZip = "False"


''' # Uncomment this section if a commandline version is desired (not tested and likely not updated to latest version)
# Passed Arguments -------------------------------------------------------------
# ST-Sim database that the user will supply to execute the Monte Carlo simulations
#   in parallel
STM_DB = sys.argv[1]

# ST-Sim database scenario ID supplied to execute the Monte Carlo simulations
#   in parallel
sid = sys.argv[2]

# The number of Monte Carlo simulations to execute for the provided scenario (must
#   match the settings in the ST-STM database).
NumJobs = sys.argv[3]

# Path of STM libraries and execuatables used to execute code via command line
STM_Lib = sys.argv[4]

# If files are transfered does user want to zip these files
# Evalauted only if transfer True
TransferZip = eval(sys.argv[5])
'''

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
def CommandLine():
    """
    This will allow someone to run the python script from command line and check
      what the commandline arguments are. This has not been tested and it currently
      is not in use.
    """
    # Set global variables required for script
    global STM_DB, SID, NumJobs, STM_Lib, TransferZip

    try:
        opts, args = getopt.getopt(sys.argv[1:],"hdsjl:",["STM_DB=","sid=", "NumJobs=", "STM_Lib=", "Zip="])
    except getopt.GetoptError:
       msg = "ST-Sim_HTCondor_setup.exe \n\t--d <ST-SIM database file> \n\t--s <Scenario ID>" + \
        "\n\t--j <Count of Monte Carlo> \n\t--l <Path to STM commandline executables>" + \
        "\n\t--z <True|False for zipping transfered data>"
       AddMsgAndPrint(msg, 2)
       sys.exit(1)
    for opt, arg in opts:
       if opt == '-h':
           msg = "ST-Sim_HTCondor_setup.exe \n\t--d <ST-SIM database file> \n\t--s <Scenario ID>" + \
           "\n\t--j <Count of Monte Carlo> \n\t--l <Path to STM commandline executables>" + \
           "\n\t--z <True|False for zipping transfered data>"
           AddMsgAndPrint(msg, 2)
           sys.exit(1)
       elif opt in ("-d", "--STM_DB"):
          STM_DB = arg
       elif opt in ("-s", "--sid"):
          sid = arg
       elif opt in ("-j", "--NumJobs"):
          NumJobs = arg
       elif opt in ("-l", "--STM_Lib"):
          STM_Lib = arg
       elif opt in ("-z", "--Zip"):
          TransferZip = arg


def main():
    """
    This is the main routine that call all subsequent routines for setting up a HTCondor
      job to run ST-SIM models.
    """

    msg = "+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++"
    AddMsgAndPrint(msg, 0)
    msg = "Starting the setup for running STM with HTCondor..."
    AddMsgAndPrint(msg, 0)

    msg = "\nCollecting information on required files before continuing..."
    AddMsgAndPrint(msg, 0)

    # Assign HTCondor Installation path
    global HTCondorBin
    HTCondorBin = GetCondorBin()

    '''
    # These may not all work, which depends on how the script is invoked
    print sys.argv[0]
    print os.getcwd()
    print sys.path[0]
    '''

    # Assign HTCondor submit template file
    global Condor_SubmitTemplate
    Condor_SubmitTemplate = os.path.join(os.path.dirname(sys.argv[0]), "SubmitTemplate.sub")

    # QAQC
    if TransferZip not in ["True", "False"]:
        msg = "\tThe transfer and zip argumnets are case sensitive and require 'True' or 'False'"
        AddMsgAndPrint(msg, 2)
        sys.exit(1)

    # Global Parameters
    global RootWS, ClientExe, TransferZip2
    RootWS = os.getcwd() # Location of script
    ClientExe = os.path.join(RootWS, "STM_runScenario.exe")
    TransferZip2 = eval(TransferZip)

    # QAQC
    if not os.path.exists(STM_DB):
        msg = "\tST-SIM database file specified by the user is missing: " + STM_DB
        AddMsgAndPrint(msg, 2)
        sys.exit(1)
    if not os.path.exists(STM_Lib):
        msg = "\tST-SIM libraries specified by the user are missing: " + STM_Lib
        AddMsgAndPrint(msg, 2)
        sys.exit(1)
    if not os.path.exists(Condor_SubmitTemplate):
        msg = "\tA template HTCondor submit files specified by the user is missing: " + \
          Condor_SubmitTemplate
        AddMsgAndPrint(msg, 2)
        sys.exit(1)
    if not os.path.exists(os.path.join(HTCondorBin, "condor_master.exe")):
        msg = "\tThe HTCondor bin workspace specified by user is missing: " + HTCondorBin
        AddMsgAndPrint(msg, 2)
        sys.exit(1)


    msg = "\nCheck for output on previous runs and delete existing folders..."
    AddMsgAndPrint(msg, 0)
    SSimJobs_WS = os.path.dirname(STM_DB)
    DelFolder = os.path.join(SSimJobs_WS, STM_JobsOutWS)
    if os.path.exists(DelFolder):
        msg = "\tDeleting: " + DelFolder
        DeleteFolder(DelFolder, msg)

    msg = "\nSplit the jobs and set up the files to run..."
    AddMsgAndPrint(msg, 0)
    SplitScenario()

    msg = "\nGenerate HTCondor DAG..."
    AddMsgAndPrint(msg, 0)
    Batch1 = GenerateDAG()

    msg = "\nGenerate workspace for HTCondor log files..."
    AddMsgAndPrint(msg, 0)
    CreateHTCondor_WS()


    msg = "\n+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++\n"
    AddMsgAndPrint(msg, 0)
    msg = "Step 1: Test that the HTCondor user has access: \n\tcondor_store_cred query\n"
    AddMsgAndPrint(msg, 0)
    msg = "Step 2: The user should execute the batch file from a cmd.exe: \n\t" + Batch1
    AddMsgAndPrint(msg, 0)
    msg = "  The user can then monitor their HTCondor jobs with condor_q -dag\n"
    AddMsgAndPrint(msg, 0)
    msg = "Step 3: After the distributed computing completes, the user should " + \
          "evaluate the results in ST-SIM software.\n"
    AddMsgAndPrint(msg, 0)



def GetCondorBin():
    """
    Determine location of HTCondor installation
    """
    # Set variables to null
    HTCondorBin = 'NotFound'
    try:
        # HTCondor is a 32bit app so this will allow a 64bit version of python to find the key
        hkey = _winreg.OpenKey(_winreg.HKEY_LOCAL_MACHINE, 'SOFTWARE\\Condor',
               0, _winreg.KEY_READ | _winreg.KEY_WOW64_32KEY)
        RegVal1 = _winreg.QueryValueEx(hkey, 'RELEASE_DIR')
        _winreg.CloseKey(hkey)
        HTCondorBin = str(os.path.join(RegVal1[0], "bin"))
        if os.path.exists(os.path.join(HTCondorBin, "condor_master.exe")):
            pass
        else:
            msg = "\tDid not find HTCondor installation path..." + HTCondorBin
            AddMsgAndPrint(msg, 0)
            msg = HTCondorBin
            AddMsgAndPrint(msg, 2)
            sys.exit(1)
    except:
        try: _winreg.CloseKey(hkey)
        except: pass
        msg = "\tDid not find HTCondor installation path (error accessing registry)..."
        AddMsgAndPrint(msg, 2)
        sys.exit(1)

    return HTCondorBin


def getDriveMappings():
    """
    Return a list of mapped drives and the UNC path for the respective drive
      This does not return local drives, but mapped network storage only.
    This is required because HTCondor may not suuport loading of users' mapped
      drives, which depends on the HTCondor configuration. Using tUNC paths will
      handle all cases, which is what the script relies on.
    """
    # Get a list of networked drives available to the user
    DriveUNClist = []
    network = win32com.client.Dispatch('WScript.Network')
    drives = network.EnumNetworkDrives()

    # EnumNetworkDrives returns an even-length array of drive/unc pairs.
    for iDrive in range(0, len(drives), 2):
        DriveUNClist.append((str(drives[iDrive]), str(drives[iDrive+1])))

    return DriveUNClist


def DeleteFolder(DirTree, msg):
    """
    Check if directory exists and if so delete everything within folder
      including itself.
    """
    if os.path.exists(DirTree):
        for root, dirs, files in os.walk(DirTree, topdown=False):
            time.sleep(1)
            for name in files:
                filename = os.path.join(root, name)
                os.chmod(filename, stat.S_IWRITE)
                os.remove(filename)
            for dir in dirs:
                os.chmod(os.path.join(root, dir), stat.S_IWRITE)
                os.rmdir(os.path.join(root, dir))
    if os.path.exists(DirTree):
        try:
            time.sleep(2)
            os.chmod(DirTree, stat.S_IWRITE)
            os.rmdir(DirTree)
        except: pass


def SplitScenario():
    """
    Executes the ST-SIM software to decompose a ST-SIM database into partial models
      so they can be executed independenlty as embarassingly parallel tasks in
      HTCondor.
    This could be set up to support a user defining an output directory, but for
      now this script uses the default. To provide support for an ouput folder,
      add another argument for running the script and update the name space
      (end of script)
    """
    Arg = os.path.join("\"" + STM_Lib, "ApexRms.SyncroSim.Split.exe\"") + \
        " --lib=\"" + STM_DB + "\" --sid=" + sid + " --jobs=" + NumJobs
    ret = subprocess.Popen(Arg, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    ret.wait()
    time.sleep(2)

    # Check for messages
    e = ret.communicate()
    ErrBool = False
    e2 = e[0].split("\r\n")
    #print str(e2)
    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")
    #print str(e3)
    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)


def GenerateDAG():
    """
    Generates all the files used for submitting a directed acyclic graph job to
      HTCondor. This routine moves a lot of files around and creates new files to
      execute everything correctly in a HTCondor environment.
    This script uses a DAG because DAGs allow easier identification of failed jobs
      and re-submitting of failed jobs versus running everything again.
    """
    # Open template HTCondor submit file
    HTCondor_templateF = open(Condor_SubmitTemplate, 'r')
    HTCondor_templateFRead = HTCondor_templateF.read()
    HTCondor_templateF_List = HTCondor_templateFRead.split("\n")
    HTCondor_templateF.close()
    del HTCondor_templateF

    # Get a list of drives and UNC mappings
    DriveUNClist = getDriveMappings()

    # Generate a list of directories created from the split (decomposition)
    for root, dirs, files in os.walk(STM_JobsOutWS):
        break
    if len(dirs) == 0:
        msg = "The setup of output directories for ST-STM changed and this script requires an update."
        AddMsgAndPrint(msg, 2)
        sys.exit(1)

    # Create a sorted list of directories; If ST-STM software changes, this may require an update
    dirs2 = []
    for idir in range(1, len(dirs)+1):
        dirs2.append("Job-" + str(idir) + ".ssim.input")

    # Define the XML file used for the final merge
    OutJobsXML = os.path.join(STM_JobsOutWS, SSimJobs_xml)

    # Copy STM commandline libraries to Condor workspace and zip if this is required by user
    if TransferZip2 == True:
        # Zip file of STM software
        if os.path.exists(STM_Lib + ".zip"):
            try:
                shutil.copyfile(STM_Lib, os.path.join(STM_JobsOutWS, os.path.basename(STM_Lib)))
            except:
                msg = "Error copying STM software: " + \
                    os.path.join(STM_JobsOutWS, os.path.basename(STM_Lib))
                AddMsgAndPrint(msg, 2)
                sys.exit(1)
        # Zip file for STM libraries does not exist so create in SSimJobs workspace
        else:
            msg = "\tGenerating: " + os.path.basename(STM_Lib) + ".zip"
            AddMsgAndPrint(msg, 0)
            Data = False
            zip_folder(STM_Lib, "#", os.path.join(STM_JobsOutWS, os.path.basename(STM_Lib)) + ".zip", Data)
    else:
        # Not transfering and zipping data
        pass


    # Convert paths of files to UNC
    ResultsWS_Dr = os.path.splitdrive(STM_JobsOutWS)[0]
    ClientExe_Dr = os.path.splitdrive(ClientExe)[0]
    STM_Lib_Dr = os.path.splitdrive(STM_Lib)[0]
    #
    # Defaults so if data stored on local drive, this will work
    ResultsWS2 = STM_JobsOutWS
    ClientExe2 = ClientExe
    STM_Lib2 = STM_Lib
    #
    for iDriveUNClist in DriveUNClist:
        if ResultsWS_Dr in iDriveUNClist[0]:
            ResultsWS2 = STM_JobsOutWS.replace(ResultsWS_Dr, iDriveUNClist[1])
        if ClientExe_Dr in iDriveUNClist[0]:
            ClientExe2 = ClientExe.replace(ClientExe_Dr, iDriveUNClist[1])
        if STM_Lib_Dr in iDriveUNClist[0]:
            STM_Lib2 = STM_Lib.replace(STM_Lib_Dr, iDriveUNClist[1])

    # Open output DAG file so we can write to it for HTCondor jobs
    DateTime = time.strftime("%Y%m%d_%Hh%Mm%Ss")
    OutDAG = os.path.join(STM_JobsOutWS, "STM_" + DateTime + ".dag")
    f2 = open(OutDAG, 'w')

    # Enumerate each Monte Carlo folder created with the ST-STM split command
    #   line software and set up HTCondor jobs.
    UniqCntr = 0
    Parents = []
    for iFolder in dirs2:
        iFolder = os.path.join(ResultsWS2, iFolder)

        # SSIM database file that is tied to the folder for the individual Monte Carlo
        SSIM_File = iFolder.replace(".input", "")

        # SSIM folder that contains the data to run the Monte Carlo on the client
        SSIM_Fld = iFolder

        # Create a zip file of the scenario data if data to be transferred via zip
        if TransferZip2 == True:
            msg = "\tGenerating: " + os.path.basename(iFolder) + ".zip"
            AddMsgAndPrint(msg, 0)
            Data = True
            zip_folder(iFolder, SSIM_File, os.path.join(STM_JobsOutWS, os.path.basename(iFolder)) + ".zip", Data)
            msg = msg = "\tDeleting: " + iFolder
            # The returned results will have the same folder name and location, so delete
            #   this interim data
            DeleteFolder(iFolder, msg)

        # Create the submit files to include the zip files and libraries
        # Include the unzip.exe, STSim.zip, batch file, scenario zip folder, and jobs_n.xml file
        OutSubmit = os.path.join(STM_JobsOutWS, os.path.basename(iFolder) + ".sub")
        f = open(OutSubmit, 'w')
        bool_exe, bool_err, bool_log, bool_out, \
            bool_args, bool_transIn, bool_transOut = True, True, True, True, True, True, True
        for iLine in HTCondor_templateF_List:
            if iLine.strip().lower() == "executable =":
                f.write(iLine + " " + ClientExe2 + "\n")
                bool_exe = True
            # ---------
            elif iLine.strip().lower() == "error =":
                f.write(iLine + " " + \
                    os.path.join(ResultsWS2, "HTcondorLogs",
                        os.path.basename(SSIM_File.replace(".ssim", "")) + ".err") + "\n")
                bool_err = True
            # ---------
            elif iLine.strip().lower() == "log =":
                f.write(iLine + " " + \
                    os.path.join(ResultsWS2, "HTcondorLogs",
                        os.path.basename(SSIM_File.replace(".ssim", "")) + ".log") + "\n")
                bool_log = True
            # ---------
            elif iLine.strip().lower() == "output =":
                f.write(iLine + " " + \
                    os.path.join(ResultsWS2, "HTcondorLogs",
                        os.path.basename(SSIM_File.replace(".ssim", "")) + ".out") + "\n")
                bool_out = True
            # ---------
            elif iLine.strip().lower() == "arguments =":
                if TransferZip2 == True:
                    f.write(iLine + " " + \
                    os.path.join(ResultsWS2, os.path.basename(STM_Lib)) + ".zip" + \
                        " " + os.path.basename(iFolder) + ".zip True\n")
                else:
                    f.write(iLine + " " + \
                    STM_Lib2 + " " + SSIM_Fld + " False\n")
                bool_args = True
            # ---------
            elif iLine.strip().lower() == "transfer_input_files =":
                if TransferZip2 == True:
                    f.write(iLine + " " + \
                        os.path.basename(STM_Lib2) + ".zip," + \
                            os.path.basename(iFolder) + ".zip" + "\n")
                else:
                    f.write(iLine + " " + STM_Lib2 + "," + SSIM_Fld + "," + SSIM_File + "\n")
                bool_transIn = True
            # ---------
            elif iLine.strip().lower() == "transfer_output_files =":
                # Will go to the root directory were submit file executed from.
                # If inputs are not transfered, then we do not need to transfer output
                # If we transfer folders, we may need to rename the output folders
                #   for this to work or create a new workspace.
                # When transfering directories, do not use a forward or backward slash as shown
                #   in the documentation, this does not work.
                if TransferZip2 == True:
                    Fld_1 = os.path.basename(iFolder)
                    Fld_2 = Fld_1.replace("input", "output")
                    f.write(iLine + " " + Fld_1 + "," + Fld_2 + "," + os.path.basename(SSIM_File) + "\n")
                bool_transOut = True
            # ---------
            else:
                f.write(iLine + "\n")
        f.close()
        del f

        # Ensure that these are included regardless of what the user provided in the template
        iCount = 0
        ExitAll = False
        for ibool in [bool_exe, bool_err, bool_log, bool_out, bool_args, bool_transIn, bool_transOut]:
            if not ibool:
                if iCount == 0:
                    msg = "HTCondor submit template file Error: Missing executable ="
                    AddMsgAndPrint(msg, 0)
                if iCount == 1:
                    msg = "HTCondor submit template file Error: error ="
                    AddMsgAndPrint(msg, 0)
                if iCount == 2:
                    msg = "HTCondor submit template file Error: log ="
                    AddMsgAndPrint(msg, 0)
                if iCount == 3:
                    msg = "HTCondor submit template file Error: output ="
                    AddMsgAndPrint(msg, 0)
                if iCount == 4:
                    msg = "HTCondor submit template file Error: arguments ="
                    AddMsgAndPrint(msg, 0)
                if iCount == 5:
                    msg = "HTCondor submit template file Error: transfer_input_files ="
                    AddMsgAndPrint(msg, 0)
                if iCount == 6:
                    msg = "HTCondor submit template file Error: transfer_output_files ="
                    AddMsgAndPrint(msg, 0)
                ExitAll = True
            iCount += 1
        if ExitAll:
            try:
                f2.close()
                del f2
            except:
                pass
            msg = "Exiting..."
            AddMsgAndPrint(msg, 2)
            sys.exit(1)

        # Create the DAG file; this will have a post-processing component too
        f2.write("JOB J" + str(UniqCntr) + " " + os.path.basename(OutSubmit) + "\n")
        Parents.append("J" + str(UniqCntr))
        UniqCntr += 1
    f2.close()
    del f2

    # Generate cmd.exe batch file so users can launch more easily
    HTC_WS = os.path.dirname(STM_DB)
    Batch = GenHTCondorBatch(HTC_WS, OutDAG, OutJobsXML)

    return Batch


def zip_folder(folder_path, SSIM_File, output_path, Data):
    """
    Zip the contents of an entire folder. Empty subfolders will be
    included in the archive as well. The .ssim database file that resides outside
      the folder being zipped will also be included in this zip file

    folder_path: Path and directory to zip
    SSIM_File: The path and file name of the Monte Carlo .ssim (If zipping data, then
      This is defined as #
    output_path: The name of the output zip file
    Data: Boolean stating whether folder contains software (False) or data (True)
    """
    try:
        zip_file = zipfile.ZipFile(output_path, 'w', zipfile.ZIP_DEFLATED)

        # Add root directory
        #zip_file.write(folder_path, os.path.basename(folder_path))
        RootFld = os.path.basename(folder_path)

        # Add all nested directoriees and files
        for root, folders, files in os.walk(folder_path):
            # Include all subfolders, including empty ones.
            for folder_name in folders:
                absolute_path = os.path.join(root, folder_name)
                if Data:
                    # Write output into folder with name of zip
                    relative_path = os.path.join(RootFld, absolute_path.replace(folder_path + '\\', ''))
                else:
                    relative_path = absolute_path.replace(folder_path + '\\', '')

                zip_file.write(absolute_path, relative_path)
            for file_name in files:
                absolute_path = os.path.join(root, file_name)
                if Data:
                    # Write output into folder with name of zip
                    relative_path = os.path.join(RootFld, absolute_path.replace(folder_path + '\\', ''))
                else:
                    relative_path = absolute_path.replace(folder_path + '\\', '')
                zip_file.write(absolute_path, relative_path)

        # Add .ssim database that is stored outside of Job folder if zipping data
        #   (not required if zipping software)
        if SSIM_File != "#":
            zip_file.write(SSIM_File, os.path.basename(SSIM_File))

        # Delete the input data since they are zipped (Do not delete software if this
        #   is being zipped)
        if Data:
            if os.path.exists(folder_path):
                msg = "Deleting: " + os.path.basename(folder_path)
                DeleteFolder(folder_path, msg)
            if os.path.exists(SSIM_File):
                os.chmod(SSIM_File, stat.S_IWRITE)
                os.remove(SSIM_File)

    except IOError, message:
        print message
        sys.exit(1)
    except OSError, message:
        print message
        sys.exit(1)
    except zipfile.BadZipfile, message:
        print message
        sys.exit(1)
    finally:
        zip_file.close()


def GenHTCondorBatch(HTC_WS, OutDAG, OutJobsXML):
    """
    Creates a batch file that the user can execute to launch the ST-SIM jobs
      in a HTCondor pool
    """
    # Get a list of drives and UNC mappings
    DriveUNClist = getDriveMappings()

    # Get drive associated to path
    OutDAG_Dr = os.path.splitdrive(OutDAG)[0]
    OutSTM_Dr = os.path.splitdrive(STM_Lib)[0]

    # Defaults so if data stored on local drive, this will work
    OutDAG2 = OutDAG
    #
    for iDriveUNClist in DriveUNClist:
        if OutDAG_Dr in iDriveUNClist[0]:
            OutDAG2 = OutDAG.replace(OutDAG_Dr, iDriveUNClist[1])
        if OutSTM_Dr in iDriveUNClist[0]:
            STM_Lib2 = OutDAG.replace(OutSTM_Dr, iDriveUNClist[1])

    # Get username
    UserName = str(os.getenv('USERNAME'))

    # Write the batch file -------------------------------------------------------
    DateTime = time.strftime("%Y%m%d_%Hh%Mm%Ss")
    Batch = os.path.join(HTC_WS, "STM_ExecuteHTcondor_" + DateTime + ".bat")
    f2 = open(Batch, 'w')

    f2.write("@echo off\n\n")
    f2.write("::--------------------------------- Submitting HTCondor jobs\n")
    f2.write("echo.\n")
    f2.write("echo Submitting HTCondor jobs...\n")
    f2.write("pushd " + os.path.dirname(OutDAG2) + "\n")
    f2.write(os.path.join(HTCondorBin, "condor_submit_dag.exe") + \
        " -no_submit -notification never " + os.path.basename(OutDAG) + "\n")
    f2.write(os.path.join(HTCondorBin, "condor_submit.exe") + \
        " " + os.path.basename(OutDAG) + ".condor.sub" + "\n")
    f2.write("popd\n")
    f2.write("sleep 15\n")
    f2.write("\n")
    f2.write("\n")

    f2.write("::--------------------------------- Monitor HTCondor job completion\n")
    f2.write("rem This assumes the user on the submit machine has no other jobs running.\n")
    f2.write("rem Determine when HTCondor finishes jobs.\n")
    f2.write("echo.\n")
    f2.write("echo Checking Job Status...\n")
    f2.write("\n")
    f2.write(":while1\n")
    f2.write("SET JobsRunning=0\n")
    f2.write("\n")
    f2.write("sleep 15\n")
    f2.write("\n")
    f2.write("rem Query Condor job status, print to stdout and to ascii file.\n")
    f2.write("echo.\n")
    f2.write(os.path.join(HTCondorBin, "condor_q") + " -dag\n")
    f2.write(os.path.join(HTCondorBin, "condor_q") + " -submitter " + UserName + \
             " -format \"%d.\" ClusterId -format \"%d\\n\" ProcId > " + \
             os.path.join(os.path.dirname(OutDAG), "JobStatus.txt") + "\n")

    f2.write("\n")
    f2.write("rem 0 jobs; 0 completed, 0 removed, 0 idle, 0 running, 0 held, 0 suspended\n")
    f2.write("rem Read output file and determine if jobs completed\n")
    f2.write("FOR /F %%i IN (" + os.path.join(os.path.dirname(OutDAG), "JobStatus.txt") + ") DO (\n")
    f2.write("  SET JobsRunning=999\n")
    f2.write(")\n")
    f2.write("\nrem Delete text file\n")
    f2.write("if %JobsRunning% == 999 (\n")
    f2.write("  del " + os.path.join(os.path.dirname(OutDAG), "JobStatus.txt") + "\n")
    f2.write("  goto :while1\n")
    f2.write(")\n")
    f2.write("\n")
    f2.write("echo.\n")
    f2.write("echo Jobs completed...\n")
    f2.write("\n")
    f2.write("\n")

    f2.write("::--------------------------------- Merge STM results\n")
    f2.write("echo.\n")
    f2.write("echo Merge STM results...\n")
    f2.write("rem Check if HTCodnor rescue file exists and if not then merge results\n")
    f2.write("rem If rescue file exists then the user needs to run batch again\n")
    f2.write("rem but I need to handle multiple resuce files somehow\n")
    f2.write("rem '*.rescue*'\n")
    f2.write("if exist " + OutDAG + ".condor.sub.rescue (\n")
    f2.write("    echo One or many HTCondor jobs failed. Investigate the resuce " + \
             "file and re-run this batch depending on reason for failed jobs.\n")
    f2.write(") else (\n")
    f2.write("    " + os.path.join(STM_Lib, "ApexRms.SyncroSim.Merge.exe") + \
             " --cfg=\"" + OutJobsXML + "\"\n")
    f2.write(")\n")
    f2.write("\n")
    f2.write("\n")

    f2.write(":: ------------------------------------- Allow user to read any STDOUT\n")
    f2.write("echo.\n")
    f2.write("echo Merge completed...\n")
    f2.write("pause\n\n")

    return Batch


def CreateHTCondor_WS():
    """
    Creates a workspace to store HTCondor log files, which are specific to each job.
    Users can review these files if jobs fail or for STDOUT information.
    """
    SSimJobs_WS = os.path.join(os.path.dirname(STM_DB), STM_JobsOutWS)
    Log_Folder = os.path.join(SSimJobs_WS, "HTcondorLogs")

    if not os.path.exists(Log_Folder):
        os.mkdir(Log_Folder)
    else:
        DelFolder = Log_Folder
        msg = "\tDeleting: " + DelFolder
        DeleteFolder(DelFolder, msg)
        try:
            os.mkdir(Log_Folder)
        except:
            print "Error deleing directory: " + \
                  os.path.join(os.path.dirname(STM_DB), "HTcondorLogs")


def AddMsgAndPrint(msg, severity=0):
    """
    Adds a Message to the stdout (standard output) to inform user of steps and
      errors
    """
    try:
        if severity == 0:
            print msg
        elif severity == 1:
            print msg
            tb = 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
        elif severity == 2:
            print msg
            tb = 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
            CleanUp()
            sys.exit(1)
    except:
        CleanUp()
        sys.exit(1)


def CleanUp():
    """
    Garbage collector...
    Delete all remaining variables.
    """
    for v in vars().copy(): del v


if __name__ == '__main__':

    # If using the commandline version of this script this needs to be uncommented
    #CommandLine()

    # STM split output folder (User will not have control for defining this)
    STM_JobsOutWS = os.path.join(os.path.dirname(STM_DB),
        os.path.basename(STM_DB) + ".temp", "Scenario-" + sid, "Parallel")

    # XML file name used for merging the results
    SSimJobs_xml = "SSimJobs.xml"

    # Execute code
    main()
    CleanUp()

