"""Pytest tests for InfectionWorldBuilder."""

import sys
from pathlib import Path


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

from infection_world import InfectionWorldBuilder, InfectionWorld
from infection import InfectionEvent


def test_create_infection_world_builder():
    """Test creating an InfectionWorldBuilder with parameters."""
    params = {
        'population_size': 100,
        'min_age': 20,
        'max_age': 80,
        'virus_name': 'TestVirus',
        'transmission_rate': 0.3,
        'incubation_period': 2,
        'infectious_period': 5,
        'recovery_period': 7,
        'base_mortality_rate': 0.01,
        'initial_infections': 5
    }

    builder = InfectionWorldBuilder(params)
    assert builder.param_dict == params


def test_build_infection_world():
    """Test building an infection world with default parameters."""
    params = {
        'population_size': 100,
        'virus_name': 'TestVirus',
        'transmission_rate': 0.3,
        'incubation_period': 2,
        'infectious_period': 5,
        'recovery_period': 7,
        'base_mortality_rate': 0.01,
        'initial_infections': 1
    }

    builder = InfectionWorldBuilder(params)
    world = builder.build()

    assert isinstance(world, InfectionWorld)
    assert world.population is not None
    assert world.schedule is not None
    assert world.virus is not None


def test_world_population():
    """Test that the built world has correct population."""
    params = {
        'population_size': 50,
        'min_age': 20,
        'max_age': 80,
        'initial_infections': 0
    }

    builder = InfectionWorldBuilder(params)
    world = builder.build()

    assert world.population.size() == 50
    assert len(world.population.get_agents()) == 50


def test_world_virus():
    """Test that the built world has correct virus parameters."""
    params = {
        'population_size': 10,
        'virus_name': 'COVID-19',
        'transmission_rate': 0.5,
        'incubation_period': 3,
        'infectious_period': 7,
        'recovery_period': 10,
        'base_mortality_rate': 0.02,
        'initial_infections': 0
    }

    builder = InfectionWorldBuilder(params)
    world = builder.build()

    assert world.virus.name == 'COVID-19'
    assert world.virus.transmission_rate == 0.5
    assert world.virus.incubation_period == 3
    assert world.virus.infectious_period == 7
    assert world.virus.recovery_period == 10
    assert world.virus.base_mortality_rate == 0.02


def test_initial_infections():
    """Test that initial infections are created correctly."""
    params = {
        'population_size': 100,
        'virus_name': 'TestVirus',
        'transmission_rate': 0.3,
        'incubation_period': 2,
        'infectious_period': 5,
        'recovery_period': 7,
        'base_mortality_rate': 0.01,
        'initial_infections': 10
    }

    builder = InfectionWorldBuilder(params)
    world = builder.build()

    # Count infected individuals
    infected_count = world.population.count_agents(
        lambda p: p.get_health_state() == 'Infected'
    )

    assert infected_count == 10

    # Verify infections have correct virus
    for agent in world.population.get_agents():
        if agent.has_active_infection():
            assert agent.infections[0].virus == world.virus
            assert agent.infections[0].time_infected == 0


def test_no_initial_infections():
    """Test that world can be created with no initial infections."""
    params = {
        'population_size': 100,
        'initial_infections': 0
    }

    builder = InfectionWorldBuilder(params)
    world = builder.build()

    infected_count = world.population.count_agents(
        lambda p: p.get_health_state() == 'Infected'
    )

    assert infected_count == 0


def test_schedule_initialized():
    """Test that schedule is properly initialized with infection checks."""
    params = {
        'population_size': 50,
        'initial_infections': 5
    }

    builder = InfectionWorldBuilder(params)
    world = builder.build()

    assert world.schedule is not None
    assert world.schedule.get_current_step() == 0
    # Should have infection check events scheduled for each initial infection
    assert len(world.schedule.get_events()) == 5


def test_infection_id_generator():
    """Test that infection ID generator works correctly."""
    params = {
        'population_size': 10,
        'initial_infections': 0
    }

    builder = InfectionWorldBuilder(params)
    world = builder.build()

    id1 = world.next_infection_id()
    id2 = world.next_infection_id()
    id3 = world.next_infection_id()

    assert id1 == "infection_1"
    assert id2 == "infection_2"
    assert id3 == "infection_3"


def test_default_parameters():
    """Test that default parameters are used when not specified."""
    params = {}  # Empty dict

    builder = InfectionWorldBuilder(params)
    world = builder.build()

    # Check defaults
    assert world.population.size() == 100  # default population_size
    assert world.virus.name == 'DefaultVirus'  # default virus_name
    assert world.virus.transmission_rate == 0.3  # default

    # Count initial infections (default is 1)
    infected_count = world.population.count_agents(
        lambda p: p.get_health_state() == 'Infected'
    )
    assert infected_count == 1

def test_random_seed_reproducibility():
    """Test that using the same random seed produces identical populations."""
    params1 = {
        'population_size': 50,
        'min_age': 20,
        'max_age': 80,
        'initial_infections': 0,
        'random_seed': 42
    }

    params2 = {
        'population_size': 50,
        'min_age': 20,
        'max_age': 80,
        'initial_infections': 0,
        'random_seed': 42
    }

    builder1 = InfectionWorldBuilder(params1)
    world1 = builder1.build()

    builder2 = InfectionWorldBuilder(params2)
    world2 = builder2.build()

    # Get ages from both worlds
    ages1 = [agent.age for agent in world1.population.get_agents()]
    ages2 = [agent.age for agent in world2.population.get_agents()]

    # Should be identical with same seed
    assert ages1 == ages2


def test_different_seed_produces_different_results():
    """Test that different random seeds produce different populations."""
    params1 = {
        'population_size': 50,
        'min_age': 20,
        'max_age': 80,
        'initial_infections': 0,
        'random_seed': 42
    }

    params2 = {
        'population_size': 50,
        'min_age': 20,
        'max_age': 80,
        'initial_infections': 0,
        'random_seed': 99
    }

    builder1 = InfectionWorldBuilder(params1)
    world1 = builder1.build()

    builder2 = InfectionWorldBuilder(params2)
    world2 = builder2.build()

    # Get ages from both worlds
    ages1 = [agent.age for agent in world1.population.get_agents()]
    ages2 = [agent.age for agent in world2.population.get_agents()]

    # Should be different with different seeds
    assert ages1 != ages2


def test_network_creation():
    """Test that exposure network is created."""
    params = {
        'population_size': 50,
        'avg_degree': 5,
        'initial_infections': 0
    }

    builder = InfectionWorldBuilder(params)
    world = builder.build()

    assert world.network_manager is not None
    assert len(world.network_manager.get_edges()) > 0


def test_network_with_different_avg_degree():
    """Test that avg_degree affects network size."""
    params1 = {
        'population_size': 100,
        'avg_degree': 2,
        'initial_infections': 0,
        'random_seed': 42
    }

    params2 = {
        'population_size': 100,
        'avg_degree': 10,
        'initial_infections': 0,
        'random_seed': 43
    }

    builder1 = InfectionWorldBuilder(params1)
    world1 = builder1.build()

    builder2 = InfectionWorldBuilder(params2)
    world2 = builder2.build()

    edges1 = len(world1.network_manager.get_edges())
    edges2 = len(world2.network_manager.get_edges())

    # Higher avg_degree should result in more edges
    assert edges2 > edges1


def test_network_neighbors():
    """Test that get_neighbors works with the network."""
    params = {
        'population_size': 20,
        'avg_degree': 4,
        'initial_infections': 0,
        'random_seed': 42
    }

    builder = InfectionWorldBuilder(params)
    world = builder.build()

    # Check that at least some agents have neighbors
    agents = world.population.get_agents()
    agents_with_neighbors = 0

    for agent in agents:
        neighbors = world.network_manager.get_neighbors(agent)
        if len(neighbors) > 0:
            agents_with_neighbors += 1

    # Most agents should have neighbors
    assert agents_with_neighbors > len(agents) * 0.5


def test_network_types():
    """Test different network types can be created."""
    network_types = ['erdos_renyi', 'barabasi_albert', 'watts_strogatz']

    for net_type in network_types:
        params = {
            'population_size': 30,
            'avg_degree': 5,
            'network_type': net_type,
            'initial_infections': 0,
            'random_seed': 42
        }

        builder = InfectionWorldBuilder(params)
        world = builder.build()

        assert world.network_manager is not None
        assert len(world.network_manager.get_edges()) > 0


def test_infection_schedule():
    """Test that infections trigger schedule events as expected."""
    params = {
                'population_size': 300,
                'avg_degree': 5,
                'network_type': 'erdos_renyi',
                'initial_infections': 0,
                'random_seed': 42
            }
    builder = InfectionWorldBuilder(params)
    world = builder.build()
    initial_infected = world.population.get_agents(fraction = 0.1) #grab 10% of pop

    for indiv in initial_infected:
        initial_infection_event = InfectionEvent(indiv, world.network_manager, world.schedule, world.next_infection_id)
        world.schedule.add_event(0, initial_infection_event)

    assert len(world.schedule.get_events()) == 0.1 * 300
    print()
    print("Events in Schedule:")
    print(len(world.schedule.get_events()))

    for i in range(15):
        world.schedule.execute_next_event().execute()

    print()
    print("Events in Schedule:")
    print(len(world.schedule.get_events()))

    print()
    print("Infections:")
    print(len(world.population.filter_agents(lambda p: p.get_health_state() == 'Infected')))