# MIPLearn, an extensible framework for Learning-Enhanced Mixed-Integer Optimization
# Copyright (C) 2019-2020 Argonne National Laboratory. All rights reserved.
# Written by Alinson S. Xavier <axavier@anl.gov>

from miplearn import BranchPriorityComponent, LearningSolver
from miplearn.problems.knapsack import KnapsackInstance
import numpy as np
import tempfile

def _get_instances():
    return [
        KnapsackInstance(
            weights=[23., 26., 20., 18.],
            prices=[505., 352., 458., 220.],
            capacity=67.,
        ),
    ] * 2


def test_branching():
    instances = _get_instances()
    component = BranchPriorityComponent()
    for instance in instances:
        component.after_solve(None, instance, None)
    component.fit(None)
    for key in ["default"]:
        assert key in component.x_train.keys()
        assert key in component.y_train.keys()
        assert component.x_train[key].shape == (8,  4)
        assert component.y_train[key].shape == (8,  1)
        
        
def test_branch_priority_save_load():
    state_file = tempfile.NamedTemporaryFile(mode="r")
    solver = LearningSolver(components={"branch-priority": BranchPriorityComponent()})
    solver.parallel_solve(_get_instances(), n_jobs=2)
    solver.fit()
    comp = solver.components["branch-priority"]
    assert comp.x_train["default"].shape == (8, 4)
    assert comp.y_train["default"].shape == (8, 1)
    assert "default" in comp.predictors.keys()
    solver.save_state(state_file.name)
    
    solver = LearningSolver(components={"branch-priority": BranchPriorityComponent()})
    solver.load_state(state_file.name)
    comp = solver.components["branch-priority"]
    assert comp.x_train["default"].shape == (8, 4)
    assert comp.y_train["default"].shape == (8, 1)
    assert "default" in comp.predictors.keys()