#  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
import os
import tempfile
from typing import List, cast

import dill

from miplearn.instance.base import Instance
from miplearn.instance.picklegz import PickleGzInstance, write_pickle_gz, read_pickle_gz
from miplearn.solvers.gurobi import GurobiSolver
from miplearn.solvers.internal import InternalSolver
from miplearn.solvers.learning import LearningSolver

# noinspection PyUnresolvedReferences
from tests.solvers.test_internal_solver import internal_solvers

logger = logging.getLogger(__name__)


def test_learning_solver(
    internal_solvers: List[InternalSolver],
) -> None:
    for mode in ["exact", "heuristic"]:
        for internal_solver in internal_solvers:
            logger.info("Solver: %s" % internal_solver)
            instance = internal_solver.build_test_instance_knapsack()
            solver = LearningSolver(
                solver=internal_solver,
                mode=mode,
            )

            solver.solve(instance)
            assert hasattr(instance, "features")

            sample = instance.training_data[0]
            assert sample.solution is not None
            assert sample.solution["x[0]"] == 1.0
            assert sample.solution["x[1]"] == 0.0
            assert sample.solution["x[2]"] == 1.0
            assert sample.solution["x[3]"] == 1.0
            assert sample.lower_bound == 1183.0
            assert sample.upper_bound == 1183.0
            assert sample.lp_solution is not None
            assert sample.lp_solution["x[0]"] is not None
            assert sample.lp_solution["x[1]"] is not None
            assert sample.lp_solution["x[2]"] is not None
            assert sample.lp_solution["x[3]"] is not None
            assert round(sample.lp_solution["x[0]"], 3) == 1.000
            assert round(sample.lp_solution["x[1]"], 3) == 0.923
            assert round(sample.lp_solution["x[2]"], 3) == 1.000
            assert round(sample.lp_solution["x[3]"], 3) == 0.000
            assert sample.lp_value is not None
            assert round(sample.lp_value, 3) == 1287.923
            assert sample.mip_log is not None
            assert len(sample.mip_log) > 100

            solver.fit([instance])
            solver.solve(instance)

            # Assert solver is picklable
            with tempfile.TemporaryFile() as file:
                dill.dump(solver, file)


def test_solve_without_lp(
    internal_solvers: List[InternalSolver],
) -> None:
    for internal_solver in internal_solvers:
        logger.info("Solver: %s" % internal_solver)
        instance = internal_solver.build_test_instance_knapsack()
        solver = LearningSolver(
            solver=internal_solver,
            solve_lp=False,
        )
        solver.solve(instance)
        solver.fit([instance])
        solver.solve(instance)


def test_parallel_solve(
    internal_solvers: List[InternalSolver],
) -> None:
    for internal_solver in internal_solvers:
        instances = [internal_solver.build_test_instance_knapsack() for _ in range(10)]
        solver = LearningSolver(solver=internal_solver)
        results = solver.parallel_solve(instances, n_jobs=3)
        assert len(results) == 10
        for instance in instances:
            data = instance.training_data[0]
            assert data.solution is not None
            assert len(data.solution.keys()) == 4


def test_solve_fit_from_disk(
    internal_solvers: List[InternalSolver],
) -> None:
    for internal_solver in internal_solvers:
        # Create instances and pickle them
        instances: List[Instance] = []
        for k in range(3):
            instance = internal_solver.build_test_instance_knapsack()
            with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as file:
                instances += [PickleGzInstance(file.name)]
                write_pickle_gz(instance, file.name)

        # Test: solve
        solver = LearningSolver(solver=internal_solver)
        solver.solve(instances[0])
        instance_loaded = read_pickle_gz(cast(PickleGzInstance, instances[0]).filename)
        assert len(instance_loaded.training_data) > 0
        assert instance_loaded.features.instance is not None
        assert instance_loaded.features.variables is not None
        assert instance_loaded.features.constraints is not None

        # Test: parallel_solve
        solver.parallel_solve(instances)
        for instance in instances:
            instance_loaded = read_pickle_gz(cast(PickleGzInstance, instance).filename)
            assert len(instance_loaded.training_data) > 0
            assert instance_loaded.features.instance is not None
            assert instance_loaded.features.variables is not None
            assert instance_loaded.features.constraints is not None

        # Delete temporary files
        for instance in instances:
            os.remove(cast(PickleGzInstance, instance).filename)


def test_simulate_perfect() -> None:
    internal_solver = GurobiSolver()
    instance = internal_solver.build_test_instance_knapsack()
    with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as tmp:
        write_pickle_gz(instance, tmp.name)
        solver = LearningSolver(
            solver=internal_solver,
            simulate_perfect=True,
        )
        stats = solver.solve(PickleGzInstance(tmp.name))
        assert stats["mip_lower_bound"] == stats["Objective: Predicted lower bound"]


def test_gap() -> None:
    assert LearningSolver._compute_gap(ub=0.0, lb=0.0) == 0.0
    assert LearningSolver._compute_gap(ub=1.0, lb=0.5) == 0.5
    assert LearningSolver._compute_gap(ub=1.0, lb=1.0) == 0.0
    assert LearningSolver._compute_gap(ub=1.0, lb=-1.0) is None
    assert LearningSolver._compute_gap(ub=1.0, lb=None) is None
    assert LearningSolver._compute_gap(ub=None, lb=1.0) is None
    assert LearningSolver._compute_gap(ub=None, lb=None) is None