import networkx as nx
import gurobipy as gp
from gurobipy import GRB


def callback(model, where):
    if where == GRB.Callback.MIPSOL:

        G2 = nx.DiGraph()

        x, offsets, G = model._vars

        vals = model.cbGetSolution(x)

        for fr, to in x.keys():
            if vals[fr, to] > 0.5:
                G2.add_edge(fr, to, weight=G[fr][to]["weight"])

        for cycle in nx.simple_cycles(G2):

            print(cycle)

            weight = sum(G2[fr][to]["weight"]
                         for fr, to in zip(cycle, cycle[1:] + [cycle[0]]))

            weight_single = weight / len(cycle)

            for fr, to in zip(cycle, cycle[1:] + [cycle[0]]):
                model.cbLazy(
                    offsets[fr, to] <= -weight_single
                )

                print((fr, to), "<=", -weight_single)


def augment_graph_lp(G):

    m = gp.Model("augment")

    x = m.addVars(G.edges(), vtype=GRB.BINARY, name="x")
    offsets = m.addVars(G.edges(), vtype=GRB.CONTINUOUS, name="offsets")

    for fr, to in G.edges():
        m.addConstr(offsets[fr, to] <= 0)

    # add deg 2 constraints

    m.addConstrs(
        x.sum(fr, "*") == x.sum("*", fr)
        for fr in G.nodes()
    )

    m.addConstrs(
        x.sum(fr, "*") <= 1 for fr in G.nodes()
    )

    m.addConstrs(
        x.sum("*", to) <= 1 for to in G.nodes()
    )

    obj = gp.quicksum(x[fr, to] * G[fr][to]["weight"] + offsets[fr, to] * x[fr, to]
                      for fr, to in G.edges())

    m.setObjective(obj, GRB.MAXIMIZE)

    m._vars = x, offsets, G
    m.params.LazyConstraints = 1
    m.params.OutputFlag = 0
    m.optimize(callback)

    print("Optimal value:", m.objVal)

    for offset in offsets.values():
        print(offset.x)


def main():
    G = nx.DiGraph()

    G.add_edge("USD", "ETH", weight=-1)
    G.add_edge("ETH", "BNB", weight=-1)
    G.add_edge("BNB", "BTC", weight=1)
    G.add_edge("BTC", "EUR", weight=-1)
    G.add_edge("EUR", "ETH", weight=2)

    # print(nx.negative_edge_cycle(G))

    # print(get_best_sol(G))

    # G2, offset = augment_graph(G)

    # print(offset)

    augment_graph_lp(G)


if __name__ == "__main__":
    main()