"""Population class - a simple extension of Collective for managing populations of agents."""

import random
from typing import List, Callable, Any, Optional
from .collective import Collective


class Population(Collective):
    """
    Population class for managing a collection of agents representing individuals.

    This is a simple subclass of Collective that provides convenience methods
    for common population-level operations.
    """

    def __init__(self, population_id: str):
        """
        Initialize a Population.

        Args:
            population_id: Unique identifier for this population
        """
        super().__init__(collective_id=population_id)

    def get_agents(self, fraction: Optional[float] = None) -> List[Any]:
        """
        Get all agents in the population, or a random fraction of them.

        Args:
            fraction: Optional fraction of agents to return (0.0 to 1.0).
                     If None, returns all agents. If specified, returns a
                     random sample of that fraction.

        Returns:
            List of agents (all or sampled)
        """
        if fraction is None:
            return self.agents

        if not 0.0 <= fraction <= 1.0:
            raise ValueError(f"fraction must be between 0.0 and 1.0, got {fraction}")

        sample_size = int(len(self.agents) * fraction)
        return random.sample(self.agents, sample_size)

    def size(self) -> int:
        """
        Get the size of the population.

        Returns:
            Number of agents in the population
        """
        return len(self.agents)

    def filter_agents(self, predicate: Callable[[Any], bool]) -> List[Any]:
        """
        Filter agents based on a predicate function.

        Args:
            predicate: Function that takes an agent and returns True/False

        Returns:
            List of agents that satisfy the predicate
        """
        return [agent for agent in self.agents if predicate(agent)]

    def count_agents(self, predicate: Callable[[Any], bool]) -> int:
        """
        Count agents that satisfy a predicate function.

        Args:
            predicate: Function that takes an agent and returns True/False

        Returns:
            Number of agents that satisfy the predicate
        """
        return sum(1 for agent in self.agents if predicate(agent))

    def __repr__(self):
        return f"Population({self.collective_id}, size={self.size()})"