# CS146 Population genomics example
from random import choice

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

def next_gen(pop):
    """
    Generate the next generation by randomly sampling from the
    current population
    
    Args:
        pop: Current population as a string
    
    Returns:
        Next generation as a string
    """
    next_pop = ""
    for i in range(len(pop)):
        next_pop += choice(pop)
    return next_pop

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