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

import random
from tempfile import TemporaryDirectory

import numpy as np
from scipy.stats import randint, uniform

from miplearn.h5 import H5File
from miplearn.problems.maxcut import MaxCutGenerator, build_maxcut_model_gurobipy


def _set_seed():
    random.seed(42)
    np.random.seed(42)


def test_maxcut_generator_not_fixed() -> None:
    _set_seed()
    gen = MaxCutGenerator(
        n=randint(low=5, high=6),
        p=uniform(loc=0.5, scale=0.0),
        fix_graph=False,
    )
    data = gen.generate(3)
    assert len(data) == 3
    assert list(data[0].graph.nodes()) == [0, 1, 2, 3, 4]
    assert list(data[0].graph.edges()) == [
        (0, 2),
        (0, 3),
        (0, 4),
        (2, 3),
        (2, 4),
        (3, 4),
    ]
    assert data[0].weights.tolist() == [-1, 1, -1, -1, -1, 1]
    assert list(data[1].graph.nodes()) == [0, 1, 2, 3, 4]
    assert list(data[1].graph.edges()) == [(0, 1), (0, 3), (0, 4), (1, 4), (3, 4)]
    assert data[1].weights.tolist() == [-1, -1, -1, 1, -1]


def test_maxcut_generator_fixed() -> None:
    random.seed(42)
    np.random.seed(42)
    gen = MaxCutGenerator(
        n=randint(low=5, high=6),
        p=uniform(loc=0.5, scale=0.0),
        fix_graph=True,
    )
    data = gen.generate(3)
    assert len(data) == 3
    assert list(data[0].graph.nodes()) == [0, 1, 2, 3, 4]
    assert list(data[0].graph.edges()) == [
        (0, 2),
        (0, 3),
        (0, 4),
        (2, 3),
        (2, 4),
        (3, 4),
    ]
    assert data[0].weights.tolist() == [-1, 1, -1, -1, -1, 1]
    assert list(data[1].graph.nodes()) == [0, 1, 2, 3, 4]
    assert list(data[1].graph.edges()) == [
        (0, 2),
        (0, 3),
        (0, 4),
        (2, 3),
        (2, 4),
        (3, 4),
    ]
    assert data[1].weights.tolist() == [-1, -1, -1, 1, -1, -1]


def test_maxcut_model():
    _set_seed()
    data = MaxCutGenerator(
        n=randint(low=10, high=11),
        p=uniform(loc=0.5, scale=0.0),
        fix_graph=True,
    ).generate(1)[0]
    model = build_maxcut_model_gurobipy(data)

    with TemporaryDirectory() as tempdir:
        with H5File(f"{tempdir}/data.h5", "w") as h5:
            model.extract_after_load(h5)
            obj_lin = h5.get_array("static_var_obj_coeffs")
            assert obj_lin is not None
            assert obj_lin.tolist() == [
                3.0,
                1.0,
                3.0,
                1.0,
                -1.0,
                0.0,
                -1.0,
                0.0,
                -1.0,
                0.0,
            ]
            obj_quad = h5.get_array("static_var_obj_coeffs_quad")
            assert obj_quad is not None
            assert obj_quad.tolist() == [
                [0.0, 0.0, -1.0, 1.0, -1.0, 0.0, 0.0, 0.0, -1.0, -1.0],
                [0.0, 0.0, 1.0, -1.0, 0.0, -1.0, -1.0, 0.0, 0.0, 1.0],
                [0.0, 0.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0, -1.0, -1.0],
                [0.0, 0.0, 0.0, 0.0, 0.0, -1.0, 1.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, 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, 1.0],
                [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, -1.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],
            ]

    model.optimize()
    assert model.inner.ObjVal == -4