from typing import List
import random


def ga_mutate(new_pop: List[List[str]],
            mutation_rate: float = .02) -> List[List[str]]:
    """Mutate the population

    Each chromosome is copied and mutated at mutation_rate.
    When a gene mutates, it's randomly assigned to a partiton.
    possibly the same partition.

    Parameters
    ----------
    new_pop: List[List[str]]
        A list of chromosomes.
    mutation_rate: float
        How often a mutation occurs. 0.02 is a good rate for
        my test sets.
    Returns
    -------
    List[List[str]]
        A list of chromosomes. Mutated chromosomes.
    """
    mutated = []
    for solution in new_pop:
        new_solution = list(solution)
        for i, gene in enumerate(new_solution):
            if random.random() < mutation_rate:
                new_solution[i] = ['train', 'valid', 'test'][random.randint(0,2)]
        mutated.append(new_solution)

    return mutated