"""
CS 150 Prelab 4 example
"""

import random

def generate_population(population_size):
    """
    Generate a random population containing two alleles
    
    Args:
        population_size: Size of the population
    
    Returns:
        Population as a string
    """
    population = ""
    for i in range(population_size):
        # choice randomly chooses from a sequence, e.g., from a string
        population += random.choice("aA")
    return population
    

def next_generation(population):
    """
    Generate the next generation by randomly sampling from the
    current population
    
    Args:
        population: Current population as a string
    
    Returns:
        Next generation as a string
    """
    next_population = ''
    for i in range(len(population)):
        next_population += random.choice('aA')
    return next_population

def simulate_population(population):
    """
    Simulate allele fixation in a population
    
    Args:
        population: Initial population as a string
    
    Returns:
        Integer number of generations need to achieve fixation
    """
    print(population)
    generations = 0
    while 'a' in population and 'A' in population:
        population = next_generation(population)
        print(population)
        generations += 1
    return generations