reason_first_program/concepts.py

"""
Stage 2: Concept Discovery

Discover semantic concepts that characterize regions of the program space.
Three complementary discovery mechanisms:

1. BehavioralConceptDiscovery (TRACED, 2306.07487):
   Concepts from execution traces — branch coverage, variable quantiles, output signatures

2. SAEConceptDiscovery (CB-SAE, 2512.10805 + DN-CBM, 2407.14499):
   Train Sparse Autoencoder on LLM hidden states; auto-name neurons as concepts

3. AbstractionConceptDiscovery (LILO, 2310.19791 / ReGAL, 2401.16467):
   Discover shared program fragments and name them via LLM
"""

from __future__ import annotations

import ast
import hashlib
import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Optional, Callable
from collections import defaultdict, Counter

import numpy as np

from reason_first_program.program_space import Program, ProgramSpace

logger = logging.getLogger(__name__)


@dataclass
class Concept:
    """
    A semantic concept that characterizes a region of the program space.

    A concept is defined by:
      - name: human-readable label
      - description: natural language description
      - type: 'behavioral', 'representational', or 'structural'
      - detector: a function Program -> float (concept score in [0, 1])
      - vector: optional embedding vector for this concept (for steering)
      - programs: set of program_ids that exhibit this concept
    """

    name: str
    description: str
    type: str  # 'behavioral', 'representational', 'structural'
    detector: Optional[Callable[[Program], float]] = None
    vector: Optional[np.ndarray] = None
    programs: set[str] = field(default_factory=set)
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def concept_id(self) -> str:
        return hashlib.sha256(
            f"{self.name}:{self.type}".encode()
        ).hexdigest()[:12]

    def score(self, program: Program) -> float:
        """Score how strongly a program exhibits this concept."""
        if self.detector is not None:
            return self.detector(program)
        return 1.0 if program.program_id in self.programs else 0.0

    def to_dict(self) -> dict[str, Any]:
        return {
            "concept_id": self.concept_id,
            "name": self.name,
            "description": self.description,
            "type": self.type,
            "n_programs": len(self.programs),
            "has_vector": self.vector is not None,
            "metadata": self.metadata,
        }


class ConceptSet:
    """A collection of concepts discovered for a program space."""

    def __init__(self):
        self._concepts: dict[str, Concept] = {}

    def add(self, concept: Concept) -> None:
        self._concepts[concept.concept_id] = concept

    def get(self, concept_id: str) -> Optional[Concept]:
        return self._concepts.get(concept_id)

    def get_by_name(self, name: str) -> Optional[Concept]:
        for c in self._concepts.values():
            if c.name == name:
                return c
        return None

    @property
    def concepts(self) -> list[Concept]:
        return list(self._concepts.values())

    @property
    def names(self) -> list[str]:
        return [c.name for c in self._concepts.values()]

    def score_program(self, program: Program) -> dict[str, float]:
        """Score a program against all concepts. Returns {concept_name: score}."""
        return {c.name: c.score(program) for c in self._concepts.values()}

    def score_matrix(self, programs: list[Program]) -> np.ndarray:
        """
        Compute concept score matrix for a list of programs.
        Returns shape (n_programs, n_concepts).
        """
        matrix = np.zeros((len(programs), len(self._concepts)))
        for i, program in enumerate(programs):
            for j, concept in enumerate(self._concepts.values()):
                matrix[i, j] = concept.score(program)
        return matrix

    def concept_lattice(self) -> list[tuple[frozenset[str], frozenset[str]]]:
        """
        Build the Formal Concept Analysis (FCA) lattice.
        Each node = (extent: set of programs, intent: set of concepts).
        
        The Galois connection:
          extent(B) = {p ∈ P : ∀c ∈ B, c.score(p) > 0.5}
          intent(A) = {c ∈ C : ∀p ∈ A, c.score(p) > 0.5}
        """
        all_programs = set()
        for c in self._concepts.values():
            all_programs |= c.programs

        lattice = []
        # For each subset of concepts, compute its extent
        concept_list = list(self._concepts.values())
        n = len(concept_list)

        # For tractability, we compute single-concept and pairwise nodes
        # Full lattice enumeration is exponential
        for i in range(n):
            ci = concept_list[i]
            intent = frozenset([ci.name])
            extent = frozenset(ci.programs)
            if extent:
                lattice.append((extent, intent))

            for j in range(i + 1, n):
                cj = concept_list[j]
                joint_extent = frozenset(ci.programs & cj.programs)
                joint_intent = frozenset([ci.name, cj.name])
                if joint_extent:
                    lattice.append((joint_extent, joint_intent))

        return lattice

    def to_dict(self) -> list[dict[str, Any]]:
        return [c.to_dict() for c in self._concepts.values()]


class ConceptDiscovery(ABC):
    """Abstract base for concept discovery methods."""

    @abstractmethod
    def discover(self, space: ProgramSpace) -> ConceptSet:
        """Discover concepts from a program space."""
        ...


class BehavioralConceptDiscovery(ConceptDiscovery):
    """
    Discover concepts from execution behavior.
    Based on TRACED (2306.07487): execution traces reveal semantic properties
    that are independent of syntactic implementation.

    Concepts discovered:
      - Algorithmic pattern (iterative, recursive, divide-and-conquer)
      - Space complexity class (constant, linear, quadratic)
      - Time complexity class (from execution time quantiles)
      - Error handling strategy (try/except, assertions, guard clauses)
      - Data structure usage (list, dict, set, heap, etc.)
      - Mutation pattern (in-place, copy-on-write, immutable)
      - Control flow pattern (early return, single return, generator)
    """

    # AST-based detectors for syntactic behavioral proxies
    SYNTACTIC_DETECTORS = {
        "uses_recursion": "Implements the solution using recursive function calls",
        "uses_iteration": "Uses explicit loops (for/while) for the main computation",
        "uses_list_comprehension": "Uses list/dict/set comprehensions",
        "uses_generator": "Uses yield or generator expressions",
        "uses_try_except": "Handles errors with try/except blocks",
        "uses_assertion": "Uses assert statements for validation",
        "uses_early_return": "Returns early on edge cases (guard clauses)",
        "uses_sorting": "Uses sort() or sorted() built-in",
        "uses_dict": "Uses dictionaries as primary data structure",
        "uses_set": "Uses sets for membership testing or deduplication",
        "uses_heap": "Uses heapq module for priority queue operations",
        "uses_functional": "Uses map/filter/reduce functional patterns",
        "uses_lambda": "Uses lambda expressions",
        "uses_mutation": "Modifies data structures in-place",
        "single_return": "Has exactly one return statement at the end",
        "uses_defaultdict": "Uses collections.defaultdict",
        "uses_counter": "Uses collections.Counter",
        "uses_enumerate": "Uses enumerate for index-value iteration",
        "uses_zip": "Uses zip to iterate multiple sequences",
        "uses_slicing": "Uses list/string slicing operations",
    }

    def _detect_uses_recursion(self, source: str, func_name: str) -> bool:
        try:
            tree = ast.parse(source)
            for node in ast.walk(tree):
                if isinstance(node, ast.Call):
                    if isinstance(node.func, ast.Name) and node.func.id == func_name:
                        return True
            return False
        except SyntaxError:
            return False

    def _detect_uses_iteration(self, source: str) -> bool:
        try:
            tree = ast.parse(source)
            for node in ast.walk(tree):
                if isinstance(node, (ast.For, ast.While)):
                    return True
            return False
        except SyntaxError:
            return False

    def _detect_uses_list_comprehension(self, source: str) -> bool:
        try:
            tree = ast.parse(source)
            for node in ast.walk(tree):
                if isinstance(node, (ast.ListComp, ast.DictComp, ast.SetComp)):
                    return True
            return False
        except SyntaxError:
            return False

    def _detect_uses_generator(self, source: str) -> bool:
        try:
            tree = ast.parse(source)
            for node in ast.walk(tree):
                if isinstance(node, (ast.Yield, ast.YieldFrom, ast.GeneratorExp)):
                    return True
            return False
        except SyntaxError:
            return False

    def _detect_uses_try_except(self, source: str) -> bool:
        try:
            tree = ast.parse(source)
            for node in ast.walk(tree):
                if isinstance(node, ast.Try):
                    return True
            return False
        except SyntaxError:
            return False

    def _detect_uses_assertion(self, source: str) -> bool:
        try:
            tree = ast.parse(source)
            for node in ast.walk(tree):
                if isinstance(node, ast.Assert):
                    return True
            return False
        except SyntaxError:
            return False

    def _detect_uses_early_return(self, source: str) -> bool:
        try:
            tree = ast.parse(source)
            for node in ast.walk(tree):
                if isinstance(node, ast.FunctionDef):
                    returns = [n for n in ast.walk(node) if isinstance(n, ast.Return)]
                    if len(returns) > 1:
                        return True
            return False
        except SyntaxError:
            return False

    def _detect_uses_sorting(self, source: str) -> bool:
        return "sorted(" in source or ".sort(" in source

    def _detect_uses_dict(self, source: str) -> bool:
        try:
            tree = ast.parse(source)
            for node in ast.walk(tree):
                if isinstance(node, ast.Dict):
                    return True
                if isinstance(node, ast.Call):
                    if isinstance(node.func, ast.Name) and node.func.id == "dict":
                        return True
            return False
        except SyntaxError:
            return False

    def _detect_uses_set(self, source: str) -> bool:
        try:
            tree = ast.parse(source)
            for node in ast.walk(tree):
                if isinstance(node, ast.Set):
                    return True
                if isinstance(node, ast.Call):
                    if isinstance(node.func, ast.Name) and node.func.id == "set":
                        return True
            return False
        except SyntaxError:
            return False

    def _detect_uses_heap(self, source: str) -> bool:
        return "heapq" in source or "heappush" in source or "heappop" in source

    def _detect_uses_functional(self, source: str) -> bool:
        return any(f in source for f in ["map(", "filter(", "reduce(", "functools"])

    def _detect_uses_lambda(self, source: str) -> bool:
        try:
            tree = ast.parse(source)
            for node in ast.walk(tree):
                if isinstance(node, ast.Lambda):
                    return True
            return False
        except SyntaxError:
            return False

    def _detect_uses_mutation(self, source: str) -> bool:
        """Detect in-place mutation: .append, .extend, .pop, .remove, augmented assign."""
        mutation_markers = [".append(", ".extend(", ".pop(", ".remove(",
                           ".insert(", ".update(", ".add(", ".discard("]
        if any(m in source for m in mutation_markers):
            return True
        try:
            tree = ast.parse(source)
            for node in ast.walk(tree):
                if isinstance(node, ast.AugAssign):
                    return True
            return False
        except SyntaxError:
            return False

    def _detect_single_return(self, source: str) -> bool:
        try:
            tree = ast.parse(source)
            for node in ast.walk(tree):
                if isinstance(node, ast.FunctionDef):
                    returns = [n for n in ast.walk(node) if isinstance(n, ast.Return)]
                    return len(returns) == 1
            return False
        except SyntaxError:
            return False

    def _detect_uses_defaultdict(self, source: str) -> bool:
        return "defaultdict" in source

    def _detect_uses_counter(self, source: str) -> bool:
        return "Counter(" in source

    def _detect_uses_enumerate(self, source: str) -> bool:
        return "enumerate(" in source

    def _detect_uses_zip(self, source: str) -> bool:
        return "zip(" in source

    def _detect_uses_slicing(self, source: str) -> bool:
        try:
            tree = ast.parse(source)
            for node in ast.walk(tree):
                if isinstance(node, ast.Subscript):
                    if isinstance(node.slice, ast.Slice):
                        return True
            return False
        except SyntaxError:
            return False

    def discover(self, space: ProgramSpace) -> ConceptSet:
        """
        Discover behavioral concepts from the program space.
        Uses AST analysis + execution trace analysis.
        """
        concept_set = ConceptSet()
        programs = space.valid_programs if space.valid_programs else space.programs
        func_name = space.stub.name

        # --- AST-based behavioral concepts ---
        detector_methods = {
            "uses_recursion": lambda s: self._detect_uses_recursion(s, func_name),
            "uses_iteration": self._detect_uses_iteration,
            "uses_list_comprehension": self._detect_uses_list_comprehension,
            "uses_generator": self._detect_uses_generator,
            "uses_try_except": self._detect_uses_try_except,
            "uses_assertion": self._detect_uses_assertion,
            "uses_early_return": self._detect_uses_early_return,
            "uses_sorting": self._detect_uses_sorting,
            "uses_dict": self._detect_uses_dict,
            "uses_set": self._detect_uses_set,
            "uses_heap": self._detect_uses_heap,
            "uses_functional": self._detect_uses_functional,
            "uses_lambda": self._detect_uses_lambda,
            "uses_mutation": self._detect_uses_mutation,
            "single_return": self._detect_single_return,
            "uses_defaultdict": self._detect_uses_defaultdict,
            "uses_counter": self._detect_uses_counter,
            "uses_enumerate": self._detect_uses_enumerate,
            "uses_zip": self._detect_uses_zip,
            "uses_slicing": self._detect_uses_slicing,
        }

        for concept_name, detect_fn in detector_methods.items():
            description = self.SYNTACTIC_DETECTORS[concept_name]
            positive_programs: set[str] = set()

            for program in programs:
                source = program.full_source
                try:
                    if detect_fn(source):
                        positive_programs.add(program.program_id)
                except Exception:
                    continue

            # Only create concept if it's discriminative (not all or nothing)
            if 0 < len(positive_programs) < len(programs):
                def make_detector(pids):
                    def detector(p: Program) -> float:
                        return 1.0 if p.program_id in pids else 0.0
                    return detector

                concept = Concept(
                    name=concept_name,
                    description=description,
                    type="behavioral",
                    detector=make_detector(positive_programs),
                    programs=positive_programs,
                    metadata={
                        "coverage": len(positive_programs) / len(programs),
                        "n_positive": len(positive_programs),
                        "n_total": len(programs),
                    },
                )
                concept_set.add(concept)

        # --- Execution-based behavioral concepts ---
        if programs and programs[0].traces:
            # Time complexity proxy: fast vs slow programs
            exec_times = []
            for p in programs:
                avg_time = np.mean([t.execution_time_ms for t in p.traces if t.succeeded])
                exec_times.append((p.program_id, avg_time))

            if exec_times:
                times = [t for _, t in exec_times]
                median_time = np.median(times)

                fast_programs = {pid for pid, t in exec_times if t < median_time}
                slow_programs = {pid for pid, t in exec_times if t >= median_time}

                if fast_programs and slow_programs:
                    concept_set.add(Concept(
                        name="fast_execution",
                        description="Executes faster than median (lower time complexity or better constants)",
                        type="behavioral",
                        programs=fast_programs,
                        metadata={"median_time_ms": float(median_time)},
                    ))

        logger.info(f"Discovered {len(concept_set.concepts)} behavioral concepts")
        return concept_set


class SAEConceptDiscovery(ConceptDiscovery):
    """
    Discover concepts via Sparse Autoencoder on LLM hidden states.
    Based on CB-SAE (2512.10805) and DN-CBM (2407.14499).

    Pipeline:
      1. Pass programs through a code LLM, extract hidden states
      2. Train SAE: z = TopK(E_sae(v - b)), v_hat = D_sae @ z + b
      3. Auto-name neurons by cosine similarity with concept vocabulary
      4. Prune neurons with low interpretability/steerability
    """

    def __init__(
        self,
        model_name: str = "deepseek-ai/deepseek-coder-1.3b-instruct",
        layer: int = 12,
        expansion_factor: int = 32,
        top_k: int = 32,
        n_training_steps: int = 5000,
        learning_rate: float = 3e-4,
        device: str = "auto",
    ):
        self.model_name = model_name
        self.layer = layer
        self.expansion_factor = expansion_factor
        self.top_k = top_k
        self.n_training_steps = n_training_steps
        self.learning_rate = learning_rate
        self.device = device

        # SAE parameters (initialized during training)
        self.encoder = None  # E_sae ∈ R^{ω×d}
        self.decoder = None  # D_sae ∈ R^{d×ω}
        self.bias = None     # b ∈ R^d
        self.concept_names: dict[int, str] = {}
        self.concept_scores: dict[int, float] = {}  # interpretability score

    # Default concept vocabulary for code
    CODE_CONCEPT_VOCABULARY = [
        "recursion", "iteration", "memoization", "dynamic_programming",
        "sorting", "searching", "hashing", "tree_traversal", "graph_search",
        "backtracking", "greedy", "divide_and_conquer", "two_pointers",
        "sliding_window", "binary_search", "stack_based", "queue_based",
        "heap_based", "linked_list", "array_manipulation",
        "string_processing", "bit_manipulation", "mathematical_formula",
        "brute_force", "optimization", "caching", "lazy_evaluation",
        "eager_evaluation", "immutable_operations", "mutable_operations",
        "error_handling", "input_validation", "edge_case_handling",
        "functional_style", "imperative_style", "object_oriented",
        "generator_pattern", "accumulator_pattern", "map_reduce_pattern",
        "fold_pattern", "filter_pattern", "pipeline_pattern",
        "early_termination", "exhaustive_search", "probabilistic",
        "space_efficient", "time_efficient", "readable_code",
        "compact_code", "verbose_code", "pythonic_idiom",
    ]

    def extract_hidden_states(
        self, programs: list[Program]
    ) -> Optional[np.ndarray]:
        """
        Extract hidden states from a code LLM for each program.
        Returns shape (n_programs, hidden_dim).
        """
        try:
            import torch
            from transformers import AutoTokenizer, AutoModel
        except ImportError:
            logger.error("transformers and torch required for SAE concept discovery")
            return None

        logger.info(f"Loading model {self.model_name} for hidden state extraction...")
        tokenizer = AutoTokenizer.from_pretrained(self.model_name, trust_remote_code=True)
        model = AutoModel.from_pretrained(
            self.model_name, trust_remote_code=True,
            output_hidden_states=True,
        )
        if self.device == "auto":
            model = model.to("cuda" if torch.cuda.is_available() else "cpu")
        else:
            model = model.to(self.device)
        model.eval()

        hidden_states = []
        for program in programs:
            inputs = tokenizer(
                program.full_source,
                return_tensors="pt",
                truncation=True,
                max_length=512,
            ).to(model.device)

            with torch.no_grad():
                outputs = model(**inputs)

            # Extract hidden state from target layer, mean-pool over tokens
            layer_output = outputs.hidden_states[self.layer]
            pooled = layer_output.mean(dim=1).cpu().numpy()
            hidden_states.append(pooled[0])

        return np.array(hidden_states)

    def train_sae(self, hidden_states: np.ndarray) -> None:
        """
        Train Sparse Autoencoder on hidden states.
        SAE objective: ||v - v_hat||_2^2 + λ||z||_1
        Using TopK activation (2512.10805 recommends TopK over L1).
        """
        try:
            import torch
            import torch.nn as nn
            import torch.optim as optim
        except ImportError:
            logger.error("torch required for SAE training")
            return

        d = hidden_states.shape[1]
        omega = d * self.expansion_factor

        # Initialize SAE parameters
        device = "cuda" if torch.cuda.is_available() else "cpu"
        encoder_w = torch.randn(omega, d, device=device) * 0.01
        decoder_w = torch.randn(d, omega, device=device) * 0.01
        bias = torch.zeros(d, device=device)

        encoder_w.requires_grad_(True)
        decoder_w.requires_grad_(True)
        bias.requires_grad_(True)

        optimizer = optim.Adam([encoder_w, decoder_w, bias], lr=self.learning_rate)

        data = torch.tensor(hidden_states, dtype=torch.float32, device=device)
        n = data.shape[0]

        logger.info(f"Training SAE: d={d}, ω={omega}, TopK={self.top_k}")

        for step in range(self.n_training_steps):
            # Sample batch
            idx = torch.randint(0, n, (min(64, n),), device=device)
            v = data[idx]

            # Encode: z = TopK(E_sae(v - b))
            pre_act = (v - bias) @ encoder_w.T  # (batch, omega)
            # TopK activation
            topk_vals, topk_idx = torch.topk(pre_act, self.top_k, dim=-1)
            z = torch.zeros_like(pre_act)
            z.scatter_(1, topk_idx, topk_vals.clamp(min=0))

            # Decode: v_hat = D_sae @ z + b
            v_hat = z @ decoder_w.T + bias

            # Loss: reconstruction
            loss = ((v - v_hat) ** 2).mean()

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if step % 1000 == 0:
                logger.info(f"SAE step {step}: loss={loss.item():.6f}")

        # Store trained parameters
        self.encoder = encoder_w.detach().cpu().numpy()
        self.decoder = decoder_w.detach().cpu().numpy()
        self.bias = bias.detach().cpu().numpy()

        logger.info("SAE training complete")

    def auto_name_neurons(
        self,
        concept_vocabulary: Optional[list[str]] = None,
    ) -> dict[int, str]:
        """
        Auto-name SAE neurons using cosine similarity with concept vocabulary.
        Based on CLIP-Dissect approach from CB-SAE (2512.10805).
        """
        if self.decoder is None:
            logger.error("SAE not trained yet")
            return {}

        vocab = concept_vocabulary or self.CODE_CONCEPT_VOCABULARY

        try:
            from sentence_transformers import SentenceTransformer
        except ImportError:
            logger.warning(
                "sentence-transformers not available; using hash-based naming"
            )
            # Fallback: name by neuron index
            omega = self.decoder.shape[1]
            return {i: f"concept_{i}" for i in range(omega)}

        # Embed concept vocabulary
        embedder = SentenceTransformer("all-MiniLM-L6-v2")
        vocab_embeddings = embedder.encode(vocab)
        vocab_embeddings = vocab_embeddings / np.linalg.norm(
            vocab_embeddings, axis=1, keepdims=True
        )

        # For each SAE neuron, compute cosine similarity with vocab
        # Decoder columns = concept directions in activation space
        decoder_cols = self.decoder.T  # (omega, d)
        # Project decoder columns to same dimension as vocab embeddings
        # Use PCA or random projection if dimensions don't match
        if decoder_cols.shape[1] != vocab_embeddings.shape[1]:
            # Simple approach: encode the neuron's "meaning" via its top-activating pattern
            logger.info("Dimension mismatch; using neuron activation patterns for naming")
            omega = decoder_cols.shape[0]
            self.concept_names = {i: f"sae_concept_{i}" for i in range(omega)}
            return self.concept_names

        decoder_normed = decoder_cols / (
            np.linalg.norm(decoder_cols, axis=1, keepdims=True) + 1e-8
        )
        similarities = decoder_normed @ vocab_embeddings.T  # (omega, |vocab|)

        self.concept_names = {}
        self.concept_scores = {}
        for i in range(similarities.shape[0]):
            best_idx = similarities[i].argmax()
            best_score = similarities[i, best_idx]
            self.concept_names[i] = vocab[best_idx]
            self.concept_scores[i] = float(best_score)

        return self.concept_names

    def discover(self, space: ProgramSpace) -> ConceptSet:
        """Full SAE-based concept discovery pipeline."""
        programs = space.valid_programs if space.valid_programs else space.programs
        if not programs:
            return ConceptSet()

        # Step 1: Extract hidden states
        hidden_states = self.extract_hidden_states(programs)
        if hidden_states is None:
            return ConceptSet()

        # Step 2: Train SAE
        self.train_sae(hidden_states)

        # Step 3: Auto-name neurons
        self.auto_name_neurons()

        # Step 4: Get SAE activations for all programs
        if self.encoder is not None and self.bias is not None:
            v = hidden_states - self.bias
            pre_act = v @ self.encoder.T
            # TopK
            activations = np.zeros_like(pre_act)
            for i in range(pre_act.shape[0]):
                topk_idx = np.argsort(pre_act[i])[-self.top_k:]
                activations[i, topk_idx] = np.maximum(pre_act[i, topk_idx], 0)

        # Step 5: Build concepts from top neurons
        concept_set = ConceptSet()
        omega = activations.shape[1]

        # Find discriminative neurons (high variance in activation across programs)
        neuron_variance = activations.var(axis=0)
        top_neurons = np.argsort(neuron_variance)[-50:]  # Top 50 most discriminative

        for neuron_idx in top_neurons:
            name = self.concept_names.get(neuron_idx, f"sae_concept_{neuron_idx}")
            neuron_acts = activations[:, neuron_idx]
            median_act = np.median(neuron_acts)

            # Programs where this neuron is active above median
            active_programs = set()
            for i, program in enumerate(programs):
                if neuron_acts[i] > median_act:
                    active_programs.add(program.program_id)

            if 0 < len(active_programs) < len(programs):
                concept = Concept(
                    name=f"{name}",
                    description=f"SAE neuron {neuron_idx}: represents '{name}' pattern",
                    type="representational",
                    programs=active_programs,
                    vector=self.decoder[:, neuron_idx] if self.decoder is not None else None,
                    metadata={
                        "neuron_idx": int(neuron_idx),
                        "variance": float(neuron_variance[neuron_idx]),
                        "interpretability_score": self.concept_scores.get(neuron_idx, 0.0),
                    },
                )
                concept_set.add(concept)

        logger.info(
            f"Discovered {len(concept_set.concepts)} SAE-based concepts "
            f"from {omega} neurons"
        )
        return concept_set


class AbstractionConceptDiscovery(ConceptDiscovery):
    """
    Discover concepts as shared program fragments (abstractions).
    Based on LILO (2310.19791) / ReGAL (2401.16467).

    Pipeline:
      1. Parse all programs into ASTs
      2. Find common subtree patterns (frequent subgraph mining)
      3. Name patterns via LLM (AutoDoc)
      4. Each named pattern = a structural concept
    """

    def __init__(self, min_frequency: float = 0.1, max_frequency: float = 0.9):
        """
        Args:
            min_frequency: Minimum fraction of programs that must contain the pattern
            max_frequency: Maximum fraction (patterns in all programs aren't discriminative)
        """
        self.min_frequency = min_frequency
        self.max_frequency = max_frequency

    def _extract_ast_patterns(self, source: str) -> list[str]:
        """Extract structural patterns from AST."""
        patterns = []
        try:
            tree = ast.parse(source)
        except SyntaxError:
            return patterns

        for node in ast.walk(tree):
            # Pattern: control flow structure
            if isinstance(node, ast.For):
                # Check if it's a for-in-range, for-in-enumerate, etc.
                if isinstance(node.iter, ast.Call):
                    if isinstance(node.iter.func, ast.Name):
                        patterns.append(f"for_{node.iter.func.id}")
                    elif isinstance(node.iter.func, ast.Attribute):
                        patterns.append(f"for_{node.iter.func.attr}")
                else:
                    patterns.append("for_iter")

            elif isinstance(node, ast.While):
                patterns.append("while_loop")

            elif isinstance(node, ast.If):
                # Nested if depth
                depth = 0
                n = node
                while isinstance(n, ast.If) and n.orelse:
                    if len(n.orelse) == 1 and isinstance(n.orelse[0], ast.If):
                        depth += 1
                        n = n.orelse[0]
                    else:
                        break
                patterns.append(f"if_chain_{depth}" if depth > 0 else "if_simple")

            elif isinstance(node, ast.FunctionDef):
                # Nested function definition
                if any(isinstance(child, ast.FunctionDef) for child in ast.walk(node)):
                    inner_fns = [
                        child for child in ast.walk(node)
                        if isinstance(child, ast.FunctionDef) and child is not node
                    ]
                    if inner_fns:
                        patterns.append("nested_function")

            elif isinstance(node, ast.ListComp):
                patterns.append("list_comprehension")
                # Nested comprehension
                if len(node.generators) > 1:
                    patterns.append("nested_comprehension")

            elif isinstance(node, ast.DictComp):
                patterns.append("dict_comprehension")

            elif isinstance(node, ast.Call):
                if isinstance(node.func, ast.Name):
                    patterns.append(f"call_{node.func.id}")
                elif isinstance(node.func, ast.Attribute):
                    patterns.append(f"method_{node.func.attr}")

        return patterns

    def discover(self, space: ProgramSpace) -> ConceptSet:
        """Discover structural abstraction concepts."""
        programs = space.valid_programs if space.valid_programs else space.programs
        if not programs:
            return ConceptSet()

        n = len(programs)
        min_count = max(1, int(n * self.min_frequency))
        max_count = int(n * self.max_frequency)

        # Extract patterns from all programs
        program_patterns: dict[str, list[str]] = {}
        pattern_counter: Counter = Counter()

        for program in programs:
            patterns = self._extract_ast_patterns(program.full_source)
            program_patterns[program.program_id] = patterns
            pattern_counter.update(set(patterns))  # Count each pattern once per program

        # Find discriminative patterns
        concept_set = ConceptSet()

        # Human-readable descriptions for common patterns
        pattern_descriptions = {
            "for_range": "Uses for-range loop (index-based iteration)",
            "for_enumerate": "Uses enumerate for index-value pairs",
            "for_zip": "Uses zip to iterate multiple sequences in parallel",
            "for_iter": "Uses direct iterable iteration",
            "while_loop": "Uses while loop (condition-based iteration)",
            "if_simple": "Uses simple conditional branching",
            "if_chain_1": "Uses if-elif chain (multiple conditions)",
            "if_chain_2": "Uses deeply nested if-elif chains",
            "nested_function": "Defines helper function(s) inside main function",
            "list_comprehension": "Uses list comprehension for concise collection building",
            "dict_comprehension": "Uses dict comprehension",
            "nested_comprehension": "Uses nested comprehension (multiple generators)",
            "call_sorted": "Calls sorted() built-in",
            "call_len": "Calls len() built-in",
            "call_sum": "Calls sum() built-in",
            "call_max": "Calls max() built-in",
            "call_min": "Calls min() built-in",
            "method_append": "Uses list.append() for accumulation",
            "method_items": "Uses dict.items() for key-value iteration",
            "method_keys": "Uses dict.keys()",
            "method_values": "Uses dict.values()",
            "method_join": "Uses str.join() for string building",
            "method_split": "Uses str.split() for tokenization",
        }

        for pattern, count in pattern_counter.items():
            if min_count <= count <= max_count:
                # Find programs that have this pattern
                positive_programs = {
                    pid
                    for pid, patterns in program_patterns.items()
                    if pattern in patterns
                }

                description = pattern_descriptions.get(
                    pattern,
                    f"Uses {pattern.replace('_', ' ')} pattern",
                )

                concept = Concept(
                    name=f"struct_{pattern}",
                    description=description,
                    type="structural",
                    programs=positive_programs,
                    metadata={
                        "frequency": count / n,
                        "n_programs": count,
                        "pattern": pattern,
                    },
                )
                concept_set.add(concept)

        logger.info(f"Discovered {len(concept_set.concepts)} structural concepts")
        return concept_set


class UnifiedConceptDiscovery:
    """
    Combines all three concept discovery methods into a single pipeline.
    Merges behavioral, representational, and structural concepts.
    """

    def __init__(
        self,
        behavioral: Optional[BehavioralConceptDiscovery] = None,
        sae: Optional[SAEConceptDiscovery] = None,
        abstraction: Optional[AbstractionConceptDiscovery] = None,
    ):
        self.behavioral = behavioral or BehavioralConceptDiscovery()
        self.sae = sae  # Optional: requires GPU
        self.abstraction = abstraction or AbstractionConceptDiscovery()

    def discover(self, space: ProgramSpace) -> ConceptSet:
        """Run all discovery methods and merge results."""
        merged = ConceptSet()

        # Always run behavioral discovery
        logger.info("Running behavioral concept discovery...")
        behavioral_concepts = self.behavioral.discover(space)
        for c in behavioral_concepts.concepts:
            merged.add(c)

        # Run SAE discovery if available
        if self.sae is not None:
            logger.info("Running SAE concept discovery...")
            try:
                sae_concepts = self.sae.discover(space)
                for c in sae_concepts.concepts:
                    merged.add(c)
            except Exception as e:
                logger.warning(f"SAE discovery failed: {e}")

        # Always run abstraction discovery
        logger.info("Running abstraction concept discovery...")
        abstraction_concepts = self.abstraction.discover(space)
        for c in abstraction_concepts.concepts:
            merged.add(c)

        logger.info(
            f"Unified discovery: {len(merged.concepts)} total concepts "
            f"({len(behavioral_concepts.concepts)} behavioral, "
            f"{len(abstraction_concepts.concepts)} structural)"
        )

        return merged