self_evolving_model.py

#!/usr/bin/env python3
"""
╔══════════════════════════════════════════════════════════════════════════╗
║                    SELF-EVOLVING NEURAL NETWORK                          ║
║  A file that evolves itself based on a pre-existing model and can be     ║
║  trained locally by itself using evolutionary strategies.                ║
║                                                                          ║
║  • Architecture genome system (layers, units, activations, lr)           ║
║  • Self-generates training data if none provided                         ║
║  • Mutates & selects the best models across generations                  ║
║  • Saves checkpoints and evolution history to disk                       ║
║  • Resumes from checkpoints automatically                                ║
║                                                                          ║
║  Usage:                                                                  ║
║    python3 self_evolving_model.py                    # quick run         ║
║    python3 self_evolving_model.py --generations 100  # more generations  ║
║    python3 self_evolving_model.py --pop-size 20      # larger population ║
║    python3 self_evolving_model.py --reset            # fresh start       ║
╚══════════════════════════════════════════════════════════════════════════╝
"""

import os
import sys
import json
import copy
import random
import pickle
import hashlib
import argparse
import datetime
import threading
import time
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple, Callable

import numpy as np

# Suppress TensorFlow warnings for cleaner output
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

# ─────────────────────────────────────────────────────────────────────
#  Constants & Defaults
# ─────────────────────────────────────────────────────────────────────
CHECKPOINT_DIR = Path("evo_checkpoints")

ACTIVATION_POOL = ["relu", "tanh", "sigmoid", "elu", "selu", "swish", "linear"]
OPTIMIZER_POOL = ["adam", "sgd", "rmsprop", "adamw"]

# Genome pool bounds (tunable via --max-units / --max-layers so Colab/laptop
# runs can scale parameters without editing code)
MAX_LAYERS = 6
BASE_UNIT_POOL = [16, 32, 48, 64, 96, 128, 192, 256]
UNIT_POOL = list(BASE_UNIT_POOL)


def configure_genome_pool(max_units: int = 256, max_layers: int = 6):
    """
    Widen the genome search space. Unit sizes are generated by repeatedly
    growing the base pool until max_units is covered; layer count is capped.
    Used by --max-units / --max-layers so v2/v3 runs on GPU can scale up.
    """
    global UNIT_POOL, MAX_LAYERS
    # ⛔ Bound the tunable flags by the hard safety walls (the AI cannot exceed them)
    max_units = min(max(8, int(max_units)), SafetyGates.HARD_MAX_UNITS)
    max_layers = min(max(1, int(max_layers)), SafetyGates.HARD_MAX_LAYERS)
    units = [u for u in BASE_UNIT_POOL if u <= max_units]
    if max_units > BASE_UNIT_POOL[-1]:
        n = BASE_UNIT_POOL[-1]
        while n < max_units:
            n = min(max_units, int(n * 1.5))
            units.append(n)
    units = units or [min(max_units, BASE_UNIT_POOL[-1])]  # never empty
    UNIT_POOL = sorted(set(units))
    MAX_LAYERS = max(1, int(max_layers))
    log(f"Genome pool: units={UNIT_POOL}, max_layers={MAX_LAYERS}", "info")


LOSS_FUNCTIONS = {
    "regression": "mse",
    "classification": "categorical_crossentropy",
}

# ═════════════════════════════════════════════════════════════════════
#  ⛔ SAFETY GATES — HARD LIMITS THE AI CANNOT CHANGE
#  These constants live OUTSIDE the genome. They are not stored in
#  checkpoints, not part of any genome config, and the evolutionary
#  operators (mutate/crossover) can never touch them. They are enforced
#  as absolute ceilings at runtime on EVERY evaluation — even if a
#  loaded/corrupt checkpoint tries to exceed them.
# ═════════════════════════════════════════════════════════════════════
class SafetyGates:
    """
    Immutable safety valves & failure gates. Change these ONLY by editing
    this source file — the AI (evolution) has no path to modify them.
    """

    # ── Absolute architecture ceilings (hard walls) ──
    HARD_MAX_LAYERS = 16          # never more hidden layers than this
    HARD_MAX_UNITS  = 2048        # never more units per layer than this
    HARD_MAX_PARAMS = 5_000_000   # never more total model parameters than this

    # ── Fitness / loss sanity bounds ──
    MIN_FITNESS = 0.0             # fitness below → rejected
    MAX_FITNESS = 1e9             # absurd fitness → clamped
    MAX_VAL_LOSS = 1e8            # val_loss above → treated as failed model

    # ── Hyperparameter bounds ──
    MIN_LR = 1e-7
    MAX_LR = 1.0
    MAX_DROPOUT = 0.9
    MIN_BATCH = 8
    MAX_BATCH = 1024

    # ── Runtime failure gates ──
    DEFAULT_MAX_MINUTES = 0       # 0 = unlimited (set via --max-minutes)
    STAGNATION_LIMIT = 12         # generations w/o improvement → halt

    @staticmethod
    def enforce_genome(genome: "Genome") -> "Genome":
        """
        Clamp any genome (even corrupt/old checkpoints) to the hard limits.
        Called on EVERY evaluation, so the AI cannot bypass the walls.
        """
        cfg = genome.config
        cfg["layers"] = cfg.get("layers", [])[:SafetyGates.HARD_MAX_LAYERS]
        for l in cfg["layers"]:
            l["units"] = min(max(1, int(l.get("units", 16))), SafetyGates.HARD_MAX_UNITS)
            l["dropout"] = max(0.0, min(float(l.get("dropout", 0.0)), SafetyGates.MAX_DROPOUT))
        cfg["num_layers"] = len(cfg["layers"])
        cfg["learning_rate"] = max(SafetyGates.MIN_LR, min(float(cfg.get("learning_rate", 0.001)), SafetyGates.MAX_LR))
        cfg["batch_size"] = max(SafetyGates.MIN_BATCH, min(int(cfg.get("batch_size", 32)), SafetyGates.MAX_BATCH))
        return genome


# ═════════════════════════════════════════════════════════════════════
#  CLASS: TextVectorizer
#  Turns raw English text into integer token ids for the text path.
# ═════════════════════════════════════════════════════════════════════
class TextVectorizer:
    """
    Tokenizes English sentences into integer sequences using Keras
    TextVectorization (no external tokenizer dependency). Fit on the
    training texts once, then encode() maps text → int ids.
    """

    def __init__(self, vocab_size: int = 10000, max_len: int = 128):
        self.vocab_size = vocab_size
        self.max_len = max_len
        self.tv = layers.TextVectorization(
            max_tokens=vocab_size,
            output_mode="int",
            output_sequence_length=max_len,
            name="text_vectorizer",
        )
        self.fitted = False

    def fit(self, texts: List[str]) -> "TextVectorizer":
        self.tv.adapt(np.array(texts, dtype=object))
        self.fitted = True
        return self

    def encode(self, texts: List[str]) -> np.ndarray:
        if not self.fitted:
            raise RuntimeError("TextVectorizer.fit() must be called before encode()")
        return self.tv(np.array(texts, dtype=object)).numpy().astype(np.int32)

    def vocab_used(self) -> int:
        """Actual fitted vocabulary size (includes reserved tokens)."""
        return self.tv.vocabulary_size() if self.fitted else self.vocab_size

# ─────────────────────────────────────────────────────────────────────
#  Utility: Pretty printing
# ─────────────────────────────────────────────────────────────────────
class Colors:
    """ANSI color codes for terminal output."""
    HEADER  = "\033[95m"
    BLUE    = "\033[94m"
    CYAN    = "\033[96m"
    GREEN   = "\033[92m"
    YELLOW  = "\033[93m"
    RED     = "\033[91m"
    BOLD    = "\033[1m"
    DIM     = "\033[2m"
    END     = "\033[0m"


def banner(text: str, char: str = "═", width: int = 70):
    print(f"\n{Colors.CYAN}{char * width}")
    print(f"  {Colors.BOLD}{text}{Colors.END}")
    print(f"{Colors.CYAN}{char * width}{Colors.END}\n")


def log(msg: str, level: str = "info"):
    prefix = {
        "info":    f"{Colors.BLUE}[INFO]{Colors.END}",
        "success": f"{Colors.GREEN}[OK]{Colors.END}",
        "warn":    f"{Colors.YELLOW}[WARN]{Colors.END}",
        "error":   f"{Colors.RED}[ERR]{Colors.END}",
        "evo":     f"{Colors.CYAN}[EVO]{Colors.END}",
    }.get(level, "[LOG]")
    timestamp = datetime.datetime.now().strftime("%H:%M:%S")
    print(f"{Colors.DIM}{timestamp}{Colors.END} {prefix} {msg}")


# ═════════════════════════════════════════════════════════════════════
#  CLASS: Genome
#  Represents a neural network architecture as a mutable genome.
# ═════════════════════════════════════════════════════════════════════
class Genome:
    """
    Encodes a neural network architecture as a 'genome' that can be
    mutated, crossed-over, and evaluated for fitness.

    Genome structure:
    {
        "num_layers":       int,          # number of hidden layers
        "layers": [                       # per-layer configuration
            {
                "units":        int,
                "activation":   str,
                "dropout":      float,
                "batch_norm":   bool,
            },
            ...
        ],
        "learning_rate":    float,
        "optimizer":        str,
        "batch_size":       int,
        "top3_features":    list,         # indices of the 3 features gating layer 1
        "task_type":        str,          # "regression" or "classification"
    }
    """

    def __init__(self, config: Optional[Dict] = None):
        if config is not None:
            self.config = config
        else:
            self.config = self._random_genome()
        self.fitness: float = 0.0
        self.generation_born: int = 0
        self.id = hashlib.md5(
            json.dumps(self.config, sort_keys=True).encode()
        ).hexdigest()[:8]

    # ── Random genome factory ──────────────────────────────────────
    @staticmethod
    def _random_genome() -> Dict:
        num_layers = random.randint(1, MAX_LAYERS)
        layers_cfg = []
        for i in range(num_layers):
            layers_cfg.append({
                "units":      random.choice(UNIT_POOL),
                "activation": random.choice(ACTIVATION_POOL),
                "dropout":    round(random.uniform(0.0, 0.5), 2),
                "batch_norm": random.choice([True, False]),
            })
        return {
            "num_layers":    num_layers,
            "layers":        layers_cfg,
            "learning_rate": round(random.choice([0.0001, 0.0005, 0.001, 0.003, 0.005, 0.01]), 6),
            "optimizer":     random.choice(OPTIMIZER_POOL),
            "batch_size":    random.choice([16, 32, 64, 128]),
            "top3_features": sorted(random.sample(range(20), 3)),
            "task_type":     "regression",
        }

    # ── Mutation operators ─────────────────────────────────────────
    def mutate(self, mutation_rate: float = 0.3) -> "Genome":
        """Return a mutated clone of this genome."""
        child = self.clone()
        cfg = child.config

        # Mutate number of layers (add or remove)
        if random.random() < mutation_rate:
            if random.random() < 0.5 and len(cfg["layers"]) > 1:
                # Remove a random layer
                idx = random.randint(0, len(cfg["layers"]) - 1)
                cfg["layers"].pop(idx)
                log(f"  Mutation: removed layer {idx}", "evo")
            elif len(cfg["layers"]) < MAX_LAYERS:
                # Add a new layer at a random position
                idx = random.randint(0, len(cfg["layers"]))
                new_layer = {
                    "units":      random.choice(UNIT_POOL),
                    "activation": random.choice(ACTIVATION_POOL),
                    "dropout":    round(random.uniform(0.0, 0.5), 2),
                    "batch_norm": random.choice([True, False]),
                }
                cfg["layers"].insert(idx, new_layer)
                log(f"  Mutation: added layer at {idx} ({new_layer['units']} units)", "evo")
            cfg["num_layers"] = len(cfg["layers"])

        # Mutate individual layers
        for i, layer in enumerate(cfg["layers"]):
            if random.random() < mutation_rate:
                old_units = layer["units"]
                layer["units"] = random.choice(UNIT_POOL)
                log(f"  Mutation: layer {i} units {old_units} → {layer['units']}", "evo")
            if random.random() < mutation_rate:
                old_act = layer["activation"]
                layer["activation"] = random.choice(ACTIVATION_POOL)
                log(f"  Mutation: layer {i} activation {old_act} → {layer['activation']}", "evo")
            if random.random() < mutation_rate:
                layer["dropout"] = round(max(0, min(0.5, layer["dropout"] + random.uniform(-0.1, 0.1))), 2)
            if random.random() < mutation_rate * 0.5:
                layer["batch_norm"] = not layer["batch_norm"]

        # Mutate learning rate (log-scale perturbation)
        if random.random() < mutation_rate:
            old_lr = cfg["learning_rate"]
            factor = random.choice([0.5, 0.7, 1.0, 1.3, 1.5, 2.0])
            cfg["learning_rate"] = round(max(1e-6, min(0.1, old_lr * factor)), 6)
            log(f"  Mutation: lr {old_lr} → {cfg['learning_rate']}", "evo")

        # Mutate optimizer
        if random.random() < mutation_rate * 0.5:
            old_opt = cfg["optimizer"]
            cfg["optimizer"] = random.choice(OPTIMIZER_POOL)
            log(f"  Mutation: optimizer {old_opt} → {cfg['optimizer']}", "evo")

        # Mutate batch size
        if random.random() < mutation_rate * 0.5:
            cfg["batch_size"] = random.choice([16, 32, 64, 128])

        # Mutate top-3 feature gating (which features drive layer 1)
        # NOTE: .setdefault keeps old-format checkpoints (no top3_features) compatible
        cfg.setdefault("top3_features", sorted(random.sample(range(20), 3)))
        if random.random() < mutation_rate * 0.8:
            idx = random.randrange(len(cfg["top3_features"]))
            old = cfg["top3_features"][idx]
            # Pick a new index not already selected (no duplicates)
            pool = [v for v in range(20) if v not in cfg["top3_features"]] or [0]
            new = random.choice(pool)
            cfg["top3_features"][idx] = new
            cfg["top3_features"].sort()
            log(f"  Mutation: top-3 feature gate {old} → {new} ({cfg['top3_features']})", "evo")

        child.id = hashlib.md5(
            json.dumps(cfg, sort_keys=True).encode()
        ).hexdigest()[:8]
        return child

    # ── Crossover ──────────────────────────────────────────────────
    def crossover(self, other: "Genome") -> "Genome":
        """Produce a child genome by crossing over with another."""
        child_cfg = copy.deepcopy(self.config)
        other_cfg = other.config

        # Crossover layer-by-layer (take from either parent)
        max_layers = max(len(child_cfg["layers"]), len(other_cfg["layers"]))
        child_layers = []
        for i in range(max_layers):
            if i < len(child_cfg["layers"]) and i < len(other_cfg["layers"]):
                parent = random.choice([child_cfg, other_cfg])
                child_layers.append(copy.deepcopy(parent["layers"][i]))
            elif i < len(child_cfg["layers"]):
                child_layers.append(copy.deepcopy(child_cfg["layers"][i]))
            else:
                child_layers.append(copy.deepcopy(other_cfg["layers"][i]))

        child_cfg["layers"] = child_layers
        child_cfg["num_layers"] = len(child_layers)

        # Randomly inherit hyperparams from either parent
        if random.random() < 0.5:
            child_cfg["learning_rate"] = other_cfg["learning_rate"]
        if random.random() < 0.5:
            child_cfg["optimizer"] = other_cfg["optimizer"]
        if random.random() < 0.5:
            child_cfg["batch_size"] = other_cfg["batch_size"]
        if random.random() < 0.5 and "top3_features" in other_cfg:
            # Only inherit when the other parent has an evolved gate (avoid
            # clobbering a good selection with the default on old checkpoints)
            child_cfg["top3_features"] = sorted(other_cfg["top3_features"])

        child = Genome(child_cfg)
        child.generation_born = self.generation_born
        return child

    # ── Utility ────────────────────────────────────────────────────
    def clone(self) -> "Genome":
        c = Genome(copy.deepcopy(self.config))
        c.fitness = self.fitness
        c.generation_born = self.generation_born
        return c

    def to_dict(self) -> Dict:
        return {
            "config": self.config,
            "fitness": self.fitness,
            "generation_born": self.generation_born,
            "id": self.id,
        }

    @classmethod
    def from_dict(cls, data: Dict) -> "Genome":
        g = cls(data["config"])
        g.fitness = data.get("fitness", 0.0)
        g.generation_born = data.get("generation_born", 0)
        g.id = data.get("id", g.id)
        return g

    def summary(self) -> str:
        layers_str = ", ".join(
            f"{l['units']}({l['activation'][:3]})" for l in self.config["layers"]
        )
        return (
            f"Genome[{self.id}] layers={self.config['num_layers']} "
            f"[{layers_str}] lr={self.config['learning_rate']} "
            f"opt={self.config['optimizer']} top3={self.config.get('top3_features')} "
            f"fitness={self.fitness:.4f}"
        )

    def __repr__(self):
        return self.summary()


# ═════════════════════════════════════════════════════════════════════
#  CLASS: DataHandler
#  Manages training data. Generates synthetic data if none provided.
# ═════════════════════════════════════════════════════════════════════
class DataHandler:
    """
    Provides training data for the evolutionary process.
    If no external data is given, generates synthetic datasets
    that the models can learn from.
    """

    def __init__(self, task_type: str = "regression"):
        self.task_type = task_type

    def get_data(
        self,
        x: Optional[np.ndarray] = None,
        y: Optional[np.ndarray] = None,
        n_samples: int = 2000,
        n_features: int = 10,
        seed: int = 42,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Return (X, Y) arrays. If x/y are provided, use them.
        Otherwise generate synthetic data.
        """
        if x is not None and y is not None:
            log(f"Using provided data: X={x.shape}, Y={y.shape}", "info")
            return x, y

        log(f"Generating synthetic {self.task_type} data ({n_samples} samples, {n_features} features)...", "info")
        rng = np.random.RandomState(seed)
        X = rng.uniform(-3.0, 3.0, (n_samples, n_features)).astype(np.float32)

        if self.task_type == "regression":
            # Complex nonlinear target: sum of sin/cos combinations
            Y = (
                np.sin(X[:, 0]) * np.cos(X[:, 1])
                + 0.5 * np.sin(X[:, 2] + X[:, 3])
                + 0.3 * X[:, 4] ** 2
                + rng.normal(0, 0.1, n_samples)
            ).astype(np.float32)
            Y = Y.reshape(-1, 1)
        else:
            # Classification: threshold-based multi-class
            logits = np.sin(X[:, 0]) + np.cos(X[:, 1]) + X[:, 2] * 0.5
            Y = (logits > 0.5).astype(np.float32)
            Y = keras.utils.to_categorical(Y, num_classes=2)

        log(f"Data ready: X={X.shape}, Y={Y.shape}", "success")
        return X, Y

    @staticmethod
    def load_fable_dataset(
        n_samples: int = 10000,
        subset: str = "train",
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Load and featurize the Crownelius/Complete-FABLE.5-traces-2M dataset
        from HuggingFace. Extracts numerical features from heterogeneous
        JSON coding traces and returns (X, Y) arrays.

        Features extracted per row (10-dim vector):
          0: message_content_len   - length of user message content
          1: message_word_count    - word count of user message
          2: code_keyword_freq     - frequency of code keywords (def, class, etc.)
          3: completion_len        - length of assistant completion
          4: cot_len               - length of chain-of-thought
          5: has_tool_use          - whether output_type is tool_use
          6: output_complexity     - len(str(output)) if present
          7: is_user_turn          - whether row is a user message
          8: session_entropy       - hash-based session diversity proxy
          9: text_special_char_ratio - ratio of special chars in text

        Target Y: seen_count (how many times this trace was seen)
        """
        try:
            from datasets import load_dataset as hf_load_dataset
        except ImportError:
            log("datasets library not installed. Run: pip install datasets", "error")
            raise

        log(f"Loading FABLE.5-traces-2M dataset ({n_samples} samples)...", "info")
        ds = hf_load_dataset("Crownelius/Complete-FABLE.5-traces-2M", split=subset)
        if len(ds) > n_samples:
            # ⚠️ Same ordering-bias guard as the text path: if this split is sorted
            # by session/type/timestamp, a head-slice can be strongly biased.
            # Draw a deterministic random sample instead.
            rng = np.random.RandomState(42)
            idxs = rng.choice(len(ds), size=n_samples, replace=False)
            ds = ds.select(sorted(idxs))
        log(f"Dataset loaded: {len(ds)} rows", "success")

        features = []
        targets = []

        for row in ds:
            try:
                row_json = json.loads(row["row_json"])
            except (json.JSONDecodeError, TypeError):
                continue

            # Extract message content
            msg = row_json.get("message", {})
            msg_content = ""
            if isinstance(msg, dict):
                c = msg.get("content", "")
                if isinstance(c, str):
                    msg_content = c
                elif isinstance(c, list):
                    msg_content = " ".join(
                        item.get("text", "") if isinstance(item, dict) else str(item)
                        for item in c
                    )

            # Extract completion and chain-of-thought
            completion = str(row_json.get("completion", "") or "")
            cot = str(row_json.get("cot", "") or "")

            # Output info
            output = row_json.get("output", {})
            output_str = str(output) if output else ""
            output_type = str(row_json.get("output_type", "") or "")

            # Row type
            row_type = str(row_json.get("type", "") or "")

            # --- Build feature vector (10 dims) ---
            # 0: message_content_len (normalized by log)
            f0 = np.log1p(len(msg_content))
            # 1: message word count
            f1 = np.log1p(len(msg_content.split())) if msg_content else 0.0
            # 2: code keyword frequency in message
            code_keywords = ["def", "class", "import", "function", "return",
                             "if", "for", "while", "async", "const"]
            text_lower = msg_content.lower()
            f2 = sum(text_lower.count(kw) for kw in code_keywords)
            f2 = np.log1p(f2)
            # 3: completion length (log-normalized)
            f3 = np.log1p(len(completion))
            # 4: cot length (log-normalized)
            f4 = np.log1p(len(cot))
            # 5: has tool use
            f5 = 1.0 if output_type == "tool_use" else 0.0
            # 6: output complexity
            f6 = np.log1p(len(output_str))
            # 7: is user turn
            f7 = 1.0 if row_type == "user" else 0.0
            # 8: session entropy proxy (deterministic hash)
            session_id = str(row_json.get("sessionId", "") or row_json.get("session", ""))
            f8 = int(hashlib.md5(session_id.encode()).hexdigest(), 16) % 1000 / 1000.0
            # 9: special character ratio in all text
            all_text = msg_content + completion + cot
            if len(all_text) > 0:
                special = sum(1 for c in all_text if not c.isalnum() and not c.isspace())
                f9 = special / len(all_text)
            else:
                f9 = 0.0

            features.append([f0, f1, f2, f3, f4, f5, f6, f7, f8, f9])
            targets.append(float(row.get("seen_count", 1)))

        X = np.array(features, dtype=np.float32)
        Y = np.array(targets, dtype=np.float32).reshape(-1, 1)

        # Normalize features to zero mean, unit variance
        mean = X.mean(axis=0)
        std = X.std(axis=0) + 1e-8
        X = (X - mean) / std

        # Log-normalize target (seen_count is power-law distributed)
        Y = np.log1p(Y)

        log(f"Featurized: X={X.shape}, Y={Y.shape}", "success")
        log(f"  Feature ranges: min={X.min():.2f}, max={X.max():.2f}", "info")
        log(f"  Target range: min={Y.min():.2f}, max={Y.max():.2f}", "info")
        return X, Y

    @staticmethod
    def load_text_dataset(
        dataset: str = "rotten_tomatoes",
        n_samples: int = 3000,
        max_len: int = 128,
        vocab_size: int = 10000,
        subset: str = "train",
    ):
        """
        Load an English sentiment dataset from HuggingFace and convert it to
        (X_int_ids, Y_onehot, vectorizer, n_classes) for text-understanding
        evolution. Default: rotten_tomatoes (binary pos/neg movie reviews).
        """
        try:
            from datasets import load_dataset as hf_load_dataset
        except ImportError:
            log("datasets library not installed. Run: pip install datasets", "error")
            raise

        log(f"Loading English sentiment dataset '{dataset}' ({n_samples} samples)...", "info")
        ds = hf_load_dataset(dataset, split=subset)
        if len(ds) > n_samples:
            # ⚠️ Some HF splits are SORTED by label (e.g. rotten_tomatoes is
            # all-pos then all-neg). A head-slice train[:n] would then be a
            # single class and the model would learn "always positive" — which
            # scores ~100% on that slice but ~50% (chance) on balanced held-out
            # data. Draw a deterministic random sample instead.
            rng = np.random.RandomState(42)
            idxs = rng.choice(len(ds), size=n_samples, replace=False)
            ds = ds.select(sorted(idxs))
        texts = [str(r["text"]) for r in ds]
        labels = np.array([int(r["label"]) for r in ds])
        n_classes = int(labels.max()) + 1
        Y = keras.utils.to_categorical(labels, num_classes=n_classes).astype(np.float32)

        vec = TextVectorizer(vocab_size=vocab_size, max_len=max_len)
        vec.fit(texts)
        X = vec.encode(texts)
        log(f"📖 Text featurized: X={X.shape} (token ids), Y={Y.shape}, vocab={vec.vocab_used():,}", "success")
        return X, Y, vec, n_classes


# ═════════════════════════════════════════════════════════════════════
#  CLASS: ModelEvaluator
#  Builds TensorFlow models from genomes and evaluates their fitness.
# ═════════════════════════════════════════════════════════════════════
class ModelEvaluator:
    """
    Translates a Genome into a Keras model, trains it, and returns
    a fitness score (inverse of validation loss).
    """

    def __init__(self, input_dim: int, output_dim: int, task_type: str = "regression",
                 vectorizer: Optional[TextVectorizer] = None, embed_dim: int = 128):
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.task_type = task_type
        self.vectorizer = vectorizer
        self.embed_dim = embed_dim

    def build_model(self, genome: Genome) -> keras.Model:
        """
        Construct a Keras model from a genome's config. Two paths:
          - text_classification: Embedding + pooling + evolved dense stack
          - otherwise: GATED numeric architecture — the genome's top-3
            features feed layer 1; the rest concatenate into every layer
            after the first (including the output).
        """
        if self.task_type == "text_classification":
            return self._build_text_model(genome)
        cfg = genome.config
        top3, rest = self._resolve_feature_split(genome)
        has_rest = len(rest) > 0

        # Gated inputs: top-3 features (layer 1) + the rest (join later)
        inp_top = layers.Input(shape=(len(top3),), name="input_top3")
        h = inp_top
        if has_rest:
            inp_rest = layers.Input(shape=(len(rest),), name="input_rest")

        # Hidden layers from genome (rest joins every layer after the first)
        for i, layer_cfg in enumerate(cfg["layers"]):
            if i > 0 and has_rest:
                h = layers.Concatenate(name=f"merge_{i}")([h, inp_rest])
            h = layers.Dense(
                units=layer_cfg["units"],
                activation=layer_cfg["activation"],
                name=f"dense_{i}",
            )(h)
            if layer_cfg["batch_norm"]:
                h = layers.BatchNormalization(name=f"bn_{i}")(h)
            if layer_cfg["dropout"] > 0:
                h = layers.Dropout(layer_cfg["dropout"], name=f"drop_{i}")(h)

        # Remaining features also join the output layer
        if has_rest:
            h = layers.Concatenate(name="merge_output")([h, inp_rest])

        # Output layer
        if self.task_type == "regression":
            out = layers.Dense(self.output_dim, activation="linear", name="output")(h)
        else:
            out = layers.Dense(self.output_dim, activation="softmax", name="output")(h)

        model = keras.Model(
            inputs=[inp_top, inp_rest] if has_rest else [inp_top],
            outputs=out,
            name=f"model_{genome.id}",
        )

        # Compile
        optimizer = self._get_optimizer(cfg["optimizer"], cfg["learning_rate"])
        loss = LOSS_FUNCTIONS[self.task_type]
        model.compile(optimizer=optimizer, loss=loss, metrics=["mae"] if self.task_type == "regression" else ["accuracy"])

        return model

    def _build_text_model(self, genome: Genome) -> keras.Model:
        """
        English text-understanding path: token ids → Embedding → global
        average pooling → the genome's evolved dense stack → softmax.
        """
        cfg = genome.config
        vocab = (self.vectorizer.vocab_used() + 2) if self.vectorizer else (self.input_dim + 2)
        inp = layers.Input(shape=(self.input_dim,), dtype="int32", name="text_input")
        h = layers.Embedding(vocab, self.embed_dim, name="embedding")(inp)
        h = layers.GlobalAveragePooling1D(name="text_pool")(h)

        for i, layer_cfg in enumerate(cfg["layers"]):
            h = layers.Dense(
                units=layer_cfg["units"],
                activation=layer_cfg["activation"],
                name=f"dense_{i}",
            )(h)
            if layer_cfg["batch_norm"]:
                h = layers.BatchNormalization(name=f"bn_{i}")(h)
            if layer_cfg["dropout"] > 0:
                h = layers.Dropout(layer_cfg["dropout"], name=f"drop_{i}")(h)

        out = layers.Dense(self.output_dim, activation="softmax", name="output")(h)
        model = keras.Model(inputs=inp, outputs=out, name=f"model_{genome.id}")

        optimizer = self._get_optimizer(cfg["optimizer"], cfg["learning_rate"])
        model.compile(optimizer=optimizer, loss=LOSS_FUNCTIONS["classification"], metrics=["accuracy"])
        return model

    def _resolve_feature_split(self, genome: Genome) -> Tuple[List[int], List[int]]:
        """
        Resolve the genome's top-3 feature selection against the real input
        dimension (wraps out-of-range indices, dedupes, pads).
        Returns (top3_indices, rest_indices).
        """
        n = self.input_dim
        raw = genome.config.get("top3_features", [0, 1, 2])
        top: List[int] = []
        for idx in raw:
            idx = int(idx) % n
            if idx not in top:
                top.append(idx)
            if len(top) == 3:
                break
        for i in range(n):
            if len(top) >= 3:
                break
            if i not in top:
                top.append(i)
        top = top[:3]
        rest = [i for i in range(n) if i not in top]
        return top, rest

    def split_features(self, X: np.ndarray, genome: Genome) -> Tuple[np.ndarray, Optional[np.ndarray]]:
        """Split X into (X_top3, X_rest) for the gated architecture.
        Text path has no feature gating — returns (X, None)."""
        if self.task_type == "text_classification":
            return X, None
        top3, rest = self._resolve_feature_split(genome)
        X_rest = X[:, rest] if rest else None
        return X[:, top3], X_rest

    @staticmethod
    def _get_optimizer(name: str, lr: float):
        optimizers = {
            "adam":    keras.optimizers.Adam(learning_rate=lr),
            "sgd":     keras.optimizers.SGD(learning_rate=lr, momentum=0.9),
            "rmsprop": keras.optimizers.RMSprop(learning_rate=lr),
            "adamw":   keras.optimizers.AdamW(learning_rate=lr, weight_decay=1e-4),
        }
        return optimizers.get(name, keras.optimizers.Adam(learning_rate=lr))

    def train_and_evaluate(
        self,
        genome: Genome,
        X: np.ndarray,
        Y: np.ndarray,
        epochs: int = 15,
        verbose: int = 0,
    ) -> float:
        """
        Build, train, and evaluate a model from the genome.
        Returns a fitness score (higher is better).
        """
        try:
            # ⛔ SAFETY: clamp genome to hard limits before building
            SafetyGates.enforce_genome(genome)
            model = self.build_model(genome)
            batch_size = genome.config["batch_size"]

            # Train/val split (gated architecture: split features too)
            split = int(0.8 * len(X))
            X_top, X_rest = self.split_features(X, genome)
            Y_train, Y_val = Y[:split], Y[split:]
            X_top_train, X_top_val = X_top[:split], X_top[split:]

            if X_rest is not None:
                X_rest_train, X_rest_val = X_rest[:split], X_rest[split:]
                train_inputs = [X_top_train, X_rest_train]
                val_inputs = [X_top_val, X_rest_val]
            else:
                train_inputs = X_top_train
                val_inputs = X_top_val

            # 📉 LR scheduling: halve LR when val_loss plateaus (2 epochs patience)
            # so evolution can refine good architectures instead of overshooting.
            lr_schedule = keras.callbacks.ReduceLROnPlateau(
                monitor="val_loss", factor=0.5, patience=2, min_lr=1e-6, verbose=0
            )
            history = model.fit(
                train_inputs, Y_train,
                validation_data=(val_inputs, Y_val),
                epochs=epochs,
                batch_size=batch_size,
                verbose=verbose,
                callbacks=[lr_schedule],
            )

            # ⛔ FAILURE GATE: reject NaN/Inf/exploded validation loss
            val_losses = history.history["val_loss"]
            best_val_loss = min(val_losses)
            if not np.isfinite(best_val_loss) or best_val_loss > SafetyGates.MAX_VAL_LOSS:
                log(f"  ⛔ Safety gate: invalid val_loss {best_val_loss} for {genome.id}", "warn")
                del model
                keras.backend.clear_session()
                return 0.0

            # ⛔ FAILURE GATE: reject models beyond the hard parameter wall
            num_params = model.count_params()
            if num_params > SafetyGates.HARD_MAX_PARAMS:
                log(f"  ⛔ Safety gate: {num_params:,} params > hard ceiling for {genome.id}", "warn")
                del model
                keras.backend.clear_session()
                return 0.0

            # Fitness = inverse of best validation loss
            # Also penalize overly complex models slightly (Occam's razor)
            complexity_penalty = 1.0 + 1e-6 * num_params  # tiny penalty for huge models

            # 🎯 Classification: maximize VALIDATION ACCURACY directly. Loss-only
            # fitness rewards memorization (near-zero loss on trivial/constant fits),
            # which is exactly how the first text run "learned" 100% train / 50% test.
            # Accuracy^2 gives a ~0-100 scale and makes generalization the target.
            if self.task_type in ("classification", "text_classification"):
                val_accs = history.history.get("val_accuracy")
                best_val_acc = max(val_accs) if val_accs else 0.0
                fitness = (best_val_acc ** 2) * 100.0 / complexity_penalty
            else:
                fitness = 1.0 / ((best_val_loss + 1e-7) * complexity_penalty)
            fitness = max(SafetyGates.MIN_FITNESS, min(fitness, SafetyGates.MAX_FITNESS))  # clamp

            # Clean up
            del model
            keras.backend.clear_session()

            return fitness

        except Exception as e:
            log(f"  Model {genome.id} failed: {e}", "warn")
            return 0.0


# ═════════════════════════════════════════════════════════════════════
#  CLASS: EvolutionEngine
#  The main orchestrator: manages population, selection, mutation,
#  checkpointing, and the evolutionary training loop.
# ═════════════════════════════════════════════════════════════════════
class EvolutionEngine:
    """
    Runs the evolutionary loop:
    1. Initialize population of genomes
    2. Evaluate fitness of each genome by building & training it
    3. Select top performers (elitism)
    4. Reproduce via mutation & crossover
    5. Repeat for N generations
    6. Save checkpoints & history throughout
    """

    def __init__(
        self,
        pop_size: int = 8,
        generations: int = 20,
        elite_ratio: float = 0.3,
        mutation_rate: float = 0.3,
        train_epochs: int = 15,
        checkpoint_dir: Path = CHECKPOINT_DIR,
        task_type: str = "regression",
        vectorizer: Optional[TextVectorizer] = None,
        embed_dim: int = 128,
        max_minutes: int = SafetyGates.DEFAULT_MAX_MINUTES,
        stagnation_limit: int = SafetyGates.STAGNATION_LIMIT,
    ):
        self.pop_size = pop_size
        self.generations = generations
        self.elite_count = min(max(2, int(pop_size * elite_ratio)), pop_size - 1)
        self.mutation_rate = mutation_rate
        self.train_epochs = train_epochs
        self.checkpoint_dir = Path(checkpoint_dir)
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        self.task_type = task_type
        self.vectorizer = vectorizer
        self.embed_dim = embed_dim
        self.max_minutes = int(max_minutes)
        self.stagnation_limit = int(stagnation_limit)
        # Text-path label used when persisting vectorizer metadata (set by CLI)
        self.dataset_label = "cornell-movie-review-data/rotten_tomatoes"

        self.population: List[Genome] = []
        self.history: List[Dict] = []
        self.best_genome: Optional[Genome] = None
        self.current_generation = 0

        # Safety-gate tracking (set at run() start)
        self._start_time: float = 0.0
        self._stagnant_gens: int = 0
        self.safety_trips: List[str] = []

        # Threading support for GUI
        self.stop_event = threading.Event()
        self.on_generation_complete: Optional[Callable] = None

    # ── Population management ──────────────────────────────────────
    def initialize_population(self):
        """Create an initial random population."""
        log(f"Initializing population of {self.pop_size} genomes...", "info")
        self.population = [Genome() for _ in range(self.pop_size)]
        for g in self.population:
            g.generation_born = 0
        log(f"Population ready. Genome examples:", "success")
        for g in self.population[:3]:
            log(f"  {g}", "info")

    # ── Selection ──────────────────────────────────────────────────
    def selection(self):
        """Keep the top performers (elitism) and discard the rest."""
        self.population.sort(key=lambda g: g.fitness, reverse=True)
        elites = self.population[:self.elite_count]
        log(f"  Selected top {self.elite_count} elites:", "evo")
        for g in elites:
            log(f"    {g}", "evo")
        return elites

    # ── Reproduction ───────────────────────────────────────────────
    def reproduce(self, elites: List[Genome]):
        """Create next generation from elites via mutation & crossover."""
        next_gen = [e.clone() for e in elites]  # carry elites forward

        while len(next_gen) < self.pop_size:
            if random.random() < 0.7 and len(elites) >= 2:
                # Crossover + mutation
                p1, p2 = random.sample(elites, 2)
                child = p1.crossover(p2)
                child = child.mutate(self.mutation_rate)
            else:
                # Pure mutation from a random elite
                parent = random.choice(elites)
                child = parent.mutate(self.mutation_rate)

            child.generation_born = self.current_generation + 1
            child.fitness = 0.0
            next_gen.append(child)

        self.population = next_gen

    # ── Checkpointing ──────────────────────────────────────────────
    def save_checkpoint(self):
        """Save full evolution state to disk."""
        ckpt_path = self.checkpoint_dir / "checkpoint.pkl"
        state = {
            "current_generation": self.current_generation,
            "population": [g.to_dict() for g in self.population],
            "history": self.history,
            "best_genome": self.best_genome.to_dict() if self.best_genome else None,
            "params": {
                "pop_size": self.pop_size,
                "generations": self.generations,
                "elite_count": self.elite_count,
                "mutation_rate": self.mutation_rate,
                "train_epochs": self.train_epochs,
            },
        }
        with open(ckpt_path, "wb") as f:
            pickle.dump(state, f)
        log(f"  Checkpoint saved (gen {self.current_generation})", "info")

    def load_checkpoint(self) -> bool:
        """Resume from checkpoint if available. Returns True if loaded."""
        ckpt_path = self.checkpoint_dir / "checkpoint.pkl"
        if not ckpt_path.exists():
            return False

        try:
            with open(ckpt_path, "rb") as f:
                state = pickle.load(f)

            self.current_generation = state["current_generation"]
            self.population = [Genome.from_dict(g) for g in state["population"]]
            self.history = state["history"]
            if state["best_genome"]:
                self.best_genome = Genome.from_dict(state["best_genome"])

            log(f"Resumed from checkpoint: generation {self.current_generation}", "success")
            return True
        except Exception as e:
            log(f"Failed to load checkpoint: {e}", "warn")
            return False

    def save_best_model(self, evaluator: ModelEvaluator, X: np.ndarray, Y: np.ndarray):
        """Rebuild, retrain, and save the best genome's model."""
        if self.best_genome is None:
            return

        log("Saving best model to disk...", "info")
        SafetyGates.enforce_genome(self.best_genome)
        model = evaluator.build_model(self.best_genome)
        split = int(0.8 * len(X))
        X_top, X_rest = evaluator.split_features(X, self.best_genome)
        if X_rest is not None:
            train_inputs = [X_top[:split], X_rest[:split]]
        else:
            train_inputs = X_top[:split]
        model.fit(train_inputs, Y[:split], epochs=self.train_epochs * 2, batch_size=self.best_genome.config["batch_size"], verbose=0)

        save_path = self.checkpoint_dir / "best_model.keras"
        saved_ok = False
        # 🌍 TEXT PATH: bake the fitted TextVectorizer into the saved model so the
        # artifact accepts RAW English strings end-to-end (no separate tokenization
        # step at inference time). Numeric path saves the plain model as before.
        if self.task_type == "text_classification" and self.vectorizer is not None:
            try:
                raw_in = layers.Input(shape=(), dtype="string", name="raw_text")
                tokens = self.vectorizer.tv(raw_in)
                preds = model(tokens)
                serving = keras.Model(raw_in, preds, name=f"text_serving_{self.best_genome.id}")
                serving.save(str(save_path))
                saved_ok = True
                log(f"🌍 Text serving model saved (raw English → sentiment): {save_path}", "success")
                # Also keep the token-id core model (best-effort — never clobber the
                # already-saved serving model if this optional save fails)
                try:
                    model.save(str(self.checkpoint_dir / "best_model_ids.keras"))
                except Exception as e2:
                    log(f"Could not save token-id core model ({e2}) — serving model already saved", "warn")
                del serving
            except Exception as e:
                if not saved_ok:
                    log(f"Could not bake vectorizer into saved model ({e}); saving core model instead", "warn")
                    try:
                        model.save(str(save_path))
                        saved_ok = True
                    except Exception as e2:
                        log(f"Could not save core model either ({e2}) — skipping model save", "error")
                else:
                    log(f"Serving model saved but a later step failed ({e}); keeping serving model", "warn")
        else:
            try:
                model.save(str(save_path))
                saved_ok = True
            except Exception as e:
                log(f"Could not save model ({e}) — skipping model save", "error")
        del model
        keras.backend.clear_session()
        if saved_ok:
            log(f"Best model saved to {save_path}", "success")
        else:
            log("⚠️ No best model was saved to disk", "error")

        # Also save genome config as JSON
        config_path = self.checkpoint_dir / "best_genome.json"
        with open(config_path, "w") as f:
            json.dump(self.best_genome.to_dict(), f, indent=2)
        log(f"Best genome config saved to {config_path}", "success")

        # 🌍 TEXT PATH: persist vectorizer config so eval/loading can reconstruct it
        if self.task_type == "text_classification" and self.vectorizer is not None:
            vcfg = {
                "vocab_size": self.vectorizer.vocab_size,
                "max_len": self.vectorizer.max_len,
                "vocab_used": self.vectorizer.vocab_used(),
                "dataset": getattr(self, "dataset_label", "cornell-movie-review-data/rotten_tomatoes"),
                "vocab": list(self.vectorizer.tv.get_vocabulary()),  # exact fitted vocab
            }
            vpath = self.checkpoint_dir / "vectorizer_config.json"
            with open(vpath, "w") as f:
                json.dump(vcfg, f, indent=2)
            log(f"Vectorizer config saved to {vpath}", "success")

    # ── Diversity pressure ────────────────────────────────────────
    @staticmethod
    def _config_distance(cfg_a: Dict, cfg_b: Dict) -> int:
        """
        Cheap topology distance: sum of per-layer unit differences. Returns a huge
        value when layer counts differ (structurally very different networks).
        """
        la = [l["units"] for l in cfg_a.get("layers", [])]
        lb = [l["units"] for l in cfg_b.get("layers", [])]
        if len(la) != len(lb):
            return 10_000_000
        return sum(abs(a - b) for a, b in zip(la, lb))

    def apply_diversity_pressure(self):
        """
        🧬 Penalize (a) exact duplicate configs and (b) genomes whose layer
        topology is nearly identical to the current best, so the population does
        not collapse onto a single local optimum. In-place on self.population.
        """
        counts = {}
        for g in self.population:
            counts[g.id] = counts.get(g.id, 0) + 1
        best_cfg = self.best_genome.config if self.best_genome else None
        best_id = self.best_genome.id if self.best_genome else None
        for g in self.population:
            # The champion itself is NEVER penalized (it must stay comparable to the
            # stored best for stagnation tracking and re-selection to work).
            if best_id is not None and g.id == best_id:
                continue
            if counts[g.id] > 1:
                g.fitness *= 0.5
                log(f"  🧬 Diversity: duplicate config {g.id} — fitness halved", "warn")
            elif best_cfg is not None and self._config_distance(g.config, best_cfg) < 32:
                g.fitness *= 0.9
                log(f"  🧬 Diversity: {g.id} too similar to best — fitness x0.9", "warn")

    # ── Logging ────────────────────────────────────────────────────
    def log_generation(self, gen: int, fitnesses: List[float]):
        """Log and store generation statistics."""
        stats = {
            "generation": gen,
            "best_fitness": max(fitnesses),
            "avg_fitness": float(np.mean(fitnesses)),
            "worst_fitness": min(fitnesses),
            "std_fitness": float(np.std(fitnesses)),
            "timestamp": datetime.datetime.now().isoformat(),
        }
        self.history.append(stats)

        best = stats["best_fitness"]
        avg = stats["avg_fitness"]
        log(
            f"Gen {gen:3d} │ Best: {best:10.2f} │ Avg: {avg:10.2f} │ "
            f"Std: {stats['std_fitness']:8.2f} │ Pop: {len(self.population)}",
            "evo",
        )

        # Notify GUI callback if set
        if self.on_generation_complete:
            try:
                self.on_generation_complete(stats, self.best_genome)
            except Exception:
                pass

    def print_evolution_summary(self):
        """Print a final summary of the evolution."""
        banner("EVOLUTION COMPLETE", "━")

        # ⛔ Make gate halts VISIBLE — a safety-stop must never look like a normal finish
        if self.safety_trips:
            log(f"⛔ Halted by safety gates: {'; '.join(self.safety_trips)}", "error")

        if self.best_genome:
            log(f"Best genome found:", "success")
            log(f"  {self.best_genome}", "success")

        if self.history:
            log(f"Generations run: {len(self.history)}", "info")
            log(f"Initial best fitness: {self.history[0]['best_fitness']:.4f}", "info")
            log(f"Final best fitness:   {self.history[-1]['best_fitness']:.4f}", "info")

            improvement = 0
            if self.history[0]["best_fitness"] > 0:
                improvement = (
                    (self.history[-1]["best_fitness"] - self.history[0]["best_fitness"])
                    / self.history[0]["best_fitness"] * 100
                )
            log(f"Improvement: {improvement:+.1f}%", "success")

        # Save history as JSON
        history_path = self.checkpoint_dir / "evolution_history.json"
        with open(history_path, "w") as f:
            json.dump(self.history, f, indent=2)
        log(f"Full history saved to {history_path}", "info")

    # ── Main evolutionary loop ─────────────────────────────────────
    def run(
        self,
        X: Optional[np.ndarray] = None,
        Y: Optional[np.ndarray] = None,
        reset: bool = False,
    ):
        """
        Run the full evolutionary training loop.

        Args:
            X: Optional input data. If None, synthetic data is generated.
            Y: Optional target data. If None, synthetic data is generated.
            reset: If True, ignore checkpoints and start fresh.
        """
        banner("SELF-EVOLVING NEURAL NETWORK")

        # ── Safety-gate initialization ──
        self._start_time = time.monotonic()
        self._stagnant_gens = 0
        self.safety_trips = []
        log(
            f"⛔ Safety gates armed: max_layers<={SafetyGates.HARD_MAX_LAYERS}, "
            f"max_units<={SafetyGates.HARD_MAX_UNITS}, max_params<={SafetyGates.HARD_MAX_PARAMS:,}, "
            f"stagnation_limit={self.stagnation_limit}, max_minutes={self.max_minutes}",
            "info",
        )

        # ── Resume or initialize ──
        if not reset and self.load_checkpoint():
            log(f"Continuing from generation {self.current_generation + 1}...", "info")
        else:
            if reset:
                log("Reset requested. Starting fresh.", "warn")
            # Only generate a fresh random population if none was pre-seeded
            # (SelfTrainer.continue_evolution seeds from the best genome BEFORE
            # calling run, so we must not clobber it here).
            if not self.population:
                self.initialize_population()
            self.current_generation = 0

        # ── Prepare data ──
        data_handler = DataHandler(task_type=self.task_type)
        X, Y = data_handler.get_data(x=X, y=Y)

        input_dim = X.shape[1]
        output_dim = Y.shape[1] if len(Y.shape) > 1 else 1
        evaluator = ModelEvaluator(
            input_dim, output_dim, task_type=self.task_type,
            vectorizer=self.vectorizer, embed_dim=self.embed_dim,
        )

        # ── Evolutionary loop ──
        start_gen = self.current_generation
        for gen in range(start_gen, self.generations):
            # Check for stop signal from GUI
            if self.stop_event.is_set():
                log("Stop signal received. Halting evolution.", "warn")
                break

            # ⛔ FAILURE GATE: wall-clock budget exceeded
            if self.max_minutes > 0 and (time.monotonic() - self._start_time) / 60 >= self.max_minutes:
                msg = f"max_minutes budget ({self.max_minutes} min) exceeded"
                log(f"⛔ Safety gate: {msg}. Halting evolution.", "warn")
                self.safety_trips.append(msg)
                break

            # ⛔ FAILURE GATE: stagnation (no fitness improvement for N gens)
            if self.stagnation_limit > 0 and self._stagnant_gens >= self.stagnation_limit:
                msg = f"no improvement for {self._stagnant_gens} generations (limit {self.stagnation_limit})"
                log(f"⛔ Safety gate: {msg}. Halting evolution.", "warn")
                self.safety_trips.append(msg)
                break

            self.current_generation = gen
            banner(f"GENERATION {gen + 1} / {self.generations}", "─")

            # Evaluate fitness for each genome
            fitnesses = []
            for i, genome in enumerate(self.population):
                log(f"Evaluating genome {i+1}/{len(self.population)}: {genome.id}", "info")
                fitness = evaluator.train_and_evaluate(
                    genome, X, Y, epochs=self.train_epochs, verbose=0
                )
                genome.fitness = fitness
                fitnesses.append(fitness)
                log(f"  → Fitness: {fitness:.4f}", "success" if fitness > np.median(fitnesses) else "info")

            # 🧬 Keep the population diverse (prevent premature convergence)
            self.apply_diversity_pressure()
            fitnesses = [g.fitness for g in self.population]

            # Update best genome + stagnation tracking
            gen_best = max(self.population, key=lambda g: g.fitness)
            if self.best_genome is None or gen_best.fitness > self.best_genome.fitness:
                self.best_genome = gen_best.clone()
                self._stagnant_gens = 0
                log(f"  ★ New best genome! {self.best_genome.id} (fitness={self.best_genome.fitness:.4f})", "success")
            else:
                self._stagnant_gens += 1

            # Log generation stats
            self.log_generation(gen, fitnesses)

            # Selection & reproduction (skip for last gen)
            if gen < self.generations - 1:
                elites = self.selection()
                self.reproduce(elites)

            # Save checkpoint
            self.save_checkpoint()

        # ── Save final best model ──
        self.save_best_model(evaluator, X, Y)
        self.print_evolution_summary()


# ═════════════════════════════════════════════════════════════════════
#  CLASS: SelfTrainer
#  Wrapper that allows the system to load and continue training
#  an existing saved model, evolving it further.
# ═════════════════════════════════════════════════════════════════════
class SelfTrainer:
    """
    Loads a previously saved model or genome and continues the
    evolutionary process from that point, effectively letting the
    file 'evolve itself' from its own prior state.
    """

    def __init__(self, checkpoint_dir: Path = CHECKPOINT_DIR):
        self.checkpoint_dir = Path(checkpoint_dir)

    def load_best_genome(self) -> Optional[Genome]:
        """Load the best genome from previous evolution."""
        config_path = self.checkpoint_dir / "best_genome.json"
        if not config_path.exists():
            return None
        with open(config_path, "r") as f:
            data = json.load(f)
        genome = Genome.from_dict(data)
        log(f"Loaded previous best genome: {genome.id}", "success")
        return genome

    def continue_evolution(
        self,
        generations: int = 20,
        pop_size: int = 8,
        X: Optional[np.ndarray] = None,
        Y: Optional[np.ndarray] = None,
        task_type: str = "regression",
        vectorizer: Optional[TextVectorizer] = None,
        embed_dim: int = 128,
        max_minutes: int = SafetyGates.DEFAULT_MAX_MINUTES,
        stagnation_limit: int = SafetyGates.STAGNATION_LIMIT,
    ):
        """
        Continue evolving from the best saved genome.
        Seeds a new population with mutations of the best genome.
        """
        banner("CONTINUING EVOLUTION FROM SAVED STATE")

        parent = self.load_best_genome()
        if parent is None:
            log("No previous genome found. Starting fresh evolution.", "warn")
            engine = EvolutionEngine(pop_size=pop_size, generations=generations,
                                     task_type=task_type, vectorizer=vectorizer,
                                     embed_dim=embed_dim, checkpoint_dir=self.checkpoint_dir,
                                     max_minutes=max_minutes, stagnation_limit=stagnation_limit)
            engine.run(X=X, Y=Y, reset=True)
            return

        # Seed population with mutations of the best genome
        log(f"Seeding population with mutations of {parent.id}...", "info")
        engine = EvolutionEngine(pop_size=pop_size, generations=generations,
                                 task_type=task_type, vectorizer=vectorizer,
                                 embed_dim=embed_dim, checkpoint_dir=self.checkpoint_dir,
                                 max_minutes=max_minutes, stagnation_limit=stagnation_limit)
        engine.population = [parent.clone()]
        for _ in range(pop_size - 1):
            child = parent.mutate(mutation_rate=0.4)  # higher mutation for diversity
            engine.population.append(child)

        engine.current_generation = 0
        engine.best_genome = parent  # keep track of parent's fitness
        engine.run(X=X, Y=Y, reset=True)


# ═════════════════════════════════════════════════════════════════════
#  CLI Entry Point
# ═════════════════════════════════════════════════════════════════════
def parse_args():
    parser = argparse.ArgumentParser(
        description="🧬 Self-Evolving Neural Network — evolves its own architecture locally",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python3 self_evolving_model.py                       # Quick run (20 gens, 8 pop)
  python3 self_evolving_model.py --generations 50       # More generations
  python3 self_evolving_model.py --pop-size 15          # Larger population
  python3 self_evolving_model.py --train-epochs 25      # Longer training per eval
  python3 self_evolving_model.py --reset                 # Ignore checkpoint, start fresh
  python3 self_evolving_model.py --continue              # Continue from best saved model
  python3 self_evolving_model.py --max-units 512 --max-layers 8   # Bigger genomes (GPU)
        """,
    )
    parser.add_argument("--pop-size", type=int, default=8, help="Population size (default: 8, min: 2)")
    parser.add_argument("--generations", type=int, default=20, help="Number of generations (default: 20)")
    parser.add_argument("--train-epochs", type=int, default=15, help="Training epochs per evaluation (default: 15)")
    parser.add_argument("--mutation-rate", type=float, default=0.3, help="Mutation rate (default: 0.3)")
    parser.add_argument("--reset", action="store_true", help="Start fresh, ignoring checkpoints")
    parser.add_argument("--continue", dest="continue_evo", action="store_true", help="Continue evolving from best saved model")
    parser.add_argument("--n-samples", type=int, default=2000, help="Number of synthetic data samples (default: 2000)")
    parser.add_argument("--n-features", type=int, default=10, help="Number of input features (default: 10)")
    parser.add_argument("--dataset", type=str, default=None, help="HuggingFace dataset to use (e.g. 'fable' for Crownelius/Complete-FABLE.5-traces-2M)")
    parser.add_argument("--max-units", type=int, default=256, help="Maximum units per layer for genome search (default: 256)")
    parser.add_argument("--max-layers", type=int, default=6, help="Maximum number of hidden layers for genome search (default: 6)")
    parser.add_argument("--max-minutes", type=int, default=0, help="⛔ Safety: hard wall-clock budget in minutes (0 = unlimited)")
    parser.add_argument("--stagnation-limit", type=int, default=SafetyGates.STAGNATION_LIMIT, help=f"⛔ Safety: halt if no fitness improvement for N generations (default: {SafetyGates.STAGNATION_LIMIT})")
    parser.add_argument("--text-dataset", type=str, default=None, help="English text dataset for sentiment understanding (e.g. 'rotten_tomatoes')")
    parser.add_argument("--max-len", type=int, default=128, help="Max token length for text input (default: 128)")
    parser.add_argument("--vocab-size", type=int, default=10000, help="Vocabulary size for text tokenizer (default: 10000)")
    parser.add_argument("--embed-dim", type=int, default=128, help="Embedding dimension for text models (default: 128)")
    parser.add_argument("--checkpoint-dir", type=str, default=None, help="Override checkpoint directory (default: evo_checkpoints)")
    parser.add_argument("--gui", action="store_true", help="Launch the web GUI instead of CLI")
    parser.add_argument("--gui-port", type=int, default=5000, help="Port for the web GUI (default: 5000)")
    return parser.parse_args()


def main():
    args = parse_args()

    # Seed for reproducibility
    random.seed(42)
    np.random.seed(42)
    tf.random.set_seed(42)

    # Configure the genome search space (allows scaling params on GPU/laptop)
    configure_genome_pool(args.max_units, args.max_layers)

    # Limit TF GPU memory growth if GPU available
    gpus = tf.config.experimental.list_physical_devices("GPU")
    if gpus:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        log(f"Found {len(gpus)} GPU(s). Memory growth enabled.", "info")
    else:
        log("No GPU found. Running on CPU.", "info")

    # Validate pop_size
    if args.pop_size < 2:
        log("Population size must be at least 2 for evolution to work.", "error")
        sys.exit(1)

    # Load dataset (numeric FABLE path OR English text path)
    X_data, Y_data = None, None
    vectorizer = None
    task_type = "regression"
    if args.dataset:
        n = args.n_samples if args.n_samples else 10000
        X_data, Y_data = DataHandler.load_fable_dataset(n_samples=n)
    elif args.text_dataset:
        X_data, Y_data, vectorizer, n_classes = DataHandler.load_text_dataset(
            dataset=args.text_dataset, n_samples=args.n_samples,
            max_len=args.max_len, vocab_size=args.vocab_size,
        )
        task_type = "text_classification"
        log(f"📖 English text mode ready: {args.text_dataset} ({len(X_data):,} samples, {n_classes} classes)", "success")

    ckpt_dir = Path(args.checkpoint_dir) if args.checkpoint_dir else CHECKPOINT_DIR

    if args.gui:
        from evo_gui import launch_gui
        launch_gui(
            pop_size=args.pop_size,
            generations=args.generations,
            mutation_rate=args.mutation_rate,
            train_epochs=args.train_epochs,
            X=X_data, Y=Y_data,
            port=args.gui_port,
        )
        sys.exit(0)

    if args.continue_evo:
        # Continue from saved state (text-aware: task type, vectorizer, checkpoint dir)
        trainer = SelfTrainer(checkpoint_dir=ckpt_dir)
        trainer.continue_evolution(
            generations=args.generations,
            pop_size=args.pop_size,
            X=X_data, Y=Y_data,
            task_type=task_type,
            vectorizer=vectorizer,
            embed_dim=args.embed_dim,
            max_minutes=args.max_minutes,
            stagnation_limit=args.stagnation_limit,
        )
    else:
        # Fresh or resumed evolution (numeric FABLE or English text)
        engine = EvolutionEngine(
            pop_size=args.pop_size,
            generations=args.generations,
            mutation_rate=args.mutation_rate,
            train_epochs=args.train_epochs,
            checkpoint_dir=ckpt_dir,
            task_type=task_type,
            vectorizer=vectorizer,
            embed_dim=args.embed_dim,
            max_minutes=args.max_minutes,
            stagnation_limit=args.stagnation_limit,
        )
        engine.dataset_label = args.text_dataset or args.dataset or "synthetic"
        engine.run(X=X_data, Y=Y_data, reset=args.reset)


if __name__ == "__main__":
    main()