#!/usr/bin/env python

# The prism-auto script automatically executes PRISM on one or more
# models/properties, for the purposes of benchmarking or testing.

# The simplest usage is just "prism-auto <target>" where <target>
# is a directory, model file or properties file. For a directory,
# prism-auto finds all models and all properties files in the
# directory and then executes PRISM on each combination of them.
# For a model file, it runs against all properties, and vice versa.

# Run "prism-auto -h" for details of further options.

import os,sys,re,subprocess,signal,tempfile,functools,filecmp,logging
from optparse import OptionParser

#==================================================================================================
# Utility functions
#==================================================================================================

def isPrismModelFile(file):
    return re.match('.+(\.prism|\.nm|\.pm|\.sm)$', file)

def isPrismPropertiesFile(file):
    return re.match('.+(\.props|\.pctl|\.csl)$', file)

def isPrismModelListFile(file):
    return re.match('models$', os.path.basename(file))

def isPrismPropListFile(file):
    return re.match('.+(\.txt)$', file)

def isOutFile(file):
    return re.match('.+(\.out\.)', file) # Note that this matches anything that contains a .out anywhere in the name

def isTestFile(file):
    return re.match('.+(\.test)$', file)

# Check whether the given (args|test) file doesn't have an associated model

def isOrphan(dir, file):
    return not getMatchingModelsInDir(dir, file)

def lineIsCommentedOut(line):
    return line.startswith('#')

# Get a list of models in a directory, either from a "models" file if present,
# or by searching for all files in the directory with an appropriate extension.
# A "models" file is a list of (relative) path names, in which lines starting with # are ignored.
# Each item of the returned list is itself a tuple consisting of the name of the model file and
# a possibly empty argument list, e.g. ("model.pm", ["-const", "N=2"])

def getModelsInDir(dir):
    modelFiles = []
    # Process "models" file, if present
    if os.path.isfile(os.path.join(dir, "models")):
        for line in open(os.path.join(dir, "models"), 'r').readlines():
            line = line.strip()
            if len(line) == 0 or lineIsCommentedOut(line): continue
            first = 1
            args = []
            for item in line.split(' '):
                if first:
                    modelFile = os.path.join(dir, item)
                    first = 0
                else:
                    args.append(item)
            modelFiles.append((modelFile, args))
    # Otherwise look for all model files
    else:
        for file in os.listdir(dir):
            if os.path.isfile(os.path.join(dir, file)) and isPrismModelFile(file):
                modelFiles.append((os.path.join(dir, file), []))
    #print "Model files in " + dir + ": " + ' '.join( map(lambda pair: pair[0], modelFiles ))
    return modelFiles

# Get a list of all files in the directory that satisfy the given predicate

def getFilesInDir(dir, pred):
    resultFiles = []
    for file in os.listdir(dir):
        if os.path.isfile(os.path.join(dir, file)) and pred(file):
            resultFiles.append(os.path.join(dir, file))
    return resultFiles

# Return true iff the basename of file in the directory dir starts with prefix

def startsWith(prefix, dir, file):
    return os.path.basename(os.path.join(dir, file)).startswith(os.path.basename(prefix))

# Get a list of models in a directory matching a (property|args|test) file name.

def getMatchingModelsInDir(dir, fileToMatch):
    return getFilesInDir(dir, lambda file: isPrismModelFile(file) and startsWith(file, dir, fileToMatch))

# Get a list of properties in a directory, by searching for all files with an appropriate extension.

def getPropertiesInDir(dir):
    return getFilesInDir(dir, isPrismPropertiesFile)

# Get a list of properties in a directory with prefix matching a model file name.

def getMatchingPropertiesInDir(dir, modelFile):
    return getFilesInDir(dir, lambda file: isPrismPropertiesFile(file) and startsWith(modelFile, dir, file))

# Get a list of test files in a directory

def getTestFilesInDir(dir):
    return getFilesInDir(dir, isTestFile)

# Get a list of tests in a directory with prefix matching a model file name.

def getMatchingTestsInDir(dir, modelFile):
    return getFilesInDir(dir, lambda file: isTestFile(file) and startsWith(modelFile, dir, file))

# Get a list of all out files in a directory

def getOutFilesInDir(dir):
    return getFilesInDir(dir, isOutFile, false)

# Get a list of all out files in a directory whose prefix matches a model file name

def getMatchingOutFilesInDir(dir, modelFile):
    return getFilesInDir(dir, lambda file: isOutFile(file) and startsWith(modelFile, dir, file))

# Extract all command-line switches from an "args" file into a list
# Just combine switches on all (non-commented) lines together, delimited by spaces
# Returns an empty list if the file does not exist

def getAllArgsFromFile(file):
    args = []
    if not os.path.isfile(file): return args
    for line in open(file, 'r').readlines():
        line = line.strip()
        if len(line) == 0 or lineIsCommentedOut(line): continue
        items = line.split(' ')
        for item in items:
            if len(item) > 0: args.append(item)
    return args

# Extract command-line switches from an "args" file into a list of lists
# Switches from each (non-commented) line, delimited by spaces, are in a separate list
# Returns an empty list if the file does not exist

def getArgsListsFromFile(file):
    argsSet = []
    if not os.path.isfile(file): return argsSet
    for line in open(file, 'r').readlines():
        args = []
        line = line.strip()
        if len(line) == 0 or lineIsCommentedOut(line): continue
        items = line.split(' ')
        for item in items:
            if len(item) > 0: args.append(item)
        if len(args) > 0: argsSet.append(args)
    return argsSet

# Read the matching .args file for the given model/properties file and return a list of lists,
# each list corresponding to one line in the .args file, one argument per list item

def getMatchingArgListsForFile(file):
    if os.path.isfile(file + ".args"):
        return getArgsListsFromFile(file + ".args")
    return [[]]

# Add any extra args provided to this script to each of the given argument lists

def addExtraArgs(argLists):
    if options.extraArgs:
        return map(lambda args: args + options.extraArgs.split(' '), argLists)
    else:
        return argLists

# Returns true iff there is a possible name clash for the given filename

def possibleNameClash(fullName):
    withoutExt = fullName.rsplit('.', 1)[0]
    exts = ['lab','tra','stat','srew','trew']

    return any(map (os.path.exists, [fullName] + [withoutExt + '.' + ext for ext in exts] +  [withoutExt + '1.' + ext for ext in exts]))

# Join directory and filename to obtain a full path
# If doAddPrefix is true, a prefix is prepended to the filename as well

def expandName(dir, option):
    splitOption = option.split(':')
    fullName = os.path.join(dir, splitOption[0])
    return fullName + (":" + splitOption[1] if len(splitOption) > 1 else '')

# Prepend the given prefix to the given filename or filename:option string
# e.g. prependToFile('hello.', '/some/world:here) == '/some/hello.world:here'

def prependToFile(prefix, option):
    splitOption = option.split(':')
    fullName = os.path.join(os.path.dirname(splitOption[0]), 'tmp.' + os.path.basename(splitOption[0]))
    return fullName + (":" + splitOption[1] if len(splitOption) > 1 else '')

# Traverses an argument list, expanding all filenames in import and export switches
# and appending a prefix to each export filename to prevent PRISM from overriding the out file

def expandFilenames(args, dir=""):
    def isImportExportArg(arg):
        return (arg.startswith("-export") or arg.startswith("-import"))
    if args:
        return [args[0]] + [expandName(dir, args[i+1]) if isImportExportArg(args[i]) else args[i+1] for i in range(len(args)-1)]

# Rename all export files in the arguments by prepending prefix

def renameExports(prefix, args):
    def isExportArg(arg):
        return arg.startswith("-export")
    if args:
        return [args[0]] + [prependToFile(prefix, args[i+1]) if isExportArg(args[i]) else args[i+1] for i in range(len(args)-1)]

# Find all files that match any -export switch file argument
# This takes into account that a .all extension corresponds to five different files
# and that multiple reward structures will result in filenames extended with a number

def getExpectedOutFilesFromArgs(args):
    options = [args[i+1] for i in range(len(args)-1) if args[i].startswith("-export")]
    files = map(lambda option: option.split(':')[0], options)

    resultFiles = []
    for file in files:
        split = file.rsplit('.', 1)
        base = split[0]
        # Determine relevant extensions
        if split[1] == 'all':
            exts = ['lab','tra','sta','srew','trew']
        else:
            exts = [split[1]]
        # Find all files of the form base<number>.<ext>
        for ext in exts:
            fullName = base + '.' + ext
            foundFile = False
            if os.path.exists(fullName):
                resultFiles.append(fullName)
                foundFile = True
            else:
                i = 1
                fullName = base + str(i) + '.' + ext
                while os.path.exists(fullName):
                    resultFiles.append(fullName)
                    i += 1
                    fullName = base + str(i) + '.' + ext
                    foundFile = True
            if not foundFile:
                print '\033[93m' + "Warning: There is no file of the form " + base + "[number]." + ext + " to compare against -- will skip" + '\033[0m'
    return resultFiles

# Create a valid name for a log file based on a list of benchmark arguments

def createLogFileName(args, dir=""):
    logFile = '.'.join(args)
    if len(dir) > 0:
        logFile = re.sub(didirr+'/', '', logFile)
    logFile = re.sub('/', '_', logFile)
    logFile = re.sub('[^a-zA-Z0-9=_, \.]', '', logFile)
    logFile = re.sub('[ ]+', '.', logFile)
    logFile = re.sub('[\.]+', '.', logFile)
    logFile = re.sub('^[\._]+', '', logFile)
    return logFile + ".log"

# Walk a directory and execute a callback on each file

def walk(dir, meth):
    dir = os.path.abspath(dir)
    for file in [file for file in os.listdir(dir) if not file in [".","..",".svn"]]:
        nfile = os.path.join(dir, file)
        meth(nfile)
        if os.path.isdir(nfile):
            walk(nfile,meth)

#==================================================================================================
# Benchmarking
#==================================================================================================

# Run PRISM with a given list of command-line args

def runPrism(args, dir=""):
    if options.test:
        if options.testAll: args.append("-testall")
        else: args.append("-test")
    prismArgs = [options.prismExec] + args
    print ' '.join(prismArgs)
    if options.echo:
        return
    if options.logDir:
        logFile = os.path.join(options.logDir, createLogFileName(args, dir))
        f = open(logFile, 'w')
        exitCode = subprocess.Popen(prismArgs, stdout=f).wait()
        #exitCode = subprocess.Popen(prismArgs, cwd=dir, stdout=f).wait()
    elif options.test:
        f = tempfile.NamedTemporaryFile(delete=False)
        logFile = f.name
        exitCode = subprocess.Popen(prismArgs, stdout=f).wait()
    else:
        exitCode = subprocess.Popen(prismArgs).wait()
    # Extract test results if needed
    if options.test:
        for line in open(logFile, 'r').readlines():
            if re.match('Testing result:', line):
                print line,
    if options.test and exitCode != 0:
        for line in open(logFile, 'r').readlines():
            if re.match('Error:', line):
                print line,
        print "To see log file, run:"
        print "edit " + logFile
        if not options.testAll:
            sys.exit(1)

# Checks for each file from the outFiles list whether there is an identical file
# with the name exportPrefix + file. If so, said file is deleted. Otherwise, it is kept
# Returns true iff identical files were found for each out file

def verifyAndCleanupExports(outFiles, exportPrefix):
    result = True
    # Check for equality with out files
    for outFile in outFiles:
    	print "Testing export " + os.path.basename(outFile) + ":",
        expFile = prependToFile(exportPrefix, outFile)
        if os.path.isfile(expFile):
            if filecmp.cmp(outFile, expFile):
                # If successful, notify and delete exported file
                print "PASS"
                os.remove(expFile)
            else:
                print "FAIL (" + os.path.basename(expFile) + " does not match)"
                result = False
        else:
            print "FAIL (no " + os.path.basename(expFile) + " to compare to)"
            result = False
    return result

# Run a benchmark, specified by a list of command-line args,
# possibly iterating over further lists of args from a "bm" file

def benchmark(file, args, dir=""):
    logging.debug("Benchmarking: " + file + ", " + str(args))
    # Add extra arguments from command line, if applicable
    args = addExtraArgs(args)
    # Expand output files to full paths
    args = expandFilenames(args, dir)

    # Rename export files to avoid overriding out files
    exportPrefix = 'tmp.'
    if (not options.noRenaming):
        args = renameExports(exportPrefix, args)

    # print '\033[94m' + "EXECUTING BENCHMARK" + '\033[0m'
    # print "File: " + file
    # print "Directory: " + dir
    # print "Args: " + ' '.join(args)
    # print " "

    modelFileArg = [file] if (file != "") else []

    # Loop through benchmark options, if required
    if options.bmFile and os.path.isfile(os.path.join(options.bmFile)):
        argsLists = getArgsListsFromFile(options.bmFile)
        for bmArgs in argsLists:
            runPrism(modelFileArg + args + bmArgs, dir)
    # If none, just use existing args
    else:
        runPrism(modelFileArg + args, dir)

    # Verify that exported files are correct (if required)
    if not options.echo and outFiles and options.test:
        # Determine which out files apply to this benchmark from the -export switches
        outFiles = getExpectedOutFilesFromArgs(args)
        # print "Out files to verify exports against: " + ' '.join(outFiles)
        allEqual = verifyAndCleanupExports(outFiles, exportPrefix)
        if (not allEqual) and (not options.testAll):
            sys.exit(1)

# Execute benchmarking based on (possibly recursive) processing of a directory

def benchmarkDir(dir):
    logging.debug("Benchmarking dir " + dir)
    # Recurse first, unless asked not to
    if not options.nonRec:
        for file in [file for file in os.listdir(dir) if not file in [".","..",".svn"]]:
            if os.path.isdir(os.path.join(dir, file)):
                benchmarkDir(os.path.join(dir, file))
    # Get model files in dir
    modelFiles = getModelsInDir(dir)
    if len(modelFiles) > 0:
        for modelFile in modelFiles:
            benchmarkModelFile(modelFile[0], modelFile[1], dir)
    # Get "orphan" tests
    testFiles = filter(functools.partial(isOrphan, dir), getTestFilesInDir(dir))
    for testFile in testFiles:
        # print "Processing orphan test " + testFile + " in dir " + dir
        for args in getArgsListsFromFile(testFile):
            benchmark("", args, dir)

# Execute benchmarking based on a single file (model, property, list, test)

def benchmarkFile(file):
    if isPrismModelFile(file):
        benchmarkModelFile(file)
    elif isPrismPropertiesFile(file):
        benchmarkPropertiesFile(file)
    elif isPrismPropListFile(file):
        benchmarkPropListFile(file)
    elif isTestFile(file):
        benchmarkTestFile(file)

# Execute benchmarking based on a single model file

def benchmarkModelFile(modelFile, modelArgs=[], dir=""):
    logging.debug("Benchmarking model file " + modelFile + " " + str(modelArgs))
    if dir == "":
        dir = os.path.dirname(modelFile)
        if dir == "": dir = "."
    # Expand model file based on any .args file
    argLists = getMatchingArgListsForFile(modelFile)
    logging.debug("Arg lists: " + str(argLists))
    for args in argLists:
        # Build mode: just build
        if options.build:
            benchmark(modelFile, modelArgs + args)
        # Otherwise, find properties
        else:
            # Find and benchmark properties
            if options.matching:
                propertiesFiles = getMatchingPropertiesInDir(dir, modelFile)
            else:
                propertiesFiles = getPropertiesInDir(dir)
            logging.debug("Properties files: " + str(propertiesFiles))
            for propertiesFile in propertiesFiles:
                for argsp in getMatchingArgListsForFile(propertiesFile):
                    benchmark(modelFile, modelArgs + args + [propertiesFile] + argsp, dir)
        # Find and benchmark test files
        testFiles = getMatchingTestsInDir(dir, modelFile)
        logging.debug("Test files: " + str(testFiles))
        for testFile in testFiles:
            logging.debug("Test file: " + str(testFile))
            for testArgs in getArgsListsFromFile(testFile):
                benchmark(modelFile, modelArgs + testArgs, dir)

# Execute benchmarking on a test file, i.e. a file that has the same structure as an args file,
# but which does not have any corresponding prop file

def benchmarkTestFile(testFile):
    logging.debug("Benchmarking test file " + testFile)
    dir = os.path.dirname(testFile)
    if dir == "": dir = "."
    if options.matching:
        modelFiles = getMatchingModelsInDir(dir, testFile)
    else:
        modelFiles = getModelsInDir(dir)

    for modelFile in modelFiles:
        # Read args for the model
        for modelArgs in getMatchingArgListsForFile(modelFile):

            # Treat test file like an args file
            for argsList in getArgsListsFromFile(testFile):
                # Don't look for properties (corresponds to build mode)
                benchmark(modelFile, modelArgs + argsList, dir)
    if not modelFiles:
        # There aren't any (matching) model files, process as "orphaned" test
        for argsList in getArgsListsFromFile(testFile):
            benchmark("", argsList, dir)

# Execute benchmarking based on a single properties file

def benchmarkPropertiesFile(propertiesFile):
    logging.debug("Benchmarking properties file " + propertiesFile)
    dir = os.path.dirname(propertiesFile)
    if dir == "": dir = "."
    # Expand properties file based on any .args file
    argLists = getMatchingArgListsForFile(propertiesFile)
    for args in argLists:
        # Find models
        if options.matching:
            modelFiles = getMatchingModelsInDir(dir, propertiesFile)
        else:
            modelFiles = getModelsInDir(dir)
        for modelFile in modelFiles:
            # Expand model based on any .args file, too
            for modelArgs in getMatchingArgListsForFile(modelFile[0]):
                benchmark(modelFile[0], modelFile[1] + modelArgs + [propertiesFile] + args, dir)

# Execute benchmarking based on a property list

def benchmarkPropListFile(propListFile):
    logging.debug("Benchmarking propertiy list file " + propListFile)
    listDir = os.path.dirname(propListFile)
    if listDir == "": listDir = "."
    for line in open(propListFile, 'r').readlines():
        line = line.strip()
        if len(line) == 0 or lineIsCommentedOut(line): continue
        items = line.split(',')
        dir = os.path.join(listDir, items[0].strip())
        dir = os.path.realpath(dir)
        propFile = items[1].strip()
        benchmarkPropertiesFile(os.path.join(dir, propFile))

#==================================================================================================
# Main program
#==================================================================================================

def printUsage():
    print "Usage: prism-auto ..."

def signal_handler(signal, frame):
    sys.exit(1)

# Main program

signal.signal(signal.SIGINT, signal_handler)
parser = OptionParser(usage="usage: %prog [options] arg")
parser.add_option("-l", "--log", dest="logDir", metavar="DIR", default="", help="Store PRISM output in logs in DIR")
parser.add_option("-a", "--args", dest="bmFile", metavar="FILE", default="", help="Read argument lists for benchmarking from FILE")
parser.add_option("-e", "--echo", action="store_true", dest="echo", default=False, help="Just print out tasks, don't execute")
parser.add_option("-m", "--matching", action="store_true", dest="matching", default=False, help="Only use matching models/properties, not all files")
parser.add_option("-b", "--build", action="store_true", dest="build", default=False, help="Just build models, don't model check properties")
parser.add_option("-p", "--prog", dest="prismExec", metavar="FILE", default="prism", help="Program to execute [default=prism]")
parser.add_option("-n", "--non-recursive", action="store_true", dest="nonRec", default=False, help="Don't recurse into directories")
parser.add_option("-x", "--extra", dest="extraArgs", metavar="XXX", default="", help="Pass (single string of) extra switches to PRISM")
parser.add_option("-t", "--test", action="store_true", dest="test", default=False, help="Run in test mode")
parser.add_option("--test-all", action="store_true", dest="testAll", default=False, help="In test mode, don't stop after an error")
parser.add_option("--no-renaming", action="store_true", dest="noRenaming", default=False, help="Don't rename files to be exported")
parser.add_option("--debug", action="store_true", dest="debug", default=False, help="Enable debug mode: display debugging info")
(options, args) = parser.parse_args()
if len(args) != 1:
    parser.print_help()
    sys.exit(1)
if options.debug:
    logging.basicConfig(level=logging.DEBUG)
if options.logDir and not os.path.isdir(options.logDir):
    print "Log directory \"" + options.logDir + "\" does not exist"
    sys.exit(1)
if os.path.isdir(args[0]):
    benchmarkDir(args[0])
elif os.path.isfile(args[0]):
    benchmarkFile(args[0])