AI-Ready Data: The Foundation for Intelligent Systems

from rdflib import Graph, URIRef, Literal, Namespace
import networkx as nx
from pyvis.network import Network


# ---------------------------
# 1) Build RDF Graph
# ---------------------------
EX = Namespace("http://example.org/faa/")
g = Graph()
g.bind("faa", EX)


# Entities
flights = {
    "UA245": ("LAX", "JFK", "N123UA", "UnitedAirlines"),
    "UA678": ("SFO", "ORD", "N456UA", "UnitedAirlines"),
    "AA101": ("DFW", "MIA", "N789AA", "AmericanAirlines"),
}

aircrafts = {
    "N123UA": "Aircraft",
    "N456UA": "Aircraft",
    "N789AA": "Aircraft"
}

airports = {
    "LAX": "Los Angeles",
    "JFK": "New York JFK",
    "SFO": "San Francisco",
    "ORD": "Chicago O'Hare",
    "DFW": "Dallas-Fort Worth",
    "MIA": "Miami"
}

airlines = ["UnitedAirlines", "AmericanAirlines"]
passengers = {"Alice": "UA245", "Bob": "AA101"}


# Add flights
for flight, (dep, arr, ac, al) in flights.items():
    f_uri = URIRef(EX[flight])
    g.add((f_uri, EX.type, EX.Flight))
    g.add((f_uri, EX.departsFrom, URIRef(EX[dep])))
    g.add((f_uri, EX.arrivesAt, URIRef(EX[arr])))
    g.add((f_uri, EX.operatedBy, URIRef(EX[ac])))
    g.add((f_uri, EX.airline, URIRef(EX[al])))


# Add aircrafts
for ac, typ in aircrafts.items():
    ac_uri = URIRef(EX[ac])
    g.add((ac_uri, EX.type, EX[typ]))
    g.add((ac_uri, EX.registration, Literal(ac)))


# Add airports
for code, name in airports.items():
    ap_uri = URIRef(EX)
    g.add((ap_uri, EX.type, EX.Airport))
    g.add((ap_uri, EX.iataCode, Literal(code)))
    g.add((ap_uri, EX.name, Literal(name)))


# Add airlines
for al in airlines:
    al_uri = URIRef(EX[al])
    g.add((al_uri, EX.type, EX.Airline))
    g.add((al_uri, EX.name, Literal(al)))


# Add passengers and reservations
for person, flight in passengers.items():
    p_uri = URIRef(EX[person])
    g.add((p_uri, EX.type, EX.Passenger))
    g.add((p_uri, EX.name, Literal(person)))
    g.add((URIRef(EX[flight]), EX.hasReservation, p_uri))


# ----------------------------------
# 2) Convert RDF -> Property Graph
# ----------------------------------
lpg = nx.DiGraph()
props = {}


for s, p, o in g:
    s_id = str(s).split("/")[-1]
    p_label = str(p).split("/")[-1]

    if p_label == "type":
        o_label = str(o).split("/")[-1]
        props.setdefault(s_id, {})["label"] = o_label
        lpg.add_node(s_id)
    elif isinstance(o, URIRef):
        o_id = str(o).split("/")[-1]
        lpg.add_edge(s_id, o_id, label=p_label)
    else:
        props.setdefault(s_id, {})[p_label] = str(o)
        if s_id not in lpg:
            lpg.add_node(s_id)


for node, attributes in props.items():
    if node in lpg:
        lpg.nodes[node].update(attributes)


# ----------------------------------
# 3) Visualize with PyVis
# ----------------------------------
net = Network(height="800px", width="100%", directed=True, notebook=False, bgcolor="#ffffff")


type_colors = {
    "Flight": "#3b82f6",
    "Aircraft": "#10b981",
    "Airport": "#f59e0b",
    "Passenger": "#ef4444",
    "Airline": "#a855f7"
}


for node_id, attrs in lpg.nodes(data=True):
    node_type = attrs.get("label", "Thing")
    tooltip = "
".join(f"{k}: {v}" for k, v in attrs.items())
    display_label = f"{node_id}\\n({node_type})"
    net.add_node(
        node_id,
        label=display_label,
        title=tooltip or node_id,
        color=type_colors.get(node_type, "#9ca3af"),
        shape="ellipse"
    )


for u, v, eattrs in lpg.edges(data=True):
    elabel = eattrs.get("label", "")
    net.add_edge(u, v, label=elabel, title=elabel, arrows="to", smooth=True)


# Save to HTML
net.write_html("rdf_lpg_graph.html")
print("✅ Wrote rdf_lpg_graph.html (open it in your browser).")