# ____________________________________________________________________________________
#
# Pyomo: Python Optimization Modeling Objects
# Copyright (c) 2008-2026 National Technology and Engineering Solutions of Sandia, LLC
# Under the terms of Contract DE-NA0003525 with National Technology and Engineering
# Solutions of Sandia, LLC, the U.S. Government retains certain rights in this
# software.  This software is distributed under the 3-clause BSD License.
# ____________________________________________________________________________________

import os

from pyomo.common.tempfiles import TempfileManager
import pyomo.common.unittest as unittest

import pyomo.kernel as pmo
from pyomo.core import (
    Binary,
    ConcreteModel,
    Constraint,
    Objective,
    Var,
    Integers,
    RangeSet,
    minimize,
    quicksum,
    Suffix,
)
from pyomo.opt import ProblemFormat, convert_problem, SolverFactory, BranchDirection
from pyomo.solvers.plugins.solvers.CPLEX import (
    CPLEXSHELL,
    MockCPLEX,
    _validate_file_name,
)


class _mock_cplex_128:
    def version(self):
        return (12, 8, 0)


class _mock_cplex_126:
    def version(self):
        return (12, 6, 0)


class CPLEX_utils(unittest.TestCase):
    def test_validate_file_name(self):
        _126 = _mock_cplex_126()
        _128 = _mock_cplex_128()

        # Check plain file
        fname = 'foo.lp'
        self.assertEqual(fname, _validate_file_name(_126, fname, 'xxx'))
        self.assertEqual(fname, _validate_file_name(_128, fname, 'xxx'))

        # Check spaces in the file
        fname = 'foo bar.lp'
        with self.assertRaisesRegex(ValueError, "Space detected in CPLEX xxx file"):
            _validate_file_name(_126, fname, 'xxx')
        self.assertEqual('"%s"' % (fname,), _validate_file_name(_128, fname, 'xxx'))

        # check OK path separators
        fname = 'foo%sbar.lp' % (os.path.sep,)
        self.assertEqual(fname, _validate_file_name(_126, fname, 'xxx'))
        self.assertEqual(fname, _validate_file_name(_128, fname, 'xxx'))

        # check BAD path separators
        bad_char = '/\\'.replace(os.path.sep, '')
        fname = 'foo%sbar.lp' % (bad_char,)
        msg = r'Unallowed character \(%s\) found in CPLEX xxx file' % (
            repr(bad_char)[1:-1],
        )
        with self.assertRaisesRegex(ValueError, msg):
            _validate_file_name(_126, fname, 'xxx')
        with self.assertRaisesRegex(ValueError, msg):
            _validate_file_name(_128, fname, 'xxx')

        # check allowable characters
        fname = 'foo$$bar.lp'
        self.assertEqual(fname, _validate_file_name(_126, fname, 'xxx'))
        self.assertEqual(fname, _validate_file_name(_128, fname, 'xxx'))
        fname = 'foo_bar.lp'
        self.assertEqual(fname, _validate_file_name(_126, fname, 'xxx'))
        self.assertEqual(fname, _validate_file_name(_128, fname, 'xxx'))
        fname = 'foo&bar.lp'
        self.assertEqual(fname, _validate_file_name(_126, fname, 'xxx'))
        self.assertEqual(fname, _validate_file_name(_128, fname, 'xxx'))
        fname = 'foo~bar.lp'
        self.assertEqual(fname, _validate_file_name(_126, fname, 'xxx'))
        self.assertEqual(fname, _validate_file_name(_128, fname, 'xxx'))
        fname = 'foo-bar.lp'
        self.assertEqual(fname, _validate_file_name(_126, fname, 'xxx'))
        self.assertEqual(fname, _validate_file_name(_128, fname, 'xxx'))

        # Check unallowable character
        bad_char = '^'
        fname = 'foo%sbar.lp' % (bad_char,)
        msg = r"Unallowed character \(\^\) found in CPLEX xxx file"
        with self.assertRaisesRegex(ValueError, msg):
            _validate_file_name(_126, fname, 'xxx')
        with self.assertRaisesRegex(ValueError, msg):
            _validate_file_name(_128, fname, 'xxx')


class CPLEXShellWritePrioritiesFile(unittest.TestCase):
    """Unit test on writing of priorities via `CPLEXSHELL._write_priorities_file()`"""

    suffix_cls = Suffix

    def setUp(self):
        TempfileManager.push()
        self.mock_model = self.get_mock_model()
        self.mock_cplex_shell = self.get_mock_cplex_shell(self.mock_model)
        self.mock_cplex_shell._priorities_file_name = TempfileManager.create_tempfile(
            suffix=".cplex.ord"
        )

    def tearDown(self):
        TempfileManager.clear_tempfiles()

    def get_mock_model(self):
        model = ConcreteModel()
        model.x = Var(within=Binary)
        model.con = Constraint(expr=model.x >= 1)
        model.obj = Objective(expr=model.x)
        return model

    def get_mock_cplex_shell(self, mock_model):
        solver = MockCPLEX()
        solver._problem_files, solver._problem_format, solver._smap_id = (
            convert_problem(
                (mock_model,),
                ProblemFormat.cpxlp,
                [ProblemFormat.cpxlp],
                has_capability=lambda x: True,
                symbolic_solver_labels=True,
            )
        )
        return solver

    def get_priorities_file_as_string(self, mock_cplex_shell):
        with open(mock_cplex_shell._priorities_file_name, "r") as ord_file:
            priorities_file = ord_file.read()
        return priorities_file

    @staticmethod
    def _set_suffix_value(suffix, variable, value):
        suffix.set_value(variable, value)

    def test_write_without_priority_suffix(self):
        with self.assertRaises(ValueError):
            CPLEXSHELL._write_priorities_file(self.mock_cplex_shell, self.mock_model)

    def test_write_priority_to_priorities_file(self):
        self.mock_model.priority = self.suffix_cls(
            direction=Suffix.EXPORT, datatype=Suffix.INT
        )
        priority_val = 10
        self._set_suffix_value(
            self.mock_model.priority, self.mock_model.x, priority_val
        )

        CPLEXSHELL._write_priorities_file(self.mock_cplex_shell, self.mock_model)
        priorities_file = self.get_priorities_file_as_string(self.mock_cplex_shell)

        self.assertEqual(
            priorities_file,
            "* ENCODING=ISO-8859-1\n"
            "NAME             Priority Order\n"
            "  x 10\n"
            "ENDATA\n",
        )

    def test_write_priority_and_direction_to_priorities_file(self):
        self.mock_model.priority = self.suffix_cls(
            direction=Suffix.EXPORT, datatype=Suffix.INT
        )
        priority_val = 10
        self._set_suffix_value(
            self.mock_model.priority, self.mock_model.x, priority_val
        )

        self.mock_model.direction = self.suffix_cls(
            direction=Suffix.EXPORT, datatype=Suffix.INT
        )
        direction_val = BranchDirection.down
        self._set_suffix_value(
            self.mock_model.direction, self.mock_model.x, direction_val
        )

        CPLEXSHELL._write_priorities_file(self.mock_cplex_shell, self.mock_model)
        priorities_file = self.get_priorities_file_as_string(self.mock_cplex_shell)

        self.assertEqual(
            priorities_file,
            "* ENCODING=ISO-8859-1\n"
            "NAME             Priority Order\n"
            " DN x 10\n"
            "ENDATA\n",
        )

    def test_raise_due_to_invalid_priority(self):
        self.mock_model.priority = self.suffix_cls(
            direction=Suffix.EXPORT, datatype=Suffix.INT
        )
        self._set_suffix_value(self.mock_model.priority, self.mock_model.x, -1)
        with self.assertRaises(ValueError):
            CPLEXSHELL._write_priorities_file(self.mock_cplex_shell, self.mock_model)

        self._set_suffix_value(self.mock_model.priority, self.mock_model.x, 1.1)
        with self.assertRaises(ValueError):
            CPLEXSHELL._write_priorities_file(self.mock_cplex_shell, self.mock_model)

    def test_use_default_due_to_invalid_direction(self):
        self.mock_model.priority = self.suffix_cls(
            direction=Suffix.EXPORT, datatype=Suffix.INT
        )
        priority_val = 10
        self._set_suffix_value(
            self.mock_model.priority, self.mock_model.x, priority_val
        )

        self.mock_model.direction = self.suffix_cls(
            direction=Suffix.EXPORT, datatype=Suffix.INT
        )
        self._set_suffix_value(
            self.mock_model.direction, self.mock_model.x, "invalid_branching_direction"
        )

        CPLEXSHELL._write_priorities_file(self.mock_cplex_shell, self.mock_model)
        priorities_file = self.get_priorities_file_as_string(self.mock_cplex_shell)

        self.assertEqual(
            priorities_file,
            "* ENCODING=ISO-8859-1\n"
            "NAME             Priority Order\n"
            "  x 10\n"
            "ENDATA\n",
        )


class CPLEXShellWritePrioritiesFileKernel(CPLEXShellWritePrioritiesFile):
    suffix_cls = pmo.suffix

    @staticmethod
    def _set_suffix_value(suffix, variable, value):
        suffix[variable] = value

    def get_mock_model(self):
        model = pmo.block()
        model.x = pmo.variable(domain=Binary)
        model.con = pmo.constraint(expr=model.x >= 1)
        model.obj = pmo.objective(expr=model.x)
        return model


class CPLEXShellSolvePrioritiesFile(unittest.TestCase):
    """Integration test on the end-to-end application of priorities via the `Suffix` through a `solve()`"""

    def get_mock_model_with_priorities(self):
        m = ConcreteModel()
        m.x = Var(domain=Integers)
        m.s = RangeSet(0, 9)
        m.y = Var(m.s, domain=Integers)
        m.o = Objective(expr=m.x + sum(m.y), sense=minimize)
        m.c = Constraint(expr=m.x >= 1)
        m.c2 = Constraint(expr=quicksum(m.y[i] for i in m.s) >= 10)

        m.priority = Suffix(direction=Suffix.EXPORT, datatype=Suffix.INT)
        m.direction = Suffix(direction=Suffix.EXPORT, datatype=Suffix.INT)

        m.priority.set_value(m.x, 1)

        # Ensure tests work for both options of `expand`
        m.priority.set_value(m.y, 2, expand=False)
        m.direction.set_value(m.y, BranchDirection.down, expand=True)

        m.direction.set_value(m.y[9], BranchDirection.up)
        return m

    def test_use_variable_priorities(self):
        model = self.get_mock_model_with_priorities()
        with SolverFactory("_mock_cplex") as opt:
            opt._presolve(
                model, priorities=True, keepfiles=True, symbolic_solver_labels=True
            )

            with open(opt._priorities_file_name, "r") as ord_file:
                priorities_file = ord_file.read()

        self.assertEqual(
            priorities_file,
            (
                "* ENCODING=ISO-8859-1\n"
                "NAME             Priority Order\n"
                "  x 1\n"
                " DN y(0) 2\n"
                " DN y(1) 2\n"
                " DN y(2) 2\n"
                " DN y(3) 2\n"
                " DN y(4) 2\n"
                " DN y(5) 2\n"
                " DN y(6) 2\n"
                " DN y(7) 2\n"
                " DN y(8) 2\n"
                " UP y(9) 2\n"
                "ENDATA\n"
            ),
        )
        self.assertIn("read %s\n" % (opt._priorities_file_name,), opt._command.script)

    def test_ignore_variable_priorities(self):
        model = self.get_mock_model_with_priorities()
        with SolverFactory("_mock_cplex") as opt:
            opt._presolve(model, priorities=False, keepfiles=True)

            self.assertIsNone(opt._priorities_file_name)
            self.assertNotIn(".ord", opt._command.script)

    def test_can_use_manual_priorities_file_with_lp_solve(self):
        """Test that we can pass an LP file (not a pyomo model) along with a priorities file to `.solve()`"""
        model = self.get_mock_model_with_priorities()

        with SolverFactory("_mock_cplex") as pre_opt:
            pre_opt._presolve(model, priorities=True, keepfiles=True)
            lp_file = pre_opt._problem_files[0]
            priorities_file_name = pre_opt._priorities_file_name

            with open(priorities_file_name, "r") as ord_file:
                provided_priorities_file = ord_file.read()

        with SolverFactory("_mock_cplex") as opt:
            opt._presolve(
                lp_file,
                priorities=True,
                priorities_file=priorities_file_name,
                keepfiles=True,
            )

            self.assertIn(".ord", opt._command.script)

            with open(opt._priorities_file_name, "r") as ord_file:
                priorities_file = ord_file.read()

        self.assertEqual(priorities_file, provided_priorities_file)


class CPLEXShellSolvePrioritiesFileKernel(CPLEXShellSolvePrioritiesFile):
    def get_mock_model_with_priorities(self):
        m = pmo.block()
        m.x = pmo.variable(domain=Integers)
        m.s = range(10)

        m.y = pmo.variable_list(pmo.variable(domain=Integers) for _ in m.s)

        m.o = pmo.objective(expr=m.x + sum(m.y), sense=minimize)
        m.c = pmo.constraint(expr=m.x >= 1)
        m.c2 = pmo.constraint(expr=quicksum(m.y[i] for i in m.s) >= 10)

        m.priority = pmo.suffix(direction=Suffix.EXPORT, datatype=Suffix.INT)
        m.direction = pmo.suffix(direction=Suffix.EXPORT, datatype=Suffix.INT)

        m.priority[m.x] = 1
        m.priority[m.y] = 2
        m.direction[m.y] = BranchDirection.down
        m.direction[m.y[-1]] = BranchDirection.up
        return m


if __name__ == "__main__":
    unittest.main()
