layer_data.py

"""Layer metadata and neuron detail loading with on-demand caching."""

import json
import os
from typing import Any, Dict, List, Optional, Tuple

from datasets import get_dataset_split_names, load_dataset

from config import QWEN25_DATASET_ROOT, REPO_ID, qwen25_dataset_path

# Config names and directories in the HF dataset
CONFIG_META = f"{QWEN25_DATASET_ROOT}_meta_data"
CONFIG_DETAILS = f"{QWEN25_DATASET_ROOT}_neuron_details"

# Global cache for per-layer meta dict (lightweight)
layer_meta_cache: Dict[int, Dict[int, Dict]] = {}

# Cache for per-layer, per-chunk neuron detail datasets (non-streaming)
# Keyed by layer_id -> chunk_start (every 1000 neurons) -> dataset split object
layer_datasets: Dict[int, Dict[int, Any]] = {}

# Cache for per-neuron detailed rows (only selected neuron)
neuron_detail_cache: Dict[Tuple[int, int], Dict] = {}


def _summarize_splits(split_names: List[str], head: int = 5, tail: int = 5) -> str:
    if not split_names:
        return "[]"
    if len(split_names) <= head + tail:
        return str(split_names)
    return str(split_names[:head] + ["..."] + split_names[-tail:])


def load_layer_options() -> List[int]:
    """Fetch available layer IDs, prefer local JSON, fallback to default."""
    try:
        if os.path.exists("resources/layer_list.json"):
            with open("resources/layer_list.json", "r", encoding="utf-8") as f:
                layer_options = json.load(f)
            layer_options = sorted(int(x) for x in layer_options)
        else:
            layer_options = [20]
        print(f"[info] Available layer IDs: {layer_options}")
        return layer_options
    except Exception as e:  # pragma: no cover - small utility
        print(f"[error] Failed to load layer options: {e}")
        return [20]


def load_layer_data(layer_id: int) -> Dict[int, Dict]:
    """Load lightweight meta data for a specific layer from HF dataset.

    Uses a single-file data_files load for layer meta parquet.
    Only minimal fields are stored; heavy refs are not loaded here.
    """
    if layer_id in layer_meta_cache:
        cached = layer_meta_cache[layer_id]
        print(f"[info] Using cached meta_data for layer {layer_id} (size={len(cached)})")
        return cached

    try:
        data_file = qwen25_dataset_path(
            f"meta_data/layer_{layer_id}-00000-of-00001.parquet"
        )
        print(f"[info] Loading meta_data via data_files: layer={layer_id} file={data_file}")
        ds = load_dataset(
            REPO_ID,
            data_files={"train": [data_file]},
            split="train",
            streaming=False,
        )
        print(f"[info] meta_data loaded: file={data_file} rows={len(ds)}")

        meta_fields = {
            "neuron_id",
            "modality",
            "vision_summary",
            "language_summary",
            "consistency_score",
            "miou",
        }
        meta_dict: Dict[int, Dict] = {}

        for item in ds:
            nid = item.get("neuron_id")
            if nid is None:
                continue
            try:
                meta = {k: item.get(k) for k in meta_fields}
                meta_dict[int(nid)] = meta
            except Exception as row_err:  # pragma: no cover
                print(f"[warn] Skipping meta row for layer {layer_id}: {row_err}")

        layer_meta_cache[layer_id] = meta_dict
        print(f"[info] Loaded meta_data for {len(meta_dict)} neurons for layer {layer_id}")
        return meta_dict
    except Exception as e:  # pragma: no cover - network/cache dependent
        print(f"[error] Failed to load meta_data for layer {layer_id}: {e}")
        return {}


def _details_split_name(layer_id: int, neuron_id: int, chunk: int = 1000) -> str:
    start = (int(neuron_id) // chunk) * chunk
    end = start + (chunk - 1)
    return f"layer_{layer_id}_neuron_{start:06d}_{end:06d}"


def load_neuron_detail(layer_id: int, neuron_id: int) -> Optional[Dict]:
    """Load detailed refs (vision/text) for the selected neuron only.

    The dataset index now aligns with the neuron index, including empty neurons,
    so we can directly index into the split without an extra scan on `neuron_id`.
    """
    cache_key = (int(layer_id), int(neuron_id))
    if cache_key in neuron_detail_cache:
        print(f"[info] Using cached neuron_details for neuron {neuron_id} layer {layer_id}")
        return neuron_detail_cache[cache_key]

    # Ensure meta exists (for sanity), but don't depend on its split info
    meta_dict = load_layer_data(int(layer_id))
    if int(neuron_id) not in meta_dict:
        return None

    chunk_size = 1000
    split_name = _details_split_name(int(layer_id), int(neuron_id), chunk=chunk_size)
    print(f"[info] Loading neuron_details: layer={layer_id} neuron={neuron_id} split={split_name}")
    split_names = None
    try:
        split_names = get_dataset_split_names(path=REPO_ID, config_name=CONFIG_DETAILS)
        print(
            f"[info] neuron_details splits via config='{CONFIG_DETAILS}': count={len(split_names)} "
            f"sample={_summarize_splits(split_names)}"
        )
    except Exception as e:
        print(
            f"[warn] Failed to list neuron_details splits via config: {e}; "
            "trying the expected file directly..."
        )

    if split_names is not None and split_name not in split_names:
        print(f"[warn] neuron_details split '{split_name}' not found for layer {layer_id}")
        return None

    try:
        chunk_start = (int(neuron_id) // chunk_size) * chunk_size
        layer_cache = layer_datasets.setdefault(int(layer_id), {})

        if chunk_start in layer_cache:
            ds = layer_cache[chunk_start]
            print(
                f"[info] Using cached neuron_details split='{split_name}' for layer {layer_id} "
                f"chunk_start={chunk_start}"
            )
        else:
            data_file = qwen25_dataset_path(
                f"neuron_details/layer_{layer_id}/{split_name}-00000-of-00001.parquet"
            )
            print(
                f"[info] Loading neuron_details via data_files: layer={layer_id} split={split_name} "
                f"file={data_file}"
            )
            ds = load_dataset(
                REPO_ID,
                data_files={"train": [data_file]},
                split="train",
                streaming=False,
            )
            layer_cache[chunk_start] = ds
            print(f"[info] Cached split='{split_name}' rows={len(ds)} for layer {layer_id}")

        # Dataset index matches neuron index; empty neurons yield empty rows.
        local_index = int(neuron_id) - chunk_start
        if local_index < 0 or local_index >= len(ds):
            print(
                f"[warn] Neuron {neuron_id} maps to local index {local_index}, "
                f"but split '{split_name}' has length {len(ds)}"
            )
            return None

        row = ds[int(local_index)]
        neuron_detail_cache[cache_key] = row
        print(f"[info] Loaded neuron_details for neuron {neuron_id} layer {layer_id} via direct index")
        return row
    except Exception as e:  # pragma: no cover - network/cache dependent
        print(f"[error] Failed to load neuron_details for neuron {neuron_id} layer {layer_id}: {e}")
        return None


def get_neuron_ids(layer_id: int) -> List[int]:
    """Return all neuron IDs for the given layer from meta_data."""
    data = load_layer_data(layer_id)
    if data:
        return sorted(data.keys())
    return []


def get_consistency_neurons(layer_id: int, max_examples: int = 25) -> List[int]:
    """Return up to max_examples neurons with consistency_score == 1 from meta_data."""
    data = load_layer_data(layer_id)
    if not data:
        return []
    matches = [nid for nid, meta in data.items() if meta.get("consistency_score") == 1.0]
    return sorted(matches)[:max_examples]