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

HYPERGRAPH = {
    1: ([], ['Xanthine']),
    2: (['Xanthine'], ['1-Methylxanthine']),
    3: (['Xanthine'], ['3-Methylxanthine']),
    4: (['Xanthine'], ['7-Methylxanthine']),
    5: (['1-Methylxanthine'], ['Theophylline']),
    6: (['1-Methylxanthine'], ['Paraxanthine']),
    7: (['3-Methylxanthine'], ['Theophylline']),
    8: (['3-Methylxanthine'], ['Theobromine']),
    9: (['7-Methylxanthine'], ['Paraxanthine']),
    10: (['7-Methylxanthine'], ['Theobromine']),
    11: (['Theophylline'], ['Caffeine']),
    12: (['Paraxanthine'], ['Caffeine']),
    13: (['Theobromine'], ['Caffeine']),
    14: (['Caffeine'], []),
}

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

def build_model(name, hyperedges, vertices, ele, el1, el2,  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] * heads.count(v) for e_id, (_, heads) in hyperedges.items())
        outflow = quicksum(x[e_id] * tails.count(v) for e_id, (tails, _) in hyperedges.items())
        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 = {}


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

    #Excluding creation and destruction only three reactions for three nmr
    model.addConstr(quicksum(b[e_id] for e_id, (tails, heads) in hyperedges.items() if tails != [] and heads != []) == 3)

    model.setObjective(quicksum(1000 * ele[e_id] * b[e_id] - x[e_id] for e_id in hyperedges),GRB.MAXIMIZE)

    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():
    VERTICES = ['Xanthine', '1-Methylxanthine', '3-Methylxanthine', '7-Methylxanthine', 'Theophylline', 'Paraxanthine', 'Theobromine',  'Caffeine']

    #Results for comparison with 7-Methylxanthine bei shift von 2.5
    NMR1 = [0.46, 0.49, 0.43, 0.6, 0.49, 0.6, 0.45, 0.51]
    #Results for comparison with 3,7-Methylxanthine bei shift von 2.5
    NMR2 = [0.45, 0.52, 0.52, 0.63, 0.52, 0.57, 0.59, 0.59]

    #Chosable parameters
    modes = ["Product", "Average"]
    mode = modes[0] 
    normalize = False

    if normalize:
        NMR1 = [round(l/sum(NMR1), 2) for l in NMR1]
        NMR2 = [round(l/sum(NMR2), 2) for l in NMR2]
    #Change to only have one likelihood, where with manual order association
    NMRE = [0.0, NMR1[1], NMR1[2], NMR1[3], NMR2[4], NMR2[5], NMR2[6], 0.0]

    #Kombiniert Molekül mit Wahrscheinlichkeit für NMR1:
    VERTICE1 = {}
    for vertice, likelihood in zip(VERTICES, NMR1):
        VERTICE1[vertice] = likelihood
    #Kombiniert Molekül mit Wahrscheinlichkeit für NMR2:
    VERTICE2 = {}
    for vertice, likelihood in zip(VERTICES, NMR2):
        VERTICE2[vertice] = likelihood
    #Kombiniert Molekül mit Wahrscheinlichkeit für NMR Empiric:
    VERTICEE = {}
    for vertice, likelihood in zip(VERTICES, NMRE):
        VERTICEE[vertice] = likelihood
    
    #Kantenwahrscheinlichkeiten für NMR1:
    EDGE1 = {}
    #Kantenwahrscheinlichkeiten für NMR2:
    EDGE2 = {}
    #Maximum der verschiedenen Kantenwahrscheinlichkeiten:
    EDGEE = {}

    for edge, (tails, heads) in HYPERGRAPH.items():
        if heads == [] or tails == []:
            EDGE1[edge] = 0.0
            EDGE2[edge] = 0.0
            EDGEE[edge] = 0.0
        elif mode == "Product":
            #print([VERTICE1[head] for head in heads])
            EDGE1[edge] = round(np.prod([VERTICE1[head] for head in heads]),2)
            EDGE2[edge] = round(np.prod([VERTICE2[head] for head in heads]),2)
            EDGEE[edge] = round(np.prod([VERTICEE[head] for head in heads]),2)
        elif mode == "Average":
            EDGE1[edge] = round(sum(VERTICE1[head] for head in heads)/len(heads),2)
            EDGE2[edge] = round(sum(VERTICE2[head] for head in heads)/len(heads),2)
            EDGEE[edge] = round(sum(VERTICEE[head] for head in heads)/len(heads),2)


    model, x, b = build_model("HypergraphFlow", HYPERGRAPH, VERTICES, EDGEE, EDGE1, EDGE2)
    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, HYPERGRAPH)

    """ excluded_support = list(optimal_binary_solution.keys())
    second_model, x2, b2 = build_model("SecondBestHypergraphFlow", HYPERGRAPH, 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, HYPERGRAPH, VERTICES)
    else:
        print("No optimal solution found for the second best model.")
    """
if __name__ == "__main__":
    main()