import gurobipy as gp
from gurobipy import GRB, Model, quicksum

HYPEREDGES = {
    1: ([], ['Xanthine']),
    2: (['Xanthine'], ['p_{0,0}']),
    3: (['Xanthine'], ['p_{0,1}']),
    4: (['Xanthine'], ['p_{0,2}']),
    5: (['p_{0,0}'], ['p_{0,3}']),
    6: (['p_{0,0}'], ['p_{0,4}']),
    7: (['p_{0,1}'], ['p_{0,3}']),
    8: (['p_{0,1}'], ['p_{0,5}']),
    9: (['p_{0,2}'], ['p_{0,4}']),
    10: (['p_{0,2}'], ['p_{0,5}']),
    11: (['p_{0,3}'], ['Caffeine']),
    12: (['p_{0,4}'], ['Caffeine']),
    13: (['p_{0,5}'], ['Caffeine']),
    14: (['Caffeine'], []),
}
#Similyrity of Nodes NMR to Measured NMR
#Can how to add Values along a Path?
#Can you count position on path?
#Can you save spectra values to add hypergraphposition to spectrainformation?
#Compositespectra vs oredered set of spectra
VERTICES1 = {
    1: ([0, 0], ['Xanthine']),
    2: ([0, 0], ['p_{0,0}']),
    3: ([0, 0], ['p_{0,1}']),
    4: ([1, 0], ['p_{0,2}']),
    5: ([0, 0], ['p_{0,3}']),
    6: ([0, 0], ['p_{0,4}']),
    7: ([0, 1], ['p_{0,5}']),
    8: ([0, 0], ['Caffeine']),
}

VERTICES2 = {
    1: ([0], ['Xanthine']),
    2: ([0], ['p_{0,0}']),
    3: ([0], ['p_{0,1}']),
    4: ([1], ['p_{0,2}']),
    5: ([0], ['p_{0,3}']),
    6: ([0], ['p_{0,4}']),
    7: ([1], ['p_{0,5}']),
    8: ([0], ['Caffeine']),
}

VERTICES = ['Xanthine', 'p_{0,0}', 'p_{0,1}', 'p_{0,2}', 'p_{0,3}', 'p_{0,4}', 'p_{0,5}',  'Caffeine']

NMRLIKELYHOODS = [0.0, 0.2, 0.7, 0.1, 0.1, 0.2, 0.7, 0.0]

FIXED_FLOWS = {
    1: 1,
    14: 1,
}

def build_model(name, hyperedges, vertices, nmrlikelihoods,  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}
    n = model.addVars(vertices, vtype=GRB.CONTINUOUS, lb=0.0, ub=1.0, name="nmr")

    for v, nmr in zip(vertices, nmrlikelihoods):
        n[v] = nmr
        print(f'Vertice: {v}, Similarity: {n[v]}')

    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}")

    reaction_path = {}

    #for v_id in vertices:
        #model.addConstr()

    if excluded_support:
        model.addConstr(quicksum(b[e_id] for e_id in excluded_support) <= len(excluded_support) - 1, name="different_hyperedges",)
    
    #Multiplizier den node Wert mit infow + outflow

    #model.setObjective(quicksum(n[v_id] for v_id in vertices), GRB.MINIMIZE) #Maximize Similarity of Nodes
    model.ModelSense = GRB.MAXIMIZE
    model.setObjectiveN(
        quicksum(n[t_id[0]] * x[e_id] for e_id, (_, t_id) in hyperedges.items() if t_id != []),
        index=0,
        priority=2,
        name="maximize_nmr_similarity",
    )   

    model.setObjectiveN(
        quicksum(-1 *b[e_id] for e_id in hyperedges),
        index=1,
        priority=1,
        name="minimize_used_hyperedges",
    )

    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, VERTICES, NMRLIKELYHOODS)
    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, VERTICES)
    else:
        print("No optimal solution found for the second best model.")
 """
if __name__ == "__main__":
    main()