import networkx as nx

# Create an agent class
class Agent:
    def __init__(self, name, age):
        self.name = name
        self.age = age

    def __repr__(self):
        return f"Agent({self.name}, age={self.age})"

# Create agent instances
agent1 = Agent("Alice", 30)
agent2 = Agent("Bob", 25)
agent3 = Agent("Charlie", 35)

# Create a MultiGraph (allows multiple edges between same nodes)
G = nx.MultiGraph()
G.add_node(agent1)
G.add_nodes_from([agent2, agent3])

# Add multiple edge types between the same agents
# Alice and Bob: friendship and drug exposure pathway
G.add_edge(agent1, agent2, relationship_type='friendship', strength=0.8)
G.add_edge(agent1, agent2, relationship_type='drug_exposure', risk=0.3)

# Alice and Charlie: only professional
G.add_edge(agent1, agent3, relationship_type='work', strength=0.5)

# Bob and Charlie: friendship and household
G.add_edge(agent2, agent3, relationship_type='friendship', strength=0.9)
G.add_edge(agent2, agent3, relationship_type='household', strength=1.0)

print("All edges:")
for agent_a, agent_b, key, data in G.edges(keys=True, data=True):
    print(f"{agent_a.name} <-> {agent_b.name} (edge {key}): {data}")

print(f"\nNumber of edges between Alice and Bob: {G.number_of_edges(agent1, agent2)}")

print("\nAlice's relationships:")
for neighbor in G.neighbors(agent1):
    edge_count = G.number_of_edges(agent1, neighbor)
    print(f"  {agent1.name} <-> {neighbor.name}: {edge_count} relationship(s)")

    # Get all edges between them
    for key, data in G[agent1][neighbor].items():
        print(f"    - {data['relationship_type']}")

print("\nBob and Charlie's relationships:")
for key, data in G[agent2][agent3].items():
    print(f"  {data['relationship_type']}: strength={data.get('strength', 'N/A')}")