from sage.all import *
from sage.matroids.all import *
from sage.graphs.all import *
import gurobipy as gp
from gurobipy import GRB

env = gp.Env(empty=True)
env.setParam("OutputFlag",0)
env.start()

def base_hittingset_with_callback(M, integral=true):
    # model
    model = gp.Model("mip1",env=env)
    groundset = M.groundset()
    x = dict()
    if integral:
        for e in groundset: x[e]=model.addVar(vtype=GRB.BINARY)
    else:
        for e in groundset: x[e]=model.addVar(vtype=GRB.CONTINUOUS,lb=0,ub=1)
    # there is no lazy constraint for LP...
    for B in M.bases():
        model.addConstr(gp.quicksum([x[e] for e in B])>=1)
    model.setObjective(gp.quicksum([x[e] for e in groundset]), GRB.MINIMIZE)
    # callback
    def callback_func(model, where):
        if integral and where == GRB.Callback.MIPSOL:
            sol_values = {key: model.cbGetSolution(var) 
                          for key, var in x.items()}
            # find min weight base in the matroid
            base = frozenset()
            minweight=0
            for e in sorted(groundset, key=lambda ee: sol_values[ee]):
                if M.is_independent(base.union([e])):
                    base=base.union([e])
                    minweight=minweight+sol_values[e]
            if minweight < 1:
                model.cbLazy(gp.quicksum([x[e] for e in base])>=1)
    # solve
    if integral:
        model.Params.Heuristics = 0
        model.Params.LazyConstraints = 1
        model.optimize(callback_func)
    else:
        model.optimize()
    return model.ObjVal

# tests
if __name__ == "__main__":
    def cogirthip(bases, integral=true):
        model = gp.Model("mip1",env=env)
        groundset=frozenset()
        for B in bases: groundset=groundset|frozenset(B)
        x = dict()
        if integral:
            for e in groundset: x[e]=model.addVar(vtype=GRB.BINARY)
        else:
            for e in groundset: x[e]=model.addVar(vtype=GRB.CONTINUOUS,lb=0)
        model.setObjective(gp.quicksum([x[e] for e in groundset]), GRB.MINIMIZE)
        for B in bases:
            model.addConstr(gp.quicksum([x[e] for e in B])>=1)
        model.optimize()
        return model.ObjVal
    f = lambda g: g.is_connected()
    for N in range(4,10):
        for g in filter(f, graphs(N)):
            M=Matroid(g)
            cb=base_hittingset_with_callback(M)
            bf=cogirthip(M.bases())
            if randint(0,100)==1:
                print(f"callback:{cb}, bf:{bf}")
            if abs(cb-bf) > 0.01:
                print(f"callback:{cb}, bf:{bf}")
                exit(1)