from gurobipy import GRB, Model, quicksum


HYPEREDGES = {
    3: (["Ribulose-5-Phosphate"], ["p_{0,0}"]),
    6: (["Ribulose-5-Phosphate", "p_{0,0}"], ["p_{0,1}", "p_{0,2}"]),
    9: (["p_{0,1}", "p_{0,2}"], ["Fructose-6-Phosphate", "p_{0,3}"]),
    11: (["p_{0,0}", "p_{0,1}"], ["p_{0,3}", "p_{0,4}"]),
    13: (["Ribulose-5-Phosphate", "p_{0,2}"], ["Fructose-6-Phosphate", "p_{0,5}"]),
    14: (["Fructose-6-Phosphate", "p_{0,3}"], ["Fructose-6-Phosphate", "p_{0,3}"]),
    16: (["Fructose-6-Phosphate", "p_{0,5}"], ["p_{0,3}", "p_{0,6}"]),
    17: (["Fructose-6-Phosphate", "p_{0,2}"], ["Ribulose-5-Phosphate", "p_{0,3}"]),
    18: (["Fructose-6-Phosphate", "p_{0,0}"], ["p_{0,1}", "p_{0,3}"]),
    20: (["p_{0,3}", "p_{0,4}"], ["Fructose-6-Phosphate", "p_{0,7}"]),
    21: (["Ribulose-5-Phosphate", "p_{0,3}"], ["Fructose-6-Phosphate", "p_{0,2}"]),
    22: (["p_{0,1}", "p_{0,3}"], ["Fructose-6-Phosphate", "p_{0,0}"]),
    23: (["p_{0,4}", "p_{0,5}"], ["p_{0,6}", "p_{0,7}"]),
    24: (["p_{0,2}", "p_{0,4}"], ["Ribulose-5-Phosphate", "p_{0,7}"]),
    25: (["p_{0,0}", "p_{0,4}"], ["p_{0,1}", "p_{0,7}"]),
    26: (["Ribulose-5-Phosphate", "p_{0,5}"], ["p_{0,2}", "p_{0,6}"]),
    27: (["p_{0,1}", "p_{0,5}"], ["p_{0,0}", "p_{0,6}"]),
    29: (["p_{0,2}"], ["p_{0,8}"]),
    31: (["p_{0,0}", "p_{0,8}"], ["p_{0,9}"]),
    33: (["p_{0,2}", "p_{0,8}"], ["p_{0,10}"]),
    36: (["H2O", "p_{0,9}"], ["p_{0,11}", "p_{0,12}"]),
    37: (["H2O", "p_{0,9}"], ["p_{0,4}", "p_{0,11}"]),
    39: (["H2O", "p_{0,10}"], ["p_{0,11}", "p_{0,13}"]),
    40: (["H2O", "p_{0,10}"], ["Fructose-6-Phosphate", "p_{0,11}"]),
    41: ([], ["H2O"]),
    42: ([], ["Ribulose-5-Phosphate"]),
    43: (["Fructose-6-Phosphate"], []),
    44: (["p_{0,11}"], []),
}

FIXED_FLOWS = {
    41: 1,
    42: 6,
    43: 5,
    44: 1,
}


def build_model(name, hyperedges, excluded_support=None):
    model = Model(name)

    x = {e_id: model.addVar(vtype=GRB.INTEGER, lb=0, name=f"x_{e_id}") for e_id in hyperedges}
    b = {e_id: model.addVar(vtype=GRB.BINARY, name=f"b_{e_id}") for e_id in hyperedges}

    vertices = set(v for tails, heads in hyperedges.values() for v in tails + heads)

    for v in vertices:
        inflow = quicksum(x[e_id] for e_id, (_, heads) in hyperedges.items() if v in heads)
        outflow = quicksum(x[e_id] for e_id, (tails, _) in hyperedges.items() if v in tails)
        model.addConstr(inflow == outflow, name=f"flow_conservation_{v}")

    for e_id, value in FIXED_FLOWS.items():
        model.addConstr(x[e_id] == value, name=f"fixed_flow_{e_id}")

    for e_id in hyperedges:
        model.addGenConstrIndicator(b[e_id], 0, x[e_id] == 0, name=f"unused_implies_zero_{e_id}")
        model.addConstr(x[e_id] >= b[e_id], name=f"used_implies_positive_flow_{e_id}")

    if excluded_support:
        model.addConstr(
            quicksum(b[e_id] for e_id in excluded_support) <= len(excluded_support) - 1,
            name="different_hyperedges",
        )

    model.setObjective(quicksum(x[e_id] for e_id in hyperedges), GRB.MINIMIZE)
    return model, x, b


def positive_entries(variable_dict, threshold=0.5):
    return {e_id: var.X for e_id, var in variable_dict.items() if var.X > threshold}


def print_solution(title, flow_solution, binary_solution, hyperedges):
    print(f"\n{title}:")
    for e_id in sorted(flow_solution):
        flow = flow_solution[e_id]
        tails, heads = hyperedges[e_id]
        print(f"Hyperedge {e_id}: Flow = {flow}, Tails = {tails}, Heads = {heads}")

    print("\nBinary Variables:")
    for e_id in sorted(binary_solution):
        print(f"Binary Variable b_{e_id} = {binary_solution[e_id]}")

    print(f"\nTotal flow: {sum(flow_solution.values())}")
    print(f"Number of used hyperedges: {len(binary_solution)}")


def main():
    model, x, b = build_model("HypergraphFlow", HYPEREDGES)
    model.optimize()

    if model.status != GRB.Status.OPTIMAL:
        print("No optimal solution found for the first model.")
        return

    optimal_solution = positive_entries(x)
    optimal_binary_solution = positive_entries(b)
    print_solution("Optimal Solution", optimal_solution, optimal_binary_solution, HYPEREDGES)

    excluded_support = list(optimal_binary_solution.keys())
    second_model, x2, b2 = build_model(
        "SecondBestHypergraphFlow",
        HYPEREDGES,
        excluded_support=excluded_support,
    )
    second_model.optimize()

    if second_model.status == GRB.Status.OPTIMAL:
        second_solution = positive_entries(x2)
        second_binary_solution = positive_entries(b2)
        print_solution("Second Best Solution", second_solution, second_binary_solution, HYPEREDGES)
    else:
        print("No optimal solution found for the second best model.")


if __name__ == "__main__":
    main()