
"""PKTron Biology — Population Evolution (Genetic Algorithm)."""
import numpy as np
from .genetics import crossover, point_mutation, random_genome


def evolve_population(population, fitness_fn, n_generations=50,
                       mutation_rate=0.01, elitism=2, seed=0):
    """
    Evolve a population of string genomes using a genetic algorithm.

    Parameters
    ----------
    population   : list[str]   initial population of genomes
    fitness_fn   : callable    genome -> float (higher = better)
    n_generations: int
    mutation_rate: float       per-base mutation probability
    elitism      : int         top N individuals carried to next gen
    seed         : int

    Returns
    -------
    dict:
        final_population : list[str]
        best_genome      : str
        best_fitness     : float
        history          : list[float]   best fitness per generation
    """
    rng  = np.random.default_rng(seed)
    pop  = list(population)
    n    = len(pop)
    hist = []

    for gen in range(n_generations):
        # Evaluate
        scores = np.array([fitness_fn(g) for g in pop], dtype=np.float32)
        hist.append(float(scores.max()))

        # Elitism: keep top individuals
        elite_idx = np.argsort(scores)[::-1][:elitism]
        new_pop   = [pop[i] for i in elite_idx]

        # Selection + crossover + mutation
        # Tournament selection
        while len(new_pop) < n:
            t1 = rng.integers(0,n,size=3)
            t2 = rng.integers(0,n,size=3)
            p1 = pop[t1[np.argmax(scores[t1])]]
            p2 = pop[t2[np.argmax(scores[t2])]]
            c1, c2 = crossover(p1, p2, seed=int(rng.integers(0,1e6)))
            c1 = point_mutation(c1, mutation_rate, seed=int(rng.integers(0,1e6)))
            c2 = point_mutation(c2, mutation_rate, seed=int(rng.integers(0,1e6)))
            new_pop.extend([c1, c2])

        pop = new_pop[:n]

    scores    = np.array([fitness_fn(g) for g in pop], dtype=np.float32)
    best_idx  = int(np.argmax(scores))
    return {"final_population": pop,
            "best_genome"     : pop[best_idx],
            "best_fitness"    : float(scores[best_idx]),
            "history"         : hist}


def population_diversity(population):
    """
    Compute average pairwise Hamming distance normalised by length.
    """
    n   = len(population)
    L   = len(population[0])
    tot = 0
    cnt = 0
    for i in range(n):
        for j in range(i+1,n):
            tot += sum(a!=b for a,b in zip(population[i],population[j]))
            cnt += 1
    return float(tot/(cnt*L)) if cnt else 0.0


def hardy_weinberg(p, q=None):
    """
    Hardy-Weinberg equilibrium frequencies.
    p = allele A frequency, q = allele a frequency (1-p if not given).
    Returns (AA, Aa, aa) genotype frequencies.
    """
    if q is None:
        q = 1.0 - p
    return float(p**2), float(2*p*q), float(q**2)
