import json
from collections import Counter

include("butadienNetwork.py")

with open("gurobi_flows.json") as f:
    gurobi_flows = json.load(f)


def canonical_smiles(smiles):
    return str(Graph.fromSMILES(smiles).smiles)


def vertex_key(v):
    return str(v.graph.smiles)


vertices_by_smiles = {}
for v in dg.vertices:
    vertices_by_smiles.setdefault(vertex_key(v), []).append(v)


def current_vertex(smiles):
    matches = vertices_by_smiles.get(canonical_smiles(smiles), [])
    if len(matches) != 1:
        raise RuntimeError(
            f"{smiles}: expected one MØD vertex, found {len(matches)}"
        )
    return matches[0]


def current_reaction(educts, products):
    wanted_educts = Counter(canonical_smiles(x) for x in educts)
    wanted_products = Counter(canonical_smiles(x) for x in products)

    matches = []
    for e in dg.edges:
        current_educts = Counter(vertex_key(v) for v in e.sources)
        current_products = Counter(vertex_key(v) for v in e.targets)
        if current_educts == wanted_educts and current_products == wanted_products:
            matches.append(e)

    if len(matches) != 1:
        raise RuntimeError(
            f"{educts} -> {products}: expected one MØD edge, found {len(matches)}"
        )
    return matches[0]


def check_causality(gurobi_flow, rank):
    source_flow = {
        current_vertex(smiles): value
        for smiles, value in gurobi_flow["inputs"].items()
    }
    output_flow = {
        current_vertex(smiles): value
        for smiles, value in gurobi_flow["outputs"].items()
    }

    reaction_flow = {}
    for reaction in gurobi_flow["reactions"]:
        edge = current_reaction(reaction["educts"], reaction["products"])
        reaction_flow[edge] = reaction_flow.get(edge, 0) + reaction["flow"]

    flow = hyperflow.Model(dg)

    for v in source_flow:
        flow.addSource(v.graph)
    for v in output_flow:
        flow.addSink(v.graph)

    for e in dg.edges:
        flow.addConstraint(edgeFlow[e] == reaction_flow.get(e, 0))

    for v in dg.vertices:
        flow.addConstraint(inFlow[v] == source_flow.get(v, 0))
        flow.addConstraint(outFlow[v] == output_flow.get(v, 0))

    flow.findSolutions(maxNumSolutions=1)

    try:
        solution = next(iter(flow.solutions))
    except StopIteration:
        print(f"solution {rank}: not a valid MØD hyperflow")
        return "INVALID"

    query = causality.RealisabilityQuery(flow.dg)
    dag = query.findDAG(solution)

    if not dag:
        print(f"solution {rank}: causally realisable: NO")
        return "NO"

    print(f"solution {rank}: causally realisable: YES")

    printer = DGPrinter()
    printer.graphvizPrefix = 'layout = "dot";'

    def vertex_visible(v):
        return solution.eval(vertexFlow[v]) != 0

    printer.pushVertexVisible(vertex_visible)
    dg.print(printer=printer, data=dag.getPrintData(False))

    return "YES"


results = []
for rank, gurobi_flow in enumerate(gurobi_flows, start=1):
    results.append(check_causality(gurobi_flow, rank))

print("\n=== causality summary ===")
for rank, result in enumerate(results, start=1):
    print(f"solution {rank}: {result}")