#  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.

import logging
from io import StringIO
from warnings import warn

import pyomo.environ as pe

from miplearn.solvers import _RedirectOutput
from miplearn.solvers.gurobi import GurobiSolver
from miplearn.solvers.pyomo.base import BasePyomoSolver
from . import (
    _get_knapsack_instance,
    get_internal_solvers,
)
from ..fixtures.infeasible import get_infeasible_instance

logger = logging.getLogger(__name__)


def test_redirect_output():
    import sys

    original_stdout = sys.stdout
    io = StringIO()
    with _RedirectOutput([io]):
        print("Hello world")
    assert sys.stdout == original_stdout
    assert io.getvalue() == "Hello world\n"


def test_internal_solver_warm_starts():
    for solver_class in get_internal_solvers():
        logger.info("Solver: %s" % solver_class)
        instance = _get_knapsack_instance(solver_class)
        model = instance.to_model()
        solver = solver_class()
        solver.set_instance(instance, model)
        solver.set_warm_start({"x[0]": 1.0, "x[1]": 0.0, "x[2]": 0.0, "x[3]": 1.0})
        stats = solver.solve(tee=True)
        if stats["Warm start value"] is not None:
            assert stats["Warm start value"] == 725.0
        else:
            warn(f"{solver_class.__name__} should set warm start value")

        solver.set_warm_start({"x[0]": 1.0, "x[1]": 1.0, "x[2]": 1.0, "x[3]": 1.0})
        stats = solver.solve(tee=True)
        assert stats["Warm start value"] is None

        solver.fix({"x[0]": 1.0, "x[1]": 0.0, "x[2]": 0.0, "x[3]": 1.0})
        stats = solver.solve(tee=True)
        assert stats["Lower bound"] == 725.0
        assert stats["Upper bound"] == 725.0


def test_internal_solver():
    for solver_class in get_internal_solvers():
        logger.info("Solver: %s" % solver_class)

        instance = _get_knapsack_instance(solver_class)
        model = instance.to_model()
        solver = solver_class()
        solver.set_instance(instance, model)

        assert solver.get_variable_names() == ["x[0]", "x[1]", "x[2]", "x[3]"]

        stats = solver.solve_lp()
        assert not solver.is_infeasible()
        assert round(stats["LP value"], 3) == 1287.923
        assert len(stats["LP log"]) > 100

        solution = solver.get_solution()
        assert round(solution["x[0]"], 3) == 1.000
        assert round(solution["x[1]"], 3) == 0.923
        assert round(solution["x[2]"], 3) == 1.000
        assert round(solution["x[3]"], 3) == 0.000

        stats = solver.solve(tee=True)
        assert not solver.is_infeasible()
        assert len(stats["MIP log"]) > 100
        assert stats["Lower bound"] == 1183.0
        assert stats["Upper bound"] == 1183.0
        assert stats["Sense"] == "max"
        assert isinstance(stats["Wallclock time"], float)

        solution = solver.get_solution()
        assert solution["x[0]"] == 1.0
        assert solution["x[1]"] == 0.0
        assert solution["x[2]"] == 1.0
        assert solution["x[3]"] == 1.0

        # Add a brand new constraint
        if isinstance(solver, BasePyomoSolver):
            model.cut = pe.Constraint(expr=model.x[0] <= 0.0, name="cut")
            solver.add_constraint(model.cut)
        elif isinstance(solver, GurobiSolver):
            x = model.getVarByName("x[0]")
            solver.add_constraint(x <= 0.0, name="cut")
        else:
            raise Exception("Illegal state")

        # New constraint should affect solution and should be listed in
        # constraint ids
        assert solver.get_constraint_ids() == ["eq_capacity", "cut"]
        stats = solver.solve()
        assert stats["Lower bound"] == 1030.0

        assert solver.get_sense() == "max"
        assert solver.get_constraint_sense("cut") == "<"
        assert solver.get_constraint_sense("eq_capacity") == "<"

        # Verify slacks
        assert solver.get_inequality_slacks() == {
            "cut": 0.0,
            "eq_capacity": 3.0,
        }

        if isinstance(solver, GurobiSolver):
            # Extract the new constraint
            cobj = solver.extract_constraint("cut")

            # New constraint should no longer affect solution and should no longer
            # be listed in constraint ids
            assert solver.get_constraint_ids() == ["eq_capacity"]
            stats = solver.solve()
            assert stats["Lower bound"] == 1183.0

            # New constraint should not be satisfied by current solution
            assert not solver.is_constraint_satisfied(cobj)

            # Re-add constraint
            solver.add_constraint(cobj)

            # Constraint should affect solution again
            assert solver.get_constraint_ids() == ["eq_capacity", "cut"]
            stats = solver.solve()
            assert stats["Lower bound"] == 1030.0

            # New constraint should now be satisfied
            assert solver.is_constraint_satisfied(cobj)

            # Relax problem and make cut into an equality constraint
            solver.relax()
            solver.set_constraint_sense("cut", "=")
            stats = solver.solve()
            assert round(stats["Lower bound"]) == 1030.0
            assert round(solver.get_dual("eq_capacity")) == 0.0


def test_relax():
    for solver_class in get_internal_solvers():
        instance = _get_knapsack_instance(solver_class)
        solver = solver_class()
        solver.set_instance(instance)
        solver.relax()
        stats = solver.solve()
        assert round(stats["Lower bound"]) == 1288.0


def test_infeasible_instance():
    for solver_class in get_internal_solvers():
        instance = get_infeasible_instance(solver_class)
        solver = solver_class()
        solver.set_instance(instance)
        stats = solver.solve()

        assert solver.is_infeasible()
        assert solver.get_solution() is None
        assert stats["Upper bound"] is None
        assert stats["Lower bound"] is None

        stats = solver.solve_lp()
        assert solver.get_solution() is None
        assert stats["LP value"] is None


def test_iteration_cb():
    for solver_class in get_internal_solvers():
        logger.info("Solver: %s" % solver_class)
        instance = _get_knapsack_instance(solver_class)
        solver = solver_class()
        solver.set_instance(instance)
        count = 0

        def custom_iteration_cb():
            nonlocal count
            count += 1
            return count < 5

        solver.solve(iteration_cb=custom_iteration_cb)
        assert count == 5