# ____________________________________________________________________________________
#
# 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 pyomo.kernel as pmo
from pyomo.core import (
    ConcreteModel,
    Var,
    Objective,
    Constraint,
    RangeSet,
    ConstraintList,
)
from pyomo.solvers.tests.models.base import _BaseTestModel, register_model
from pyomo.repn.beta.matrix import compile_block_linear_constraints

has_numpy = False
try:
    import numpy

    has_numpy = True
except:
    pass

has_scipy = False
try:
    import scipy
    import scipy.sparse

    has_scipy = True
except:
    pass


@register_model
class LP_compiled(_BaseTestModel):
    """
    A continuous linear model that is compiled into a
    MatrixConstraint object.
    """

    description = "LP_compiled"
    capabilities = set(['linear'])
    test_pickling = False
    size = (13, 22, None)

    def __init__(self):
        _BaseTestModel.__init__(self)
        self.add_results(self.description + ".json")
        self.disable_suffix_tests = True

    def _generate_model(self):
        self.model = ConcreteModel()
        model = self.model
        model._name = self.description

        model.s = RangeSet(1, 12)
        model.x = Var(model.s)
        model.x[1].setlb(-1)
        model.x[1].setub(1)
        model.x[2].setlb(-1)
        model.x[2].setub(1)
        model.obj = Objective(
            expr=sum(model.x[i] * ((-1) ** (i + 1)) for i in model.x.index_set())
        )
        model.c = ConstraintList()
        # to make the variable used in the constraint match the name
        model.c.add(Constraint.Skip)
        model.c.add(Constraint.Skip)
        model.c.add(model.x[3] >= -1.0)
        model.c.add(model.x[4] <= 1.0)
        model.c.add(model.x[5] == -1.0)
        model.c.add(model.x[6] == -1.0)
        model.c.add(model.x[7] == 1.0)
        model.c.add(model.x[8] == 1.0)
        model.c.add((-1.0, model.x[9], -1.0))
        model.c.add((-1.0, model.x[10], -1.0))
        model.c.add((1.0, model.x[11], 1.0))
        model.c.add((1.0, model.x[12], 1.0))
        cdata = model.c.add((0, 1, 3))
        assert cdata.lower == 0
        assert cdata.upper == 3
        assert cdata.body() == 1
        assert not cdata.equality
        cdata = model.c.add((0, 2, 3))
        assert cdata.lower == 0
        assert cdata.upper == 3
        assert cdata.body() == 2
        assert not cdata.equality
        cdata = model.c.add((0, 1, None))
        assert cdata.lower is None
        assert cdata.upper == 1
        assert cdata.body() == 0
        assert not cdata.equality
        cdata = model.c.add((None, 0, 1))
        assert cdata.lower is None
        assert cdata.upper == 1
        assert cdata.body() == 0
        assert not cdata.equality
        cdata = model.c.add((1, 1))
        assert cdata.lower == 1
        assert cdata.upper == 1
        assert cdata.body() == 1
        assert cdata.equality

        model.fixed_var = Var()
        model.fixed_var.fix(1.0)
        cdata = model.c.add((0, 1 + model.fixed_var, 3))
        cdata = model.c.add((0, 2 + model.fixed_var, 3))
        cdata = model.c.add((0, model.fixed_var, None))
        cdata = model.c.add((None, model.fixed_var, 1))
        cdata = model.c.add((model.fixed_var, 1))

        model.c_inactive = ConstraintList()
        # to make the variable used in the constraint match the name
        model.c_inactive.add(Constraint.Skip)
        model.c_inactive.add(Constraint.Skip)
        model.c_inactive.add(model.x[3] >= -2.0)
        model.c_inactive.add(model.x[4] <= 2.0)

        compile_block_linear_constraints(model, 'Amatrix')

    def warmstart_model(self):
        assert self.model is not None
        model = self.model
        for i in model.s:
            model.x[i].value = None


if has_numpy and has_scipy:
    # TODO: we need to somehow label this as a skip rather
    #       than not defining the test class

    @register_model
    class LP_compiled_dense_kernel(LP_compiled):
        def _get_dense_data(self):
            assert has_numpy and has_scipy
            A = numpy.array(
                [
                    [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                    [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                ],
                dtype=float,
            )
            lb = numpy.array(
                [
                    -1.0,
                    -numpy.inf,
                    -1.0,
                    -1.0,
                    1.0,
                    1.0,
                    -1.0,
                    -1.0,
                    1.0,
                    1.0,
                    -1.0,
                    -2.0,
                    -1.0,
                    -numpy.inf,
                    0.0,
                    -2.0,
                    -3.0,
                    -1.0,
                    -numpy.inf,
                    0.0,
                    -2.0,
                    -numpy.inf,
                ]
            )
            ub = numpy.array(
                [
                    numpy.inf,
                    1.0,
                    -1.0,
                    -1.0,
                    1.0,
                    1.0,
                    -1.0,
                    -1.0,
                    1.0,
                    1.0,
                    2.0,
                    1.0,
                    numpy.inf,
                    1.0,
                    0.0,
                    1.0,
                    0.0,
                    numpy.inf,
                    0.0,
                    0.0,
                    numpy.inf,
                    2.0,
                ]
            )
            eq_index = [2, 3, 4, 5, 14, 19]

            return A, lb, ub, eq_index

        def _generate_base_model(self):
            self.model = pmo.block()
            model = self.model
            model._name = self.description

            model.s = list(range(1, 13))
            model.x = pmo.variable_dict(((i, pmo.variable()) for i in model.s))
            model.x[1].lb = -1
            model.x[1].ub = 1
            model.x[2].lb = -1
            model.x[2].ub = 1
            model.obj = pmo.objective(
                expr=sum(model.x[i] * ((-1) ** (i + 1)) for i in model.s)
            )
            variable_order = [
                model.x[3],
                model.x[4],
                model.x[5],
                model.x[6],
                model.x[7],
                model.x[8],
                model.x[9],
                model.x[10],
                model.x[11],
                model.x[12],
            ]

            return variable_order

        def _generate_model(self):
            x = self._generate_base_model()
            model = self.model
            A, lb, ub, eq_index = self._get_dense_data()
            model.Amatrix = pmo.matrix_constraint(A, lb=lb, ub=ub, x=x, sparse=False)
            for i in eq_index:
                assert model.Amatrix[i].lb == model.Amatrix[i].ub
                model.Amatrix[i].rhs = model.Amatrix[i].lb

    @register_model
    class LP_compiled_sparse_kernel(LP_compiled_dense_kernel):
        def _generate_model(self):
            x = self._generate_base_model()
            model = self.model
            A, lb, ub, eq_index = self._get_dense_data()
            model.Amatrix = pmo.matrix_constraint(A, lb=lb, ub=ub, x=x, sparse=True)
            for i in eq_index:
                assert model.Amatrix[i].lb == model.Amatrix[i].ub
                model.Amatrix[i].rhs = model.Amatrix[i].lb
