#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
#  Released under the modified BSD license. See COPYING.md for more details.

from miplearn.features import (
    FeaturesExtractor,
    InstanceFeatures,
    Constraint,
    VariableFeatures,
)
from miplearn.solvers.gurobi import GurobiSolver
from miplearn.solvers.tests import (
    assert_equals,
    _round_constraints,
    _round,
)

inf = float("inf")


def test_knapsack() -> None:
    solver = GurobiSolver()
    instance = solver.build_test_instance_knapsack()
    model = instance.to_model()
    solver.set_instance(instance, model)
    solver.solve_lp()

    features = FeaturesExtractor().extract(instance, solver)
    assert features.variables is not None
    assert features.constraints is not None
    assert features.instance is not None

    assert_equals(
        _round(features.variables),
        VariableFeatures(
            names=("x[0]", "x[1]", "x[2]", "x[3]", "z"),
            basis_status=("U", "B", "U", "L", "U"),
            categories=("default", "default", "default", "default", None),
            lower_bounds=(0.0, 0.0, 0.0, 0.0, 0.0),
            obj_coeffs=(505.0, 352.0, 458.0, 220.0, 0.0),
            reduced_costs=(193.615385, 0.0, 187.230769, -23.692308, 13.538462),
            sa_lb_down=(-inf, -inf, -inf, -0.111111, -inf),
            sa_lb_up=(1.0, 0.923077, 1.0, 1.0, 67.0),
            sa_obj_down=(311.384615, 317.777778, 270.769231, -inf, -13.538462),
            sa_obj_up=(inf, 570.869565, inf, 243.692308, inf),
            sa_ub_down=(0.913043, 0.923077, 0.9, 0.0, 43.0),
            sa_ub_up=(2.043478, inf, 2.2, inf, 69.0),
            types=("B", "B", "B", "B", "C"),
            upper_bounds=(1.0, 1.0, 1.0, 1.0, 67.0),
            user_features=(
                (23.0, 505.0),
                (26.0, 352.0),
                (20.0, 458.0),
                (18.0, 220.0),
                None,
            ),
            values=(1.0, 0.923077, 1.0, 0.0, 67.0),
            alvarez_2017=[
                [1.0, 0.32899, 0.0, 0.0, 1.0, 1.0, 5.265874, 46.051702],
                [1.0, 0.229316, 0.0, 0.076923, 1.0, 1.0, 3.532875, 5.388476],
                [1.0, 0.298371, 0.0, 0.0, 1.0, 1.0, 5.232342, 46.051702],
                [1.0, 0.143322, 0.0, 0.0, 1.0, -1.0, 46.051702, 3.16515],
                [0.0, 0.0, 0.0, 0.0, 1.0, -1.0, 0.0, 0.0],
            ],
        ),
    )
    assert_equals(
        _round_constraints(features.constraints),
        {
            "eq_capacity": Constraint(
                basis_status="N",
                category="eq_capacity",
                dual_value=13.538462,
                lazy=False,
                lhs={"x[0]": 23.0, "x[1]": 26.0, "x[2]": 20.0, "x[3]": 18.0, "z": -1.0},
                rhs=0.0,
                sa_rhs_down=-24.0,
                sa_rhs_up=1.9999999999999987,
                sense="=",
                slack=0.0,
                user_features=[0.0],
            )
        },
    )
    assert_equals(
        features.instance,
        InstanceFeatures(
            user_features=[67.0, 21.75],
            lazy_constraint_count=0,
        ),
    )