import wntr
import json

# =========================
# CARICA NET3
# =========================
wn = wntr.network.WaterNetworkModel("Net3.inp")

nodes = []
edges = []

# =========================
# NODI (con coordinate se presenti)
# =========================
for name, node in wn.nodes():

    coords = getattr(node, "coordinates", None)

    if coords is not None:
        x, y = coords
    else:
        x, y = None, None

    nodes.append({
        "_key": name,
        "type": node.node_type,
        "x": x,
        "y": y,
        "elevation": getattr(node, "elevation", None),
        "base_demand": getattr(node, "base_demand", None)
    })

# =========================
# ARCHI (pipes, pumps, valves)
# =========================
for name, link in wn.links():

    edges.append({
        "_key": name,
        "_from": f"nodes/{link.start_node_name}",
        "_to": f"nodes/{link.end_node_name}",
        "type": link.link_type,
        "length": getattr(link, "length", None),
        "diameter": getattr(link, "diameter", None),
        "roughness": getattr(link, "roughness", None)
    })

# =========================
# EXPORT JSON PER ARANGO
# =========================
output = {
    "nodes": nodes,
    "edges": edges
}

with open("net3_arango.json", "w") as f:
    json.dump(output, f, indent=2)

print("✔ Net3 convertito con coordinate → net3_arango.json")