"""
Simulation of population of individuals and how they reproduce, etc.

Original author: Prof. Victor Norman, August 7, 2019
Revisions and refactoring by Ken Arnold, Fall 2021
Bugfixes and extensions: YOUR NAME (yourID)
"""

### Imports

import random


### Constants

DEBUG = False
NUM_INITIAL_INDIVIDUALS = 50
HEALTH_VARIATION = 2.0  # plus or minus
INITIAL_AGE_MAX = 6
DIE_AGE = 8
REPRODUCE_AGE = 1
CHANCE_OF_REPRODUCING = 0.35  # 35% chance of reproducing
NUM_ITERATIONS = 50


### Classes

class DayStats:
    def __init__(self):
        self.num_died_of_old_age = 0
        self.population_size = 0
        
    def increment_old_age(self):
        self.num_died_of_old_age += 1
        if DEBUG:
            print("died of old age")
    
    def update_population(self, individuals):
        self.population_size = len(individuals)
        
    def __str__(self):
        return (
            f'Died of old age: {self.num_died_of_old_age}, ' +
            f'Total population: {self.population_size}'
        )


class Individual:
    '''A member of the population.

    Data:
    
    gender: "M" or "F"
    age:  2 (or some other number)
    is_alive: True, or perhaps False
    '''
    def __init__(self):
        """Make a new individual with random values for gender, age, and health."""
        self.gender = random.choice(["M", "F"])
        self.age = random.randint(0, INITIAL_AGE_MAX)
        self.is_alive = True
        
    def do_timestep(self, stats):
        """Advance the simulation by one time step for this individual.

        An individual dies if it reaches the max age.
        """

        self.age += 1
        if self.age > DIE_AGE:
            self.is_alive = False
            stats.increment_old_age()


### Functions

def get_reproducing_pairs(individuals):
    """Go through the population and pair up females and males. Do
    not assign either a female or male when it is already assigned.
    Return a pair of the two lists -- females in the first list and 
    males in the second list. These are indices of the individuals
    in the main list.
    Return two empty lists if no pairs are found.
    """
    female_indices = []
    male_indices = []

    # Look for unique females of sufficient age.
    # Since we only look for a single match for each female,
    # no females will get doubly assigned. But we'll have to be careful about males.
    for female_index in range(len(individuals)):
        # If this wasn't actually a female, keep looking.
        if individuals[female_index].gender != "F":
            continue
        # If they're too young, keep looking.
        if individuals[female_index].age < REPRODUCE_AGE:
            continue

        # Look for a candidate match.
        male_index = find_male(individuals, male_indices)
        if male_index != None:
            female_indices.append(female_index)
            male_indices.append(male_index)
        else:
            # No males are left, abort.
            break
    return female_indices, male_indices


def find_male(individuals, ineligibile_indices):
    """Find the first available male individual.

    Parameters:
    - individuals: the population (a list of dictionaries)
    - ineligible_indices: the indices of males that shouldn't be
        considered because they have already been paired
        
    Returns: the index of an eligible male, or None if none are found.
    """
    for male_index in range(len(individuals)):
        if individuals[male_index].gender != "M":
            continue

        # If they're too young, keep looking.
        if individuals[male_index].age < REPRODUCE_AGE:
            continue

        # If they're already taken, keep looking:
        if male_index in ineligibile_indices:
            continue

        # Otherwise, they're a candidate.
        return male_index

    # Return a special object to indicate that no candidate male was found.
    return None


def reproduce(individuals, stats):
    """
    Get reproducing pairs of females and males and create a new offspring based
    on CHANCE_OF_REPRODUCING.
    """
    female_indices, male_indices = get_reproducing_pairs(individuals)
    if DEBUG:
        print("Females:", female_indices)
        print("Males:  ", male_indices)

    for i in range(len(female_indices)):
        # Someday, perhaps: consider health of the female and male in the pair?

        if random.random() < CHANCE_OF_REPRODUCING:
            # Create an offspring.
            offspring = Individual()
            individuals.append(offspring)


def remove_dead_individuals(individuals):
    """Remove any individuals from the population if they are not alive."""
    result = []
    for x in individuals:
        if x.is_alive:
            result.append(x)
    return result


### Main code

# Initialize first population of individuals
individuals = []
for i in range(NUM_INITIAL_INDIVIDUALS):
    individuals.append(Individual())

# Initialize stats, keep a daily log.
stat_log = []
stats = DayStats()
stats.update_population(individuals)
print(stats)


# Main Loop
for i in range(NUM_ITERATIONS):
    # Start a day
    print("-" * 25, "Day", i, "-" * 25)

    # Update the simulation for each individual.
    for x in individuals:
        x.do_timestep(stats)

    # Clean up dead individuals, stop if none are left.
    individuals = remove_dead_individuals(individuals)
    if len(individuals) == 0:
        break
    
    # Allow individuals to reproduce.
    reproduce(individuals, stats)

    # Collect logs.
    stats.update_population(individuals)
    print(stats)
    stat_log.append(vars(stats))
    stats = DayStats()