NOVA/nova_evo/evolution.py

"""
NOVA-EVO: Genetic algorithm for hyperparameter and architecture search
"""

import random
import json
from pathlib import Path
from typing import List, Tuple, Optional
import time
from tqdm import tqdm
import copy

from .config import EvolutionConfig, Individual
from .fitness import FitnessEvaluator


class EvolutionEngine:
    """
    Genetic algorithm engine for evolving NOVA configurations

    Features:
    - Multi-objective fitness (loss, latency, memory, quality)
    - Elitism with Pareto selection
    - Mutation and crossover
    - Hall of Fame for best individuals
    - Rollback on regression
    """

    def __init__(
        self,
        config: EvolutionConfig,
        fitness_evaluator: FitnessEvaluator,
    ):
        """
        Args:
            config: Evolution configuration
            fitness_evaluator: Fitness evaluation engine
        """
        self.config = config
        self.evaluator = fitness_evaluator

        # Population
        self.population: List[Individual] = []
        self.generation = 0

        # Hall of Fame - best individuals
        self.hall_of_fame: List[Individual] = []
        self.max_hof_size = 10

        # Tracking
        self.evolution_history = []
        self.start_time = None

        # Setup
        Path(config.save_dir).mkdir(parents=True, exist_ok=True)
        random.seed(config.seed)

    def initialize_population(self) -> List[Individual]:
        """Create initial random population"""
        print(f"Initializing population of {self.config.population_size}...")

        population = []

        for i in range(self.config.population_size):
            individual = Individual(
                learning_rate=random.uniform(self.config.lr_min, self.config.lr_max) if self.config.search_learning_rate else 3e-4,
                batch_size=random.choice(self.config.batch_size_options) if self.config.search_batch_size else 8,
                warmup_steps=random.randint(self.config.warmup_min, self.config.warmup_max) if self.config.search_warmup_steps else 1000,
                weight_decay=random.uniform(self.config.wd_min, self.config.wd_max) if self.config.search_weight_decay else 0.1,
                rope_theta=random.choice(self.config.rope_theta_options) if self.config.search_rope_theta else 10000.0,
                hidden_act=random.choice(self.config.activation_options) if self.config.search_activation else "swiglu",
                norm_type=random.choice(self.config.norm_options) if self.config.search_norm else "rmsnorm",
                generation=0,
            )
            population.append(individual)

        return population

    def evaluate_population(self, population: List[Individual]) -> List[Individual]:
        """Evaluate fitness for all individuals in population"""
        print(f"\nEvaluating {len(population)} individuals...")

        for idx, individual in enumerate(tqdm(population, desc="Evaluating")):
            # Skip if already evaluated
            if individual.fitness is not None:
                continue

            # Evaluate
            metrics = self.evaluator.evaluate(individual)

            # Store metrics
            individual.loss = metrics['loss']
            individual.perplexity = metrics.get('perplexity')
            individual.latency_ms = metrics.get('latency_ms')
            individual.memory_mb = metrics.get('memory_mb')
            individual.quality_score = metrics.get('quality_score', 0.0)

            # Calculate multi-objective fitness
            individual.fitness = self._calculate_fitness(individual)

        return population

    def _calculate_fitness(self, individual: Individual) -> float:
        """
        Calculate multi-objective fitness score

        Lower is better (we're minimizing)
        """
        fitness = 0.0

        # Loss component (lower is better)
        if individual.loss is not None:
            fitness += individual.loss * self.config.loss_weight

        # Latency component (lower is better, normalized)
        if individual.latency_ms is not None:
            normalized_latency = individual.latency_ms / 1000.0  # Normalize to seconds
            fitness += normalized_latency * self.config.latency_weight

        # Memory component (lower is better, normalized)
        if individual.memory_mb is not None:
            normalized_memory = individual.memory_mb / 1000.0  # Normalize to GB
            fitness += normalized_memory * self.config.memory_weight

        # Quality component (higher is better, so negate)
        if individual.quality_score is not None:
            fitness -= individual.quality_score * self.config.quality_weight

        return fitness

    def select_parents(self, population: List[Individual]) -> List[Individual]:
        """
        Select parents for next generation using elitism

        Args:
            population: Current population (should be evaluated)

        Returns:
            Elite individuals to keep
        """
        # Sort by fitness (lower is better)
        sorted_pop = sorted(population, key=lambda x: x.fitness if x.fitness is not None else float('inf'))

        # Select top performers
        num_elite = max(1, int(len(population) * self.config.elite_ratio))
        elite = sorted_pop[:num_elite]

        return elite

    def crossover(self, parent1: Individual, parent2: Individual) -> Individual:
        """
        Create offspring by combining two parents

        Uses uniform crossover - randomly picks from each parent
        """
        child = Individual(
            learning_rate=random.choice([parent1.learning_rate, parent2.learning_rate]),
            batch_size=random.choice([parent1.batch_size, parent2.batch_size]),
            warmup_steps=random.choice([parent1.warmup_steps, parent2.warmup_steps]),
            weight_decay=random.choice([parent1.weight_decay, parent2.weight_decay]),
            rope_theta=random.choice([parent1.rope_theta, parent2.rope_theta]),
            hidden_act=random.choice([parent1.hidden_act, parent2.hidden_act]),
            norm_type=random.choice([parent1.norm_type, parent2.norm_type]),
            generation=self.generation + 1,
            parent_ids=[id(parent1), id(parent2)],
        )

        return child

    def mutate(self, individual: Individual) -> Individual:
        """
        Mutate an individual with random changes

        Args:
            individual: Individual to mutate

        Returns:
            Mutated copy
        """
        mutated = copy.deepcopy(individual)
        mutated.generation = self.generation + 1

        # Mutate each gene with some probability
        if random.random() < self.config.mutation_rate:
            mutated.learning_rate = random.uniform(self.config.lr_min, self.config.lr_max)

        if random.random() < self.config.mutation_rate:
            mutated.batch_size = random.choice(self.config.batch_size_options)

        if random.random() < self.config.mutation_rate:
            mutated.warmup_steps = random.randint(self.config.warmup_min, self.config.warmup_max)

        if random.random() < self.config.mutation_rate:
            mutated.weight_decay = random.uniform(self.config.wd_min, self.config.wd_max)

        if random.random() < self.config.mutation_rate:
            mutated.rope_theta = random.choice(self.config.rope_theta_options)

        if random.random() < self.config.mutation_rate:
            mutated.hidden_act = random.choice(self.config.activation_options)

        if random.random() < self.config.mutation_rate:
            mutated.norm_type = random.choice(self.config.norm_options)

        # Reset fitness (needs re-evaluation)
        mutated.fitness = None
        mutated.loss = None

        return mutated

    def create_next_generation(self, parents: List[Individual]) -> List[Individual]:
        """Create next generation from parents"""
        next_gen = []

        # Keep elite unchanged
        next_gen.extend(copy.deepcopy(parents))

        # Fill rest with offspring
        while len(next_gen) < self.config.population_size:
            # Select two random parents
            parent1, parent2 = random.sample(parents, 2)

            # Crossover
            child = self.crossover(parent1, parent2)

            # Mutate
            child = self.mutate(child)

            next_gen.append(child)

        return next_gen

    def update_hall_of_fame(self, population: List[Individual]):
        """Update hall of fame with best individuals"""
        # Add current best to hall of fame
        for ind in population:
            if ind.fitness is not None:
                self.hall_of_fame.append(copy.deepcopy(ind))

        # Sort by fitness
        self.hall_of_fame.sort(key=lambda x: x.fitness if x.fitness is not None else float('inf'))

        # Keep only top N
        self.hall_of_fame = self.hall_of_fame[:self.max_hof_size]

    def save_checkpoint(self):
        """Save evolution state"""
        checkpoint_path = Path(self.config.save_dir) / f"generation_{self.generation}.json"

        checkpoint = {
            'generation': self.generation,
            'population': [ind.to_dict() for ind in self.population],
            'hall_of_fame': [ind.to_dict() for ind in self.hall_of_fame],
            'config': self.config.__dict__,
        }

        with open(checkpoint_path, 'w') as f:
            json.dump(checkpoint, f, indent=2)

        print(f"  Checkpoint saved: {checkpoint_path}")

    def run(self):
        """Run the evolution process"""
        print("=" * 60)
        print("NOVA-EVO: Genetic Algorithm Evolution")
        print("=" * 60)

        self.start_time = time.time()

        # Initialize population
        self.population = self.initialize_population()

        # Evolution loop
        for gen in range(self.config.num_generations):
            self.generation = gen
            print(f"\n{'='*60}")
            print(f"Generation {gen + 1}/{self.config.num_generations}")
            print(f"{'='*60}")

            # Evaluate
            self.population = self.evaluate_population(self.population)

            # Select parents
            parents = self.select_parents(self.population)

            # Update hall of fame
            self.update_hall_of_fame(self.population)

            # Report best individual
            best = self.hall_of_fame[0] if self.hall_of_fame else None
            if best:
                print(f"\n🏆 Best individual so far:")
                print(f"  Fitness: {best.fitness:.4f}")
                print(f"  Loss: {best.loss:.4f}")
                print(f"  LR: {best.learning_rate:.2e}, BS: {best.batch_size}")
                print(f"  Activation: {best.hidden_act}, Norm: {best.norm_type}")

            # Checkpoint
            if (gen + 1) % self.config.checkpoint_every_n_generations == 0:
                self.save_checkpoint()

            # Create next generation
            if gen < self.config.num_generations - 1:
                self.population = self.create_next_generation(parents)

        # Final checkpoint
        self.save_checkpoint()

        print("\n" + "=" * 60)
        print("Evolution Complete!")
        print("=" * 60)
        print(f"Total time: {(time.time() - self.start_time) / 3600:.2f} hours")
        print(f"\nTop 3 individuals:")
        for i, ind in enumerate(self.hall_of_fame[:3]):
            print(f"\n{i+1}. Fitness: {ind.fitness:.4f}")
            print(f"   Loss: {ind.loss:.4f}, LR: {ind.learning_rate:.2e}")
            print(f"   Batch size: {ind.batch_size}, Warmup: {ind.warmup_steps}")
            print(f"   Activation: {ind.hidden_act}, Norm: {ind.norm_type}")