"""Pytest tests for Individual and Population classes."""

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 individual import Individual
from layer1_core.managers.population import Population


def test_create_population():
    """Test creating a population with individuals."""
    population = Population(population_id="test_population")
    pop_size : int = 100

    # Create 100 individuals
    for i in range(pop_size):
        age = 20 + (i % 60) #sequential ages between 20 and 79
        individual = Individual(agent_id=f"person_{i}", age=age)
        population.add_agent(individual)

    assert population.size() == pop_size
    assert population.count_agents(lambda p: p.alive) == pop_size
    assert population.count_agents(lambda p: not p.alive) == 0


def test_individual_death():
    """Test that individuals can die."""
    individual = Individual(agent_id="person_1", age=30)

    assert individual.alive is True
    assert individual.get_health_state() == 'Susceptible'

    individual.die()

    assert individual.alive is False
    assert individual.get_health_state() == 'Dead'


def test_kill_population_sample():
    """Test killing a random sample of the population."""
    population = Population(population_id="test_population")

    # Create 100 individuals
    pop_size: int = 100
    for i in range(pop_size):
        individual = Individual(agent_id=f"person_{i}", age=25)
        population.add_agent(individual)

    # Kill half of them
    sample = population.get_agents(fraction=0.5)
    for agent in sample:
        agent.die()

    # Verify counts
    alive_count = population.count_agents(lambda p: p.alive)
    dead_count = population.count_agents(lambda p: not p.alive)

    assert population.size() == pop_size
    assert alive_count == pop_size / 2
    assert dead_count == pop_size / 2


def test_health_state_distribution():
    """Test that health states are correctly reported."""
    population = Population(population_id="test_population")
    pop_size : int = 100

    # Create 100 individuals
    for i in range(pop_size):
        individual = Individual(agent_id=f"person_{i}", age=30)
        population.add_agent(individual)

    # Kill 30 of them
    sample = population.get_agents(fraction=0.3)
    for agent in sample:
        agent.die()

    # Check health state distribution
    susceptible = population.count_agents(lambda p: p.get_health_state() == 'Susceptible')
    infected = population.count_agents(lambda p: p.get_health_state() == 'Infected')
    dead = population.count_agents(lambda p: p.get_health_state() == 'Dead')

    assert susceptible == 70
    assert infected == 0  # No infections yet
    assert dead == 30


def test_get_agents_without_fraction():
    """Test that get_agents() returns all agents when no fraction specified."""
    population = Population(population_id="test_population")
    pop_size: int = 10

    for i in range(pop_size):
        individual = Individual(agent_id=f"person_{i}", age=25)
        population.add_agent(individual)

    all_agents = population.get_agents()

    assert len(all_agents) == pop_size
    assert all_agents == population.get_agents()  # Should be same each time
    assert population.get_agents(fraction=0.1) != population.get_agents(fraction=0.1) # should be different each time


print("hello")


def test_get_agents_with_fraction():
    """Test that get_agents(fraction) returns correct number of agents."""
    population = Population(population_id="test_population")
    pop_size: int = 100

    for i in range(pop_size):
        individual = Individual(agent_id=f"person_{i}", age=25)
        population.add_agent(individual)

    # Get 25% sample
    sample = population.get_agents(fraction=0.25)
    assert len(sample) == pop_size * 0.25

    # Get 75% sample
    sample = population.get_agents(fraction=0.75)
    assert len(sample) == pop_size * 0.75


def test_filter_agents():
    """Test filtering agents by predicate."""
    population = Population(population_id="test_population")
    pop_size: int = 100

    # Create individuals with different ages
    for i in range(pop_size):
        age = 20 + i  # Ages 20-119
        individual = Individual(agent_id=f"person_{i}", age=age)
        population.add_agent(individual)

    # Filter agents over 65
    elderly = population.filter_agents(lambda p: p.age >= 65)

    assert len(elderly) == 55  # Ages 65-119
    assert all(agent.age >= 65 for agent in elderly)