#
# This Python module defines a plugin that can be used by pyutilib_test_driver 
# (and related unittest drivers) to execute tests on AMPL solvers.
#

import os
import os.path
import pyutilib.testdriver
import coopr.opt
from pyutilib.component.core import *
import pyutilib.services
import re
import StringIO

#
# Register executables
#
pyutilib.services.register_executable('ampl')


#
# Plugin that defines actions in the setUp, execution and tearDown
# of tests.
#
class AmplTestDriver(Plugin):

    implements(pyutilib.testdriver.ITestDriver)
    alias('ampl.colin')

    def setUpClass(self, cls, options):
        """Set-up the class that defines the suite of tests"""

    def tearDownClass(self, cls, options):
        """Tear-down the class that defines the suite of tests"""

    def setUp(self, testcase, options):
        """Set-up a single test in the suite"""
        #
        # Set the current working directory
        #
        global tmpdir
        tmpdir = os.getcwd()
        os.chdir(options.currdir)
        #
        # Setup temporary file manager
        #
        pyutilib.services.TempfileManager.sequential_files(0)
        pyutilib.services.TempfileManager.tempdir = options.currdir
        #
        # Get the solver specified in the experiment
        #
        testcase.opt =  pyutilib.services.registered_executable(options.solver)
        if testcase.opt is None:
            testcase.skipTest('Solver %s is not available' % options.solver)
        testcase.ampl =  pyutilib.services.registered_executable('ampl')
        if testcase.ampl is None:
            testcase.skipTest('AMPL executable is not available')

    def tearDown(self, testcase, options):
        """Tear-down a single test in the suite"""
        #
        # Delete temporary files
        #
        pyutilib.services.TempfileManager.clear_tempfiles()
        pyutilib.services.TempfileManager.unique_files()
        #
        # Move back to the original working directory
        #
        global tmpdir
        os.chdir(tmpdir)

    def run_test(self, testcase, name, options):
        """Execute a single test in the suite"""
        #
        # Debugging I/O
        #
        if options.verbose or options.debug:
            print "Test %s - Running AMPL with solver %s and options %s" % (name, options.solver, str(options))
        #
        # Create the command-line
        #
        modfile = pyutilib.services.TempfileManager.create_tempfile(suffix='.mod')
        cmd = testcase.ampl.get_path() +" "+modfile
        if options.verbose or options.debug:
            print "Command Line: ",cmd
        #
        # Create MOD file
        #
        OUTPUT = open(modfile, 'w')
        print >>OUTPUT, 'model %s;' % options.model_file
        if options.data:
            print >>OUTPUT, 'data %s;' % options.data_file
        print >>OUTPUT, 'option solver "%s";' % testcase.opt.get_path()
        if options.solver_options:
            print >>OUTPUT, 'option colintest_options (" %s ");' % options.solver_options
        print >>OUTPUT, 'solve;'
        OUTPUT.close()
        #
        # Execute the PEBBL command
        #
        try:
            pyutilib.subprocess.run(cmd, options.currdir+name+".log", cwd=options.currdir)
        except OSError:
            pass
        #
        # Parse the logfile
        #
        results = self.process_log(options.currdir+name+".log")
        #
        # Apply assertions
        #
        testcase.assertEqual(results.value, options.optimum)
        testcase.assertLessEqual(results.valgrind_errors, 0)
        testcase.assertLessEqual(results.valgrind_leaks, 0)
        #
        # Misc cleanup
        #
        if os.path.exists(options.currdir+name+".log"):
            os.remove(options.currdir+name+".log")
        if os.path.exists(options.currdir+name+".sol.txt"):
            os.remove(options.currdir+name+".sol.txt")

    def process_log(self, logfile):
        """Process a logfile generated by a PEBBL solver"""
        if not os.path.exists(logfile):
            return None
        INPUT = open(logfile,"r")
        rstr = ""
        flag=False
        for line in INPUT:
            if 'Solver Results - BEGIN' in line:
                flag=True
            if 'Solver Results - END' in line:
                flag=False
            if flag:
                rstr += line
        #
        ostream = StringIO.StringIO(rstr)
        results = SolverResults()
        results.read(ostream)
        return ostream

