import networkx as nx

# Create an agent class for nodes
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 a base edge agent class
class EdgeAgent:
    def __init__(self):
        self.strength = 0.5  # Default strength
        self.active = True

    def __repr__(self):
        return f"{self.__class__.__name__}(strength={self.strength})"

# Subclass for friendship edges
class FriendshipEdge(EdgeAgent):
    def __init__(self):
        super().__init__()
        self.relationship_type = 'friendship'
        self.provides_mental_health_support = False
        self.trust_level = 0.5

    def __repr__(self):
        return f"FriendshipEdge(strength={self.strength}, trust={self.trust_level})"

# Subclass for work edges
class WorkEdge(EdgeAgent):
    def __init__(self):
        super().__init__()
        self.relationship_type = 'work'
        self.professionalism = 0.5
        self.contact_frequency = 'weekly'

    def __repr__(self):
        return f"WorkEdge(strength={self.strength}, professionalism={self.professionalism})"

# Subclass for drug exposure edges
class DrugExposureEdge(EdgeAgent):
    def __init__(self):
        super().__init__()
        self.relationship_type = 'drug_exposure'
        self.risk_level = 'low'

    def __repr__(self):
        return f"DrugExposureEdge(strength={self.strength}, risk={self.risk_level})"

# Subclass for household edges
class HouseholdEdge(EdgeAgent):
    def __init__(self):
        super().__init__()
        self.relationship_type = 'household'
        self.contact_frequency = 'daily'
        self.shared_resources = True

    def __repr__(self):
        return f"HouseholdEdge(strength={self.strength}, contact={self.contact_frequency})"

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

# Create edge agents using subclasses
friendship_edge = FriendshipEdge()
friendship_edge.strength = 0.8
friendship_edge.provides_mental_health_support = True
friendship_edge.trust_level = 0.9

drug_exposure_edge = DrugExposureEdge()
drug_exposure_edge.strength = 0.3
drug_exposure_edge.risk_level = 'moderate'

work_edge = WorkEdge()
work_edge.strength = 0.5
work_edge.professionalism = 0.8

household_edge = HouseholdEdge()
household_edge.strength = 1.0
household_edge.contact_frequency = 'daily'

# Create a MultiGraph
G = nx.MultiGraph()
G.add_node(agent1)
G.add_nodes_from([agent2, agent3])

# Add edges with edge-agent objects
G.add_edge(agent1, agent2, edge_agent=friendship_edge)
G.add_edge(agent1, agent2, edge_agent=drug_exposure_edge)
G.add_edge(agent1, agent3, edge_agent=work_edge)
G.add_edge(agent2, agent3, edge_agent=household_edge)

print("All relationships:")
for agent_a, agent_b, key, data in G.edges(keys=True, data=True):
    edge = data['edge_agent']
    print(f"{agent_a.name} <-> {agent_b.name}: {edge}")
    print(f"  Type: {edge.relationship_type}")
    print(f"  Strength: {edge.strength}")
    if hasattr(edge, 'risk_level'):
        print(f"  Risk level: {edge.risk_level}")
    if hasattr(edge, 'contact_frequency'):
        print(f"  Contact frequency: {edge.contact_frequency}")
    print()

print(f"Alice's relationships with Bob:")
for key, data in G[agent1][agent2].items():
    edge = data['edge_agent']
    print(f"  - {edge.relationship_type} (strength: {edge.strength})")

print(f"\nBob and Charlie's household relationship:")
for key, data in G[agent2][agent3].items():
    edge = data['edge_agent']
    if edge.relationship_type == 'household':
        print(f"  Contact frequency: {edge.contact_frequency}")
        print(f"  Strength: {edge.strength}")
