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

_AgentInfo = tuple  # (genes, best_step)


def tournament_select(genes_list: list[np.ndarray], fitnesses: np.ndarray,
                      n_parents: int, k: int = 5) -> list[np.ndarray]:
    n = len(fitnesses)
    selected = []
    for _ in range(n_parents):
        indices = np.random.choice(n, size=k, replace=False)
        winner = indices[np.argmax(fitnesses[indices])]
        selected.append(genes_list[winner].copy())
    return selected


def single_point_crossover(g1: np.ndarray, g2: np.ndarray) -> np.ndarray:
    point = np.random.randint(1, len(g1))
    return np.concatenate([g1[:point], g2[point:]])


def mutate_seq(genes: np.ndarray, rate: float) -> np.ndarray:
    """Replace each gene with a random action with probability rate."""
    mask = np.random.rand(len(genes)) < rate
    noise = np.random.randint(0, 4, size=len(genes), dtype=np.uint8)
    result = genes.copy()
    result[mask] = noise[mask]
    return result


def mutate_two_rate(genes: np.ndarray, pivot: int,
                    low_rate: float = 0.005, high_rate: float = 0.30) -> np.ndarray:
    """Prefix-aware mutation: low rate before pivot (preserve good path),
    high rate after pivot (aggressively explore the stuck region)."""
    result = genes.copy()
    n = len(genes)
    pivot = min(pivot, n)
    # Prefix: preserve
    mask_pre = np.random.rand(pivot) < low_rate
    result[:pivot][mask_pre] = np.random.randint(0, 4, mask_pre.sum(), dtype=np.uint8)
    # Tail: explore
    tail = n - pivot
    if tail > 0:
        mask_tail = np.random.rand(tail) < high_rate
        result[pivot:][mask_tail] = np.random.randint(0, 4, mask_tail.sum(), dtype=np.uint8)
    return result


def next_generation(agents: list[Agent], cfg: MazeConfig) -> list[Genome]:
    fitnesses = np.array([a.fitness for a in agents], dtype=np.float32)
    genes_list = [a.genome.genes for a in agents]

    # Elites: top n_elite survive unchanged
    n_elite = cfg.n_elite
    elite_idx = np.argsort(fitnesses)[::-1][:n_elite]
    new_genomes: list[Genome] = [Genome(genes_list[i].copy()) for i in elite_idx]

    # Parents pool: top elite_frac
    n_pool = max(2, int(len(agents) * cfg.elite_frac))
    pool_idx = np.argsort(fitnesses)[::-1][:n_pool]
    pool = [genes_list[i] for i in pool_idx]
    pool_fits = fitnesses[pool_idx]

    # Build pool with (genes, best_step) pairs
    pool_agents = [agents[i] for i in pool_idx]
    pool_genes = [a.genome.genes for a in pool_agents]
    pool_pivots = [a._best_step for a in pool_agents]

    n_offspring = cfg.population - n_elite
    parent_indices = [
        pool_idx[np.argmax(pool_fits[np.random.choice(len(pool_fits), 5, replace=False)])]
        for _ in range(n_offspring)
    ]

    for idx in parent_indices:
        parent_agent = agents[idx]
        child = mutate_two_rate(
            parent_agent.genome.genes,
            pivot=parent_agent._best_step,
            low_rate=0.003,
            high_rate=0.30,
        )
        new_genomes.append(Genome(child))

    return new_genomes