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 edge agent classes
class EdgeAgent:
    def __init__(self):
        self.strength = 0.5
        self.active = True

class FriendshipEdge(EdgeAgent):
    def __init__(self):
        super().__init__()
        self.trust_level = 0.5
        self.provides_mental_health_support = False

class WorkEdge(EdgeAgent):
    def __init__(self):
        super().__init__()
        self.professionalism = 0.5
        self.contact_frequency = 'weekly'

class HouseholdEdge(EdgeAgent):
    def __init__(self):
        super().__init__()
        self.contact_frequency = 'daily'
        self.shared_resources = True

class DrugExposureEdge(EdgeAgent):
    def __init__(self):
        super().__init__()
        self.risk_level = 'low'

# Create a Simulation class that manages separate networks
class Simulation:
    def __init__(self):
        self.agents = []
        self.friendship_network = nx.Graph()
        self.work_network = nx.Graph()
        self.household_network = nx.Graph()
        self.drug_exposure_network = nx.Graph()

    def add_agent(self, agent):
        """Add an agent to the simulation"""
        self.agents.append(agent)
        # Add agent to all networks
        self.friendship_network.add_node(agent)
        self.work_network.add_node(agent)
        self.household_network.add_node(agent)
        self.drug_exposure_network.add_node(agent)

    def add_friendship(self, agent1, agent2, edge_agent):
        """Add a friendship edge"""
        self.friendship_network.add_edge(agent1, agent2, edge_agent=edge_agent)

    def add_work_connection(self, agent1, agent2, edge_agent):
        """Add a work connection edge"""
        self.work_network.add_edge(agent1, agent2, edge_agent=edge_agent)

    def add_household_connection(self, agent1, agent2, edge_agent):
        """Add a household edge"""
        self.household_network.add_edge(agent1, agent2, edge_agent=edge_agent)

    def add_drug_exposure(self, agent1, agent2, edge_agent):
        """Add a drug exposure edge"""
        self.drug_exposure_network.add_edge(agent1, agent2, edge_agent=edge_agent)

    def get_friendships_of_agent(self, agent):
        """Get all friends of an agent"""
        return list(self.friendship_network.neighbors(agent))

    def get_household_members(self, agent):
        """Get all household members of an agent"""
        return list(self.household_network.neighbors(agent))

    def get_work_contacts(self, agent):
        """Get all work contacts of an agent"""
        return list(self.work_network.neighbors(agent))

    def get_drug_exposure_contacts(self, agent):
        """Get all drug exposure contacts of an agent"""
        return list(self.drug_exposure_network.neighbors(agent))

    def friendship_centrality(self):
        """Calculate eigenvector centrality for friendships"""
        if len(self.friendship_network.edges()) == 0:
            return {}
        return nx.eigenvector_centrality(self.friendship_network)

    def household_centrality(self):
        """Calculate eigenvector centrality for household connections"""
        if len(self.household_network.edges()) == 0:
            return {}
        return nx.eigenvector_centrality(self.household_network)

    def print_networks_summary(self):
        """Print a summary of all networks"""
        print(f"Agents: {len(self.agents)}")
        print(f"Friendship edges: {self.friendship_network.number_of_edges()}")
        print(f"Work edges: {self.work_network.number_of_edges()}")
        print(f"Household edges: {self.household_network.number_of_edges()}")
        print(f"Drug exposure edges: {self.drug_exposure_network.number_of_edges()}")

# Create simulation
sim = Simulation()

# Create agents
alice = Agent("Alice", 30)
bob = Agent("Bob", 25)
charlie = Agent("Charlie", 35)

sim.add_agent(alice)
sim.add_agent(bob)
sim.add_agent(charlie)

# Add relationships
friendship_ab = FriendshipEdge()
friendship_ab.strength = 0.8
friendship_ab.trust_level = 0.9
sim.add_friendship(alice, bob, friendship_ab)

work_ac = WorkEdge()
work_ac.strength = 0.5
work_ac.professionalism = 0.8
sim.add_work_connection(alice, charlie, work_ac)

household_bc = HouseholdEdge()
household_bc.strength = 1.0
sim.add_household_connection(bob, charlie, household_bc)

drug_ab = DrugExposureEdge()
drug_ab.strength = 0.3
drug_ab.risk_level = 'moderate'
sim.add_drug_exposure(alice, bob, drug_ab)

# Demonstrate usage
print("=== Simulation Summary ===")
sim.print_networks_summary()

print("\n=== Alice's Relationships ===")
print(f"Friends: {sim.get_friendships_of_agent(alice)}")
print(f"Work contacts: {sim.get_work_contacts(alice)}")
print(f"Household members: {sim.get_household_members(alice)}")
print(f"Drug exposure contacts: {sim.get_drug_exposure_contacts(alice)}")

print("\n=== Friendship Network Centrality ===")
centrality = sim.friendship_centrality()
for agent, score in centrality.items():
    print(f"{agent.name}: {score:.3f}")

print("\n=== Household Network Centrality ===")
centrality = sim.household_centrality()
for agent, score in centrality.items():
    print(f"{agent.name}: {score:.3f}")

print("\n=== Bob and Charlie's Relationship ===")
for agent_a, agent_b, data in sim.household_network.edges(data=True):
    if (agent_a == bob and agent_b == charlie) or (agent_a == charlie and agent_b == bob):
        edge = data['edge_agent']
        print(f"Household edge: {agent_a.name} <-> {agent_b.name}")
        print(f"  Strength: {edge.strength}")
        print(f"  Contact frequency: {edge.contact_frequency}")
        print(f"  Shared resources: {edge.shared_resources}")