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

using GLPK

function gen_tsp()
    np = pyimport("numpy")
    uniform = pyimport("scipy.stats").uniform
    randint = pyimport("scipy.stats").randint
    np.random.seed(42)

    gen = TravelingSalesmanGenerator(
        x=uniform(loc=0.0, scale=1000.0),
        y=uniform(loc=0.0, scale=1000.0),
        n=randint(low=20, high=21),
        gamma=uniform(loc=1.0, scale=0.25),
        fix_cities=true,
        round=true,
    )
    data = gen.generate(1)
    data_filenames = write_pkl_gz(data, "$BASEDIR/../fixtures", prefix="tsp-n20-")
    collector = BasicCollector()
    collector.collect(
        data_filenames,
        data -> build_tsp_model_jump(data, optimizer=GLPK.Optimizer),
        progress=true,
        verbose=true,
    )
end

function test_lazy()
    data_filenames = ["$BASEDIR/../fixtures/tsp-n20-00000.pkl.gz"]
    clf = pyimport("sklearn.dummy").DummyClassifier()
    extractor = H5FieldsExtractor(
        instance_fields=["static_var_obj_coeffs"],
    )
    comp = MemorizingLazyComponent(clf=clf, extractor=extractor)
    solver = LearningSolver(components=[comp])
    solver.fit(data_filenames)
    stats = solver.optimize(
        data_filenames[1],
        data -> build_tsp_model_jump(data, optimizer=GLPK.Optimizer),
    )
    @test stats["Lazy Constraints: AOT"] > 0
end