import numpy as np
from agent import Agent, NeuralNet
from config import MazeConfig


def tournament_select(weights: np.ndarray, fitnesses: np.ndarray, n_parents: int, k: int = 5) -> np.ndarray:
    n = len(fitnesses)
    selected = np.empty((n_parents, weights.shape[1]), dtype=np.float32)
    for i in range(n_parents):
        indices = np.random.choice(n, size=k, replace=False)
        winner = indices[np.argmax(fitnesses[indices])]
        selected[i] = weights[winner]
    return selected


def uniform_crossover(p1: np.ndarray, p2: np.ndarray) -> np.ndarray:
    mask = np.random.rand(len(p1)) < 0.5
    child = np.where(mask, p1, p2)
    return child.astype(np.float32)


def mutate(weights: np.ndarray, std: float = 0.05, rate: float = 0.20) -> np.ndarray:
    mask = np.random.rand(len(weights)) < rate
    noise = np.random.randn(len(weights)).astype(np.float32) * std
    return (weights + mask * noise).astype(np.float32)


def next_generation(agents: list[Agent], cfg: MazeConfig) -> list[NeuralNet]:
    fitnesses = np.array([a.fitness for a in agents], dtype=np.float32)
    all_weights = np.stack([a.net.weights for a in agents], axis=0)

    n_elite = 10
    elite_indices = np.argsort(fitnesses)[::-1][:n_elite]
    elite_weights = all_weights[elite_indices]

    n_pool = max(1, int(len(agents) * cfg.elite_frac))
    pool_indices = np.argsort(fitnesses)[::-1][:n_pool]
    pool_weights = all_weights[pool_indices]
    pool_fitnesses = fitnesses[pool_indices]

    n_offspring = cfg.population - n_elite
    parents = tournament_select(pool_weights, pool_fitnesses, n_parents=n_offspring * 2)

    new_nets: list[NeuralNet] = []

    for w in elite_weights:
        new_nets.append(NeuralNet(w.copy()))

    for i in range(n_offspring):
        p1 = parents[i * 2]
        p2 = parents[i * 2 + 1]
        child = uniform_crossover(p1, p2)
        child = mutate(child, std=cfg.mutation_std, rate=cfg.mutation_rate)
        new_nets.append(NeuralNet(child))

    return new_nets