"""Evolutionary algorithms for evolving axioms.

Provides a small GeneticAlgorithm class that mutates and recombines axioms (strings)
to produce candidate new axioms. Pure-Python fallback implementation so the
project can run without external evolutionary frameworks.
"""
import random
from typing import List, Callable


class GeneticAlgorithm:
    def __init__(self, population: List[str], fitness_fn: Callable[[str], float],
                 mutate_rate: float = 0.2, crossover_rate: float = 0.5):
        self.population = list(population)
        self.fitness_fn = fitness_fn
        self.mutate_rate = mutate_rate
        self.crossover_rate = crossover_rate

    def _mutate(self, axiom: str) -> str:
        if random.random() > self.mutate_rate:
            return axiom
        # simple mutation: insert/delete/replace a character or word
        ops = [self._insert_char, self._delete_char, self._replace_char, self._swap_words]
        op = random.choice(ops)
        return op(axiom)

    def _insert_char(self, s: str) -> str:
        pos = random.randint(0, len(s))
        ch = random.choice('abcdefghijklmnopqrstuvwxyz ')
        return s[:pos] + ch + s[pos:]

    def _delete_char(self, s: str) -> str:
        if not s:
            return s
        pos = random.randint(0, len(s)-1)
        return s[:pos] + s[pos+1:]

    def _replace_char(self, s: str) -> str:
        if not s:
            return s
        pos = random.randint(0, len(s)-1)
        ch = random.choice('abcdefghijklmnopqrstuvwxyz ')
        return s[:pos] + ch + s[pos+1:]

    def _swap_words(self, s: str) -> str:
        parts = s.split()
        if len(parts) < 2:
            return s
        i, j = random.sample(range(len(parts)), 2)
        parts[i], parts[j] = parts[j], parts[i]
        return ' '.join(parts)

    def _crossover(self, a: str, b: str) -> str:
        # single-point crossover at word boundary
        wa = a.split()
        wb = b.split()
        if not wa or not wb:
            return a
        pa = random.randint(0, len(wa))
        pb = random.randint(0, len(wb))
        child = ' '.join(wa[:pa] + wb[pb:])
        return child

    def step(self, k: int = 10) -> List[str]:
        """Run one generation, return top-k candidates."""
        scored = [(p, self.fitness_fn(p)) for p in self.population]
        scored.sort(key=lambda x: x[1], reverse=True)
        survivors = [p for p, _ in scored[: max(2, len(scored)//2)]]

        # produce new population
        new_pop = list(survivors)
        while len(new_pop) < len(self.population):
            if random.random() < self.crossover_rate:
                a, b = random.sample(survivors, 2)
                child = self._crossover(a, b)
            else:
                child = random.choice(survivors)
            child = self._mutate(child)
            new_pop.append(child)

        self.population = new_pop
        result = [(p, self.fitness_fn(p)) for p in self.population]
        result.sort(key=lambda x: x[1], reverse=True)
        return [p for p, _ in result[:k]]