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

HYPERGRAPH = {
     0: (['Butadien'], []),
     2: (['p_{0,0}'], []),
     3: (['p_{0,1}'], []),
     4: (['Butadien', 'Butadien'], ['p_{0,0}', 'p_{0,1}']),
     6: (['p_{0,2}'], []),
     7: (['Butadien', 'Butadien'], ['p_{0,2}']),
     10: (['p_{0,0}', 'p_{0,1}'], ['Butadien', 'Butadien']),
     12: (['p_{0,3}'], []),
     13: (['p_{0,0}', 'p_{0,2}'], ['p_{0,3}']),
     17: (['p_{0,4}'], []),
     18: (['p_{0,1}', 'p_{0,1}'], ['Butadien', 'p_{0,4}']),
     19: (['p_{0,5}'], []),
     20: (['p_{0,1}', 'p_{0,2}'], ['Butadien', 'p_{0,5}']),
     23: (['Butadien', 'p_{0,1}'], ['p_{0,0}', 'p_{0,4}']),
     26: (['Butadien', 'p_{0,2}'], ['p_{0,0}', 'p_{0,5}']),
     27: (['p_{0,6}'], []),
     28: (['p_{0,0}', 'p_{0,1}'], ['p_{0,6}']),
     29: (['p_{0,7}'], []),
     30: (['Butadien', 'p_{0,0}'], ['p_{0,7}']),
     31: (['p_{0,8}'], []),
     32: (['Butadien', 'p_{0,1}'], ['p_{0,8}']),
     33: (['Butadien', 'p_{0,1}'], ['p_{0,5}']),
     34: (['p_{0,9}'], []),
     35: (['Butadien', 'p_{0,1}'], ['p_{0,9}']),
     36: (['p_{0,10}'], []),
     37: (['Butadien', 'p_{0,1}'], ['p_{0,10}']),
     38: (['p_{0,3}'], ['p_{0,0}', 'p_{0,2}']),
     41: (['p_{0,11}'], []),
     42: (['p_{0,4}'], ['p_{0,0}', 'p_{0,11}']),
     43: (['p_{0,12}'], []),
     44: (['p_{0,5}'], ['p_{0,12}']),
     63: (['p_{0,13}'], []),
     64: (['p_{0,4}', 'p_{0,6}'], ['p_{0,1}', 'p_{0,13}']),
     65: (['p_{0,2}', 'p_{0,4}'], ['p_{0,1}', 'p_{0,5}']),
     66: (['p_{0,0}', 'p_{0,4}'], ['Butadien', 'p_{0,1}']),
     67: (['Butadien', 'p_{0,4}'], ['p_{0,1}', 'p_{0,1}']),
     77: (['p_{0,5}', 'p_{0,6}'], ['p_{0,2}', 'p_{0,13}']),
     79: (['p_{0,0}', 'p_{0,5}'], ['Butadien', 'p_{0,2}']),
     80: (['Butadien', 'p_{0,5}'], ['p_{0,1}', 'p_{0,2}']),
     83: (['p_{0,1}', 'p_{0,5}'], ['p_{0,2}', 'p_{0,4}']),
     96: (['Butadien', 'p_{0,6}'], ['p_{0,0}', 'p_{0,13}']),
     97: (['p_{0,1}', 'p_{0,6}'], ['Butadien', 'p_{0,13}']),
     99: (['p_{0,14}'], []),
     100: (['p_{0,0}', 'p_{0,6}'], ['p_{0,14}']),
     101: (['p_{0,15}'], []),
     102: (['p_{0,0}', 'p_{0,7}'], ['p_{0,15}']),
     121: (['p_{0,0}', 'p_{0,4}'], ['p_{0,13}']),
     122: (['p_{0,16}'], []),
     123: (['p_{0,0}', 'p_{0,4}'], ['p_{0,16}']),
     124: (['p_{0,17}'], []),
     125: (['Butadien', 'p_{0,7}'], ['p_{0,17}']),
     127: (['p_{0,12}'], ['p_{0,5}']),
     130: (['p_{0,14}'], ['p_{0,0}', 'p_{0,6}']),
     132: (['p_{0,15}'], ['p_{0,0}', 'p_{0,7}']),
     135: (['p_{0,0}', 'p_{0,11}'], ['p_{0,4}']),
     153: (['Butadien', 'p_{0,13}'], ['p_{0,1}', 'p_{0,6}']),
     155: (['p_{0,0}', 'p_{0,13}'], ['Butadien', 'p_{0,6}']),
     157: (['p_{0,1}', 'p_{0,13}'], ['p_{0,4}', 'p_{0,6}']),
     158: (['p_{0,2}', 'p_{0,13}'], ['p_{0,5}', 'p_{0,6}']),
     210: (['p_{0,18}'], []),
     211: (['Butadien', 'p_{0,11}'], ['p_{0,18}']),
     213: ([], ['Butadien']),
}

HYPERGRAPH2 = {
    2: (['p_{0,0}'], []),
    3: (['p_{0,1}'], []),
    4: (['Butadien', 'Butadien'], ['p_{0,0}', 'p_{0,1}']),
    6: (['p_{0,2}'], []),
    7: (['Butadien', 'Butadien'], ['p_{0,2}']),
    17: (['p_{0,4}'], []),
    23: (['Butadien', 'p_{0,1}'], ['p_{0,0}', 'p_{0,4}']),
    41: (['p_{0,11}'], []),
    42: (['p_{0,4}'], ['p_{0,0}', 'p_{0,11}']),
    213: ([], ['Butadien']),
}

HYPERGRAPH3 = {
    0: (['Butadien'], []),
    213: ([], ['Butadien']),
    4: (['Butadien', 'Butadien'], ['p_{0,0}', 'Butadien']),
    2: (['p_{0,0}'], []),
}


FIXED_FLOWS = {
    # three Butadiene molecules enter the network
    213: 3,

    # one benzene molecule leaves the network
    41: 1,

    # do not fix edge 4; otherwise part of the expected pathway is already given
    # 4: 1,
}


def build_model(name, hyperedges, vertices, vl1, vl2, vl3, 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
    }

    # NMR variables are now on molecules, not reaction edges
    candidate_vertices = [v for v in vertices if v != "Butadien"]

    n1 = {
        v: model.addVar(vtype=GRB.BINARY, name=f"n1_{v}")
        for v in candidate_vertices
    }
    n2 = {
        v: model.addVar(vtype=GRB.BINARY, name=f"n2_{v}")
        for v in candidate_vertices
    }
    n3 = {
        v: model.addVar(vtype=GRB.BINARY, name=f"n3_{v}")
        for v in candidate_vertices
    }

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

    # flow conservation, unchanged from the original model
    for v in all_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}",
        )

    # link integer flow x[e] and binary "reaction is used" variable b[e]
    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}",
        )

    # each NMR observation is explained by one candidate molecule
    model.addConstr(
        quicksum(n1[v] for v in candidate_vertices) == 1,
        name="one_molecule_for_nmr1",
    )
    model.addConstr(
        quicksum(n2[v] for v in candidate_vertices) == 1,
        name="one_molecule_for_nmr2",
    )
    model.addConstr(
        quicksum(n3[v] for v in candidate_vertices) == 1,
        name="one_molecule_for_nmr3",
    )

    # an NMR-supported molecule must actually be produced by a selected reaction
    for v in candidate_vertices:
        producing_edges = [
            e_id
            for e_id, (tails, heads) in hyperedges.items()
            if tails != [] and v in heads
        ]

        if producing_edges:
            produced = quicksum(b[e_id] for e_id in producing_edges)
            model.addConstr(
                n1[v] <= produced,
                name=f"nmr1_molecule_must_be_produced_{v}",
            )
            model.addConstr(
                n2[v] <= produced,
                name=f"nmr2_molecule_must_be_produced_{v}",
            )
            model.addConstr(
                n3[v] <= produced,
                name=f"nmr3_molecule_must_be_produced_{v}",
            )
        else:
            model.addConstr(n1[v] == 0)
            model.addConstr(n2[v] == 0)
            model.addConstr(n3[v] == 0)

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

    # no direct forward/reverse reaction pairs
    for e_id1, (tails1, heads1) in hyperedges.items():
        for e_id2, (tails2, heads2) in hyperedges.items():
            if e_id1 < e_id2 and (heads1, tails1) == (tails2, heads2):
                model.addConstr(
                    b[e_id1] + b[e_id2] <= 1,
                    name=f"no_reverse_pair_{e_id1}_{e_id2}",
                )

    # NMR scores stay attached directly to candidate molecules
    nmr_score = quicksum(
        vl1[v] * n1[v] + vl2[v] * n2[v] + vl3[v] * n3[v]
        for v in candidate_vertices
    )

    internal_edges = [
        e_id
        for e_id, (tails, heads) in hyperedges.items()
        if tails != [] and heads != []
    ]

    number_of_reactions = quicksum(
        b[e_id] for e_id in internal_edges
    )

    total_internal_flow = quicksum(
        x[e_id] for e_id in internal_edges
    )

    return (
        model,
        x,
        b,
        n1,
        n2,
        n3,
        nmr_score,
        number_of_reactions,
        total_internal_flow,
    )


def optimize_model(
    model,
    nmr_score,
    number_of_reactions,
    total_internal_flow,
):
    # 1. maximize NMR evidence
    model.setObjective(nmr_score, GRB.MAXIMIZE)
    model.optimize()

    if model.status != GRB.Status.OPTIMAL:
        return False

    best_nmr = model.ObjVal
    model.addConstr(
        nmr_score >= best_nmr - 1e-8,
        name="keep_best_nmr_score",
    )

    # 2. among equally good NMR solutions, use as few reactions as possible
    model.setObjective(number_of_reactions, GRB.MINIMIZE)
    model.optimize()

    if model.status != GRB.Status.OPTIMAL:
        return False

    best_number_of_reactions = int(round(model.ObjVal))
    model.addConstr(
        number_of_reactions <= best_number_of_reactions,
        name="keep_minimum_number_of_reactions",
    )

    # 3. then minimize total internal flow
    model.setObjective(total_internal_flow, GRB.MINIMIZE)
    model.optimize()

    return model.status == GRB.Status.OPTIMAL


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


def print_solution(
    title,
    flow_solution,
    binary_solution,
    nmr1_solution,
    nmr2_solution,
    nmr3_solution,
    vertex1,
    vertex2,
    vertex3,
    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}, "
            f"Tails = {tails}, Heads = {heads}"
        )

    print("\nNMR assignments:")
    for label, solution, scores in [
        ("NMR1", nmr1_solution, vertex1),
        ("NMR2", nmr2_solution, vertex2),
        ("NMR3", nmr3_solution, vertex3),
    ]:
        for molecule in solution:
            print(
                f"{label}: {molecule}, "
                f"similarity = {scores[molecule]:.4f}"
            )

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


def solve_once(
    name,
    hypergraph,
    vertices,
    vertex1,
    vertex2,
    vertex3,
    excluded_support=None,
):
    (
        model,
        x,
        b,
        n1,
        n2,
        n3,
        nmr_score,
        number_of_reactions,
        total_internal_flow,
    ) = build_model(
        name,
        hypergraph,
        vertices,
        vertex1,
        vertex2,
        vertex3,
        excluded_support=excluded_support,
    )

    if not optimize_model(
        model,
        nmr_score,
        number_of_reactions,
        total_internal_flow,
    ):
        return None

    return (
        model,
        x,
        b,
        n1,
        n2,
        n3,
    )


def main():
    VERTICES = [
        'Butadien',
        'p_{0,0}',
        'p_{0,1}',
        'p_{0,2}',
        'p_{0,3}',
        'p_{0,4}',
        'p_{0,5}',
        'p_{0,6}',
        'p_{0,7}',
        'p_{0,8}',
        'p_{0,9}',
        'p_{0,10}',
        'p_{0,11}',
        'p_{0,12}',
        'p_{0,13}',
        'p_{0,14}',
        'p_{0,15}',
        'p_{0,16}',
        'p_{0,17}',
        'p_{0,18}',
    ]

    # same molecule similarities as in the current butadiensynthesis.py
    # the difference is that they are used directly as molecule scores

    # comparison with the NMR of ethylene and hexatriene
    NMR1 = [
        0.32, 0.5, 0.87, 0.0, 0.11, 0.58, 0.06, 0.0, 0.0, 0.0,
        0.0, 0.0, 0.0, 0.0, 0.09, 0.0, 0.0, 0.17, 0.0, 0.09,
    ]

    # comparison with the NMR of ethylene and octatriene
    NMR2 = [
        0.33, 0.43, 0.6, 0.0, 0.13, 0.9, 0.07, 0.01, 0.0, 0.0,
        0.01, 0.0, 0.0, 0.01, 0.1, 0.0, 0.0, 0.22, 0.0, 0.13,
    ]

    # comparison with the NMR of ethylene and benzene
    NMR3 = [
        0.0, 0.32, 0.0, 0.0, 0.0, 0.0, 0.05, 0.0, 0.0, 0.02,
        0.0, 0.29, 0.95, 0.17, 0.0, 0.0, 0.0, 0.0, 0.0, 0.05,
    ]

    normalize = False

    if normalize:
        NMR1 = [value / sum(NMR1) for value in NMR1]
        NMR2 = [value / sum(NMR2) for value in NMR2]
        NMR3 = [value / sum(NMR3) for value in NMR3]

    VERTICE1 = dict(zip(VERTICES, NMR1))
    VERTICE2 = dict(zip(VERTICES, NMR2))
    VERTICE3 = dict(zip(VERTICES, NMR3))

    result = solve_once(
        "HypergraphFlowSpeciesNMR",
        HYPERGRAPH,
        VERTICES,
        VERTICE1,
        VERTICE2,
        VERTICE3,
    )

    if result is None:
        print("No optimal solution found.")
        return

    model, x, b, n1, n2, n3 = result

    optimal_solution = positive_entries(x)
    optimal_binary_solution = positive_entries(b)
    optimal_nmr1 = positive_entries(n1)
    optimal_nmr2 = positive_entries(n2)
    optimal_nmr3 = positive_entries(n3)

    print_solution(
        "Optimal Solution",
        optimal_solution,
        optimal_binary_solution,
        optimal_nmr1,
        optimal_nmr2,
        optimal_nmr3,
        VERTICE1,
        VERTICE2,
        VERTICE3,
        HYPERGRAPH,
    )

    # optional second-best network support
    excluded_support = list(optimal_binary_solution.keys())

    result2 = solve_once(
        "SecondBestHypergraphFlowSpeciesNMR",
        HYPERGRAPH,
        VERTICES,
        VERTICE1,
        VERTICE2,
        VERTICE3,
        excluded_support=excluded_support,
    )

    if result2 is None:
        print("\nNo second solution found.")
        return

    model2, x2, b2, n12, n22, n32 = result2

    second_solution = positive_entries(x2)
    second_binary_solution = positive_entries(b2)
    second_nmr1 = positive_entries(n12)
    second_nmr2 = positive_entries(n22)
    second_nmr3 = positive_entries(n32)

    print_solution(
        "Second Best Solution",
        second_solution,
        second_binary_solution,
        second_nmr1,
        second_nmr2,
        second_nmr3,
        VERTICE1,
        VERTICE2,
        VERTICE3,
        HYPERGRAPH,
    )


if __name__ == "__main__":
    main()