import random
import copy
import torch


def mutate(model, mutation_rate=0.01):
    new_model = copy.deepcopy(model)
    for param in new_model.parameters():
        if random.random() < mutation_rate:
            noise = torch.randn_like(param) * 0.1
            param.data += noise
    return new_model


def crossover(parent1, parent2):
    child = copy.deepcopy(parent1)
    for p_child, p2 in zip(child.parameters(), parent2.parameters()):
        mask = torch.rand_like(p_child) < 0.5
        p_child.data[mask] = p2.data[mask]
    return child


def evolve(population, fitnesses, retain_ratio=0.2, mutation_rate=0.1):
    # rank by fitness (higher is better)
    paired = sorted(zip(fitnesses, population), key=lambda x: x[0], reverse=True)
    retain_len = int(len(paired) * retain_ratio)
    parents = [ind for _, ind in paired[:retain_len]]
    next_gen = parents.copy()
    while len(next_gen) < len(population):
        p1, p2 = random.sample(parents, 2)
        child = crossover(p1, p2)
        child = mutate(child, mutation_rate)
        next_gen.append(child)
    return next_gen