# ____________________________________________________________________________________
#
# 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 logging

from pyomo.common.collections import Bunch
from pyomo.opt import SolverFactory
from pyomo.opt.base.solvers import UnknownSolver

import pyomo.environ
from pyomo.solvers.plugins.solvers.GUROBI import GUROBISHELL
from pyomo.solvers.plugins.solvers.BARON import BARONSHELL
from pyomo.solvers.plugins.solvers.mosek_direct import MOSEKDirect

# ----------------------------------------------------------------

_test_solver_cases = {}


# ----------------------------------------------------------------

licensed_solvers_with_demo_mode = {'baron'}


#
# NOTE: we initialize the test case, since different
# interfaces may be used for the same "solver"
#
def initialize(**kwds):
    obj = Bunch(**kwds)
    #
    # Set obj.available
    #
    try:
        opt = SolverFactory(obj.name, solver_io=obj.io)
    except:
        opt = None

    if opt is None or isinstance(opt, UnknownSolver):
        obj.available = False
    elif not opt.available(exception_flag=False):
        obj.available = False
    elif hasattr(opt, 'executable') and opt.executable() is None:
        obj.available = False
    elif not opt.license_is_valid() and obj.name not in licensed_solvers_with_demo_mode:
        obj.available = False
    else:
        obj.available = True
    #
    # Set the limits for the solver's "demo" (unlicensed) mode:
    #   ( nVars, nCons, nNonZeros )
    obj.demo_limits = (None, None, None)
    if obj.available:
        if obj.name == "baron" and not opt.license_is_valid():
            obj.demo_limits = (10, 10, 50)
    #
    # Check capabilities, even if the solver is not available
    #
    if not (opt is None or isinstance(opt, UnknownSolver)):
        for _c in obj.capabilities:
            if not _c in opt._capabilities:
                raise ValueError(
                    "Solver %s does not support capability %s!" % (obj.name, _c)
                )
    #
    # Get version
    #
    if obj.available:
        obj.version = opt.version()
    return obj


def test_solver_cases(*args):
    """
    A function for accessing _test_solver_casess as global state
    """
    if len(_test_solver_cases) == 0:
        logging.disable(logging.WARNING)

        #
        # MOSEK
        #

        _mosek_capabilities = set(
            [
                'linear',
                'integer',
                'quadratic_objective',
                'quadratic_constraint',
                'conic_constraints',
            ]
        )

        _test_solver_cases['mosek', 'python'] = initialize(
            name='mosek',
            io='python',
            capabilities=_mosek_capabilities,
            import_suffixes=['dual', 'rc', 'slack'],
        )

        #
        # MOSEK Persistent
        #
        _test_solver_cases['mosek', 'persistent'] = initialize(
            name='mosek',
            io='persistent',
            capabilities=_mosek_capabilities,
            import_suffixes=['dual', 'rc', 'slack'],
        )

        #
        # CPLEX
        #

        _cplex_capabilities = set(
            [
                'linear',
                'integer',
                'quadratic_objective',
                'quadratic_constraint',
                'sos1',
                'sos2',
            ]
        )

        _test_solver_cases['cplex', 'lp'] = initialize(
            name='cplex',
            io='lp',
            capabilities=_cplex_capabilities,
            import_suffixes=['slack', 'dual', 'rc'],
        )

        _test_solver_cases['cplex', 'mps'] = initialize(
            name='cplex',
            io='mps',
            capabilities=_cplex_capabilities,
            import_suffixes=['slack', 'dual', 'rc'],
        )

        _test_solver_cases['cplex', 'nl'] = initialize(
            name='cplex',
            io='nl',
            capabilities=_cplex_capabilities,
            import_suffixes=['dual'],
        )

        _test_solver_cases['cplex', 'python'] = initialize(
            name='cplex',
            io='python',
            capabilities=_cplex_capabilities,
            import_suffixes=['slack', 'dual', 'rc'],
        )

        #
        # CPLEX PERSISTENT
        #

        _test_solver_cases['cplex_persistent', 'python'] = initialize(
            name='cplex_persistent',
            io='python',
            capabilities=_cplex_capabilities,
            import_suffixes=['slack', 'dual', 'rc'],
        )

        #
        # GAMS
        #

        _gams_capabilities = set(
            ['linear', 'integer', 'quadratic_objective', 'quadratic_constraint']
        )

        _test_solver_cases['gams', 'gms'] = initialize(
            name='gams',
            io='gms',
            capabilities=_gams_capabilities,
            import_suffixes=['dual', 'rc'],
        )

        _test_solver_cases['gams', 'python'] = initialize(
            name='gams',
            io='python',
            capabilities=_gams_capabilities,
            import_suffixes=['dual', 'rc'],
        )

        #
        # GAMS V2
        #

        _gams_v2_capabilities = set(['linear', 'integer'])

        _test_solver_cases['gams_v2', 'gms'] = initialize(
            name='gams_v2',
            io='gms',
            capabilities=_gams_v2_capabilities,
            import_suffixes=['dual', 'rc'],
        )

        _test_solver_cases['gams_v2', 'python'] = initialize(
            name='gams_v2',
            io='python',
            capabilities=_gams_v2_capabilities,
            import_suffixes=['dual', 'rc'],
        )

        #
        # GUROBI
        #
        # **NOTE: Gurobi does not handle quadratic constraints before
        #         Major Version 5
        #
        _gurobi_capabilities = set(
            [
                'linear',
                'integer',
                'quadratic_objective',
                'quadratic_constraint',
                'sos1',
                'sos2',
            ]
        )

        _test_solver_cases['gurobi', 'lp'] = initialize(
            name='gurobi',
            io='lp',
            capabilities=_gurobi_capabilities,
            import_suffixes=['slack', 'dual', 'rc'],
        )

        _test_solver_cases['gurobi', 'mps'] = initialize(
            name='gurobi',
            io='mps',
            capabilities=_gurobi_capabilities,
            import_suffixes=['slack', 'dual', 'rc'],
        )

        _test_solver_cases['gurobi', 'nl'] = initialize(
            name='gurobi',
            io='nl',
            capabilities=_gurobi_capabilities,
            options={'qcpdual': 1, 'simplex': 1},
            import_suffixes=['dual'],
        )

        _test_solver_cases['gurobi', 'python'] = initialize(
            name='gurobi',
            io='python',
            capabilities=_gurobi_capabilities,
            import_suffixes=['slack', 'dual', 'rc'],
        )

        #
        # Gurobi PERSISTENT
        #

        _test_solver_cases['gurobi_persistent', 'python'] = initialize(
            name='gurobi_persistent',
            io='python',
            capabilities=_gurobi_capabilities,
            import_suffixes=['slack', 'dual', 'rc'],
        )

        #
        # GLPK
        #
        _glpk_capabilities = set(['linear', 'integer'])

        if 'GLPKSHELL_old' in str(pyomo.solvers.plugins.solvers.GLPK.GLPK().__class__):
            glpk_import_suffixes = ['dual']
        else:
            glpk_import_suffixes = ['rc', 'dual']

        _test_solver_cases['glpk', 'lp'] = initialize(
            name='glpk',
            io='lp',
            capabilities=_glpk_capabilities,
            import_suffixes=glpk_import_suffixes,
        )

        _test_solver_cases['glpk', 'mps'] = initialize(
            name='glpk',
            io='mps',
            capabilities=_glpk_capabilities,
            import_suffixes=glpk_import_suffixes,
            io_options={"skip_objective_sense": True},
        )

        #
        # CBC
        #
        _cbc_lp_capabilities = set(['linear', 'integer'])

        _test_solver_cases['cbc', 'lp'] = initialize(
            name='cbc',
            io='lp',
            capabilities=_cbc_lp_capabilities,
            import_suffixes=['dual', 'rc'],
        )

        _cbc_nl_capabilities = set(['linear', 'integer', 'sos1', 'sos2'])

        _test_solver_cases['cbc', 'nl'] = initialize(
            name='cbc',
            io='nl',
            capabilities=_cbc_nl_capabilities,
            import_suffixes=['dual'],
        )

        # _cbc_mps_capabilities = set(['linear', 'integer', 'sos1', 'sos2'])

        # _test_solver_cases['cbc', 'mps'] = initialize(
        # name='cbc',
        # io='mps',
        # capabilities=_cbc_mps_capabilities,
        # import_suffixes=['dual', 'rc'])

        #
        # XPRESS
        #
        _xpress_capabilities = set(
            [
                'linear',
                'integer',
                'quadratic_objective',
                'quadratic_constraint',
                'sos1',
                'sos2',
            ]
        )

        _test_solver_cases['xpress', 'python'] = initialize(
            name='xpress_direct',
            io='python',
            capabilities=_xpress_capabilities,
            import_suffixes=['dual', 'rc', 'slack'],
            options={'bargapstop': 1e-9},
        )

        #
        # XPRESS PERSISTENT
        #

        _test_solver_cases['xpress_persistent', 'python'] = initialize(
            name='xpress_persistent',
            io='python',
            capabilities=_xpress_capabilities,
            import_suffixes=['slack', 'dual', 'rc'],
            options={'bargapstop': 1e-9},
        )

        #
        # IPOPT
        #
        _ipopt_capabilities = set(
            ['linear', 'quadratic_objective', 'quadratic_constraint']
        )

        _test_solver_cases['ipopt', 'nl'] = initialize(
            name='ipopt',
            io='nl',
            capabilities=_ipopt_capabilities,
            import_suffixes=['dual'],
        )

        #
        # SCIP
        #
        _scip_capabilities = set(
            [
                'linear',
                'integer',
                'quadratic_objective',
                'quadratic_constraint',
                'sos1',
                'sos2',
            ]
        )

        _test_solver_cases['scip', 'nl'] = initialize(
            name='scip', io='nl', capabilities=_scip_capabilities, import_suffixes=[]
        )

        #
        # CONOPT
        #
        _conopt_capabilities = set(
            [
                'linear',
                'integer',
                'quadratic_objective',
                'quadratic_constraint',
                'sos1',
                'sos2',
            ]
        )
        _test_solver_cases['conopt', 'nl'] = initialize(
            name='conopt',
            io='nl',
            capabilities=_conopt_capabilities,
            import_suffixes=[],
        )

        #
        # BARON
        #
        _baron_capabilities = set(
            ['linear', 'integer', 'quadratic_objective', 'quadratic_constraint']
        )

        _test_solver_cases['baron', 'bar'] = initialize(
            name='baron',
            io='bar',
            capabilities=_baron_capabilities,
            import_suffixes=['rc', 'dual'],
        )

        #
        # KNITROAMPL
        #
        _knitroampl_capabilities = set(
            ['linear', 'integer', 'quadratic_objective', 'quadratic_constraint']
        )

        _test_solver_cases['knitroampl', 'nl'] = initialize(
            name='knitroampl',
            io='nl',
            capabilities=_knitroampl_capabilities,
            import_suffixes=['dual'],
        )

        logging.disable(logging.NOTSET)

        #
        # CUOPT
        #
        _cuopt_capabilities = set(['linear', 'integer'])

        _test_solver_cases['cuopt', 'python'] = initialize(
            name='cuopt',
            io='python',
            capabilities=_cuopt_capabilities,
            import_suffixes=['rc', 'dual'],
        )

        #
        # Error Checks
        #
        for sc in _test_solver_cases.values():
            if sc.capabilities is None:
                sc.capabilities = set([])
            if sc.export_suffixes is None:
                sc.export_suffixes = []
            if sc.import_suffixes is None:
                sc.import_suffixes = []
            if sc.options is None:
                sc.options = {}
            if sc.io_options is None:
                sc.io_options = {}
            assert (sc.io is not None) and (type(sc.io) is str)
            assert type(sc.export_suffixes) in [list, tuple]
            assert type(sc.import_suffixes) in [list, tuple]
            assert type(sc.options) is dict
            for tag in sc.export_suffixes:
                assert type(tag) is str
            for tag in sc.import_suffixes:
                assert type(tag) is str
            assert type(sc.capabilities) is set
            for tag in sc.capabilities:
                assert type(tag) is str

    if len(args) == 0:
        return _test_solver_cases.keys()
    return _test_solver_cases[args]
