archive.py

"""
Archive management for Bio-ACDC using Dominated Novelty Search (DNS).
Inspired by MAP-Elites and Quality Diversity algorithms.
"""

import json
import logging
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass, field, asdict
import numpy as np
from pathlib import Path

logger = logging.getLogger(__name__)


@dataclass
class BioSolution:
    """Represents a solution (biological LM) in the archive."""
    
    model_path: str
    fitness: float
    skill_vector: Dict[str, float] = field(default_factory=dict)
    
    # Metadata
    generation: int = 0
    parent_paths: Optional[List[str]] = field(default=None, repr=False)
    is_gibberish: bool = False
    
    # Novelty
    novelty_score: Optional[float] = None
    
    def to_dict(self) -> dict:
        return {
            "model_path": self.model_path,
            "fitness": float(self.fitness),
            "skill_vector": {k: float(v) for k, v in self.skill_vector.items()},
            "generation": self.generation,
            "parent_paths": self.parent_paths,
            "is_gibberish": self.is_gibberish,
            "novelty_score": float(self.novelty_score) if self.novelty_score is not None else None,
        }
    
    @classmethod
    def from_dict(cls, data: dict) -> "BioSolution":
        return cls(
            model_path=data["model_path"],
            fitness=data["fitness"],
            skill_vector=data.get("skill_vector", {}),
            generation=data.get("generation", 0),
            parent_paths=data.get("parent_paths"),
            is_gibberish=data.get("is_gibberish", False),
            novelty_score=data.get("novelty_score"),
        )


class BioArchive:
    """Archive with Dominated Novelty Search."""
    
    def __init__(
        self,
        archive_size: int = 50,
        k_neighbors: int = 5,
        dominated_score: float = 999.0,
        use_skill_ratio: bool = True,
        use_difficulty_weights: bool = False,
    ):
        self.archive_size = archive_size
        self.k_neighbors = k_neighbors
        self.dominated_score = dominated_score
        self.use_skill_ratio = use_skill_ratio
        self.use_difficulty_weights = use_difficulty_weights
        
        self.solutions: List[BioSolution] = []
        self._task_order: Optional[List[str]] = None
    
    def add_solution(self, solution: BioSolution) -> None:
        """Add a solution (used during initialization)."""
        self.solutions.append(solution)
        self._update_task_order()
    
    def update(self, new_solutions: List[BioSolution]) -> None:
        """Update archive with new solutions using DNS."""
        if not self.solutions and not new_solutions:
            return
        
        all_solutions = self.solutions + new_solutions
        
        # Compute difficulty weights if enabled
        difficulty_weights = None
        if self.use_difficulty_weights:
            difficulty_weights = self.compute_difficulty_weights(all_solutions)
        
        # Compute novelty scores
        solution_scores = []
        fittest_model_found = False
        
        for i, solution in enumerate(all_solutions):
            # Find fitter solutions
            if not fittest_model_found:
                fitter = [s for s in all_solutions if s.fitness > solution.fitness]
            else:
                fitter = [
                    s for s in all_solutions
                    if s.fitness >= solution.fitness and s.model_path != solution.model_path
                ]
            
            novelty, fittest_model_found = self._compute_novelty_score(
                solution, fitter, difficulty_weights
            )
            solution.novelty_score = novelty
            solution_scores.append((i, novelty))
        
        # Sort by novelty (descending) and keep top
        solution_scores.sort(key=lambda x: x[1], reverse=True)
        
        selected_indices = [
            idx for idx, _ in solution_scores[: self.archive_size]
        ]
        self.solutions = [all_solutions[i] for i in selected_indices]
        
        logger.info(
            f"Archive updated: {len(self.solutions)} solutions. "
            f"Novelty range: {solution_scores[-1][1]:.4f} to {solution_scores[0][1]:.4f}"
        )
    
    def _compute_novelty_score(
        self,
        solution: BioSolution,
        fitter_solutions: List[BioSolution],
        difficulty_weights: Optional[List[float]] = None,
    ) -> Tuple[float, bool]:
        """Compute dominated novelty score for a solution."""
        if not fitter_solutions:
            return self.dominated_score, True
        
        # Build skill vectors with consistent task ordering
        if self._task_order is None:
            self._update_task_order()
        
        sol_vector = self._to_vector(solution.skill_vector)
        
        if self.use_skill_ratio:
            skill_scores = []
            for fitter in fitter_solutions:
                fitter_vector = self._to_vector(fitter.skill_vector)
                
                if self.use_difficulty_weights and difficulty_weights is not None:
                    unique_solved = sum(
                        difficulty_weights[i]
                        for i, (weak, strong) in enumerate(zip(sol_vector, fitter_vector))
                        if weak and not strong
                    )
                    total_unsolved = sum(
                        difficulty_weights[i]
                        for i, strong in enumerate(fitter_vector)
                        if not strong
                    )
                    if total_unsolved == 0:
                        skill_scores.append(0.0)
                    else:
                        skill_scores.append(unique_solved / total_unsolved * 100.0)
                else:
                    unique_solved = sum(1 for w, s in zip(sol_vector, fitter_vector) if w and not s)
                    total_unsolved = sum(1 for s in fitter_vector if not s)
                    if total_unsolved == 0:
                        skill_scores.append(0.0)
                    else:
                        skill_scores.append(unique_solved / total_unsolved * 100.0)
            
            skill_scores.sort()
            k = min(self.k_neighbors, len(skill_scores))
            if k == 0:
                return self.dominated_score, True
            return sum(skill_scores[:k]) / k, False
        
        else:
            # Hamming distance
            distances = []
            for fitter in fitter_solutions:
                fitter_vector = self._to_vector(fitter.skill_vector)
                dist = sum(1 for a, b in zip(sol_vector, fitter_vector) if a != b)
                distances.append(dist)
            
            distances.sort()
            k = min(self.k_neighbors, len(distances))
            if k == 0:
                return self.dominated_score, True
            return sum(distances[:k]) / k, False
    
    def _update_task_order(self) -> None:
        """Update the consistent task ordering from all solutions."""
        all_tasks = set()
        for sol in self.solutions:
            all_tasks.update(sol.skill_vector.keys())
        self._task_order = sorted(list(all_tasks))
    
    def _to_vector(self, skill_vector: Dict[str, float]) -> List[bool]:
        """Convert skill vector to binary list with consistent ordering."""
        if self._task_order is None:
            self._update_task_order()
        
        threshold = 0.5
        return [skill_vector.get(t, 0.0) >= threshold for t in self._task_order]
    
    def compute_difficulty_weights(self, solutions: List[BioSolution]) -> List[float]:
        """Compute per-task difficulty based on failure rate."""
        if not solutions or not self._task_order:
            return []
        
        failure_counts = [0] * len(self._task_order)
        for sol in solutions:
            vector = self._to_vector(sol.skill_vector)
            for i, success in enumerate(vector):
                if not success:
                    failure_counts[i] += 1
        
        n = len(solutions)
        return [count / n for count in failure_counts]
    
    def get_best(self) -> Optional[BioSolution]:
        """Return the solution with highest fitness."""
        if not self.solutions:
            return None
        return max(self.solutions, key=lambda s: s.fitness)
    
    def get_most_novel(self) -> Optional[BioSolution]:
        """Return the solution with highest novelty score."""
        if not self.solutions:
            return None
        return max(self.solutions, key=lambda s: s.novelty_score or 0.0)
    
    def get_coverage(self, threshold: float = 0.5) -> float:
        """Compute task coverage (fraction of tasks solved by at least one model)."""
        if not self.solutions or not self._task_order:
            return 0.0
        
        solved = [False] * len(self._task_order)
        for sol in self.solutions:
            vector = self._to_vector(sol.skill_vector)
            for i, success in enumerate(vector):
                if success:
                    solved[i] = True
        
        return sum(solved) / len(solved)
    
    def save(self, path: str) -> None:
        """Save archive to JSON."""
        data = {
            "archive_size": self.archive_size,
            "solutions": [s.to_dict() for s in self.solutions],
            "task_order": self._task_order,
        }
        with open(path, "w") as f:
            json.dump(data, f, indent=2)
        logger.info(f"Saved archive to {path}")
    
    def load(self, path: str) -> None:
        """Load archive from JSON."""
        with open(path, "r") as f:
            data = json.load(f)
        self.archive_size = data.get("archive_size", self.archive_size)
        self.solutions = [BioSolution.from_dict(d) for d in data.get("solutions", [])]
        self._task_order = data.get("task_order")
        logger.info(f"Loaded archive with {len(self.solutions)} solutions from {path}")


class DominatedNoveltySearch:
    """Alias for backward compatibility."""
    pass