"""Exposure network creation utilities for infection simulations."""

import sys
from pathlib import Path
from typing import List, Dict
import networkx as nx

# Add framework src to path for imports
framework_src_path = Path(__file__).parent.parent.parent.parent / "src"
sys.path.insert(0, str(framework_src_path))

from layer1_core.managers.network_manager import NetworkManager
from layer1_core.entities.edge_agent import EdgeAgent


def create_random_network(
    agents: List,
    avg_degree: float,
    network_type: str = "erdos_renyi"
) -> NetworkManager:
    """
    Create a random exposure network among agents.

    Args:
        agents: List of Individual agents
        avg_degree: Average number of contacts per individual
        network_type: Type of random network ("erdos_renyi", "barabasi_albert", "watts_strogatz")

    Returns:
        NetworkManager with edges representing exposure relationships
    """
    n = len(agents)

    # Create NetworkX graph based on type
    if network_type == "erdos_renyi":
        # Probability of edge creation to achieve avg_degree
        p = avg_degree / (n - 1) if n > 1 else 0
        G = nx.erdos_renyi_graph(n, p)

    elif network_type == "barabasi_albert":
        # Number of edges to attach from new node
        m = int(avg_degree / 2)
        m = max(1, m)  # Must be at least 1
        G = nx.barabasi_albert_graph(n, m)

    elif network_type == "watts_strogatz":
        # k must be even and less than n
        k = int(avg_degree)
        if k % 2 != 0:
            k += 1
        k = min(k, n - 1)
        k = max(2, k)  # Must be at least 2
        rewire_prob = 0.1
        G = nx.watts_strogatz_graph(n, k, rewire_prob)

    else:
        raise ValueError(f"Unknown network_type: {network_type}. Use 'erdos_renyi', 'barabasi_albert', or 'watts_strogatz'.")

    # Convert NetworkX graph to NetworkManager with EdgeAgents
    network_manager = NetworkManager()

    for i, j in G.edges():
        edge = EdgeAgent(
            agent_1=agents[i],
            agent_2=agents[j],
            edge_id=f"edge_{i}_{j}"
        )
        network_manager.add_edge(edge)

    return network_manager


def get_network_stats(network_manager: NetworkManager) -> Dict[str, float]:
    """
    Calculate statistics about the exposure network.

    Args:
        network_manager: NetworkManager to analyze

    Returns:
        Dictionary with network statistics (num_edges, avg_degree)
    """
    num_edges = len(network_manager.get_edges())

    # Count degree for each agent
    agent_degrees = {}
    for edge in network_manager.get_edges():
        agent_degrees[edge.agent_1.agent_id] = agent_degrees.get(edge.agent_1.agent_id, 0) + 1
        agent_degrees[edge.agent_2.agent_id] = agent_degrees.get(edge.agent_2.agent_id, 0) + 1

    avg_degree = sum(agent_degrees.values()) / len(agent_degrees) if agent_degrees else 0

    return {
        'num_edges': num_edges,
        'avg_degree': avg_degree,
        'num_nodes': len(agent_degrees)
    }