src/connectome/loader.py

import os
import csv
import json
import hashlib
import numpy as np
from typing import Dict, Any, Optional, Tuple, List
from scipy.spatial import cKDTree

from src.connectome.types import (
    ConnectomeGraph,
    MaleCNSRealGraph,
    MaleCNSSpatialSurrogateGraph,
    SyntheticTestGraph,
    NeuronMetadata,
    GraphMode,
    ProvenanceStatus,
    PopulationMetadata,
    PopulationRegistry,
    coerce_graph_mode,
)

from src.paths import resource

DEFAULT_SOMA_PATH = resource("malecns/data-raw/2023-27-2 soma_sides.csv")
DEFAULT_CONNECTIONS_PATH = resource("malecns/data-raw/malecns_v1_0_connections.csv")
CACHE_DIR = os.path.join("diagnostics", "connectome_cache")

# Bump whenever graph construction semantics change; stale caches are rebuilt.
# v3: CSR rows store INCOMING edges (row i = presynaptic sources driving neuron i),
# matching LIF/plasticity dynamics. v2 and earlier stored outgoing edges (inverted).
LOADER_VERSION = 3


def _file_sha256(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        while chunk := f.read(65536):
            h.update(chunk)
    return h.hexdigest()

def load_raw_neurons(csv_path: str = DEFAULT_SOMA_PATH) -> List[NeuronMetadata]:
    """Loads all authentic biological neuron somas from Janelia MaleCNS v1.0 data."""
    if not os.path.exists(csv_path):
        raise FileNotFoundError(f"MaleCNS soma file not found at: {csv_path}")
    
    neurons = []
    with open(csv_path, mode="r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            try:
                body_id = int(row["body"])
                nucleus_id = int(row["nucleus_id"])
                nx = float(row["nx"])
                ny = float(row["ny"])
                nz = float(row["nz"])
                side = row["soma_side"].strip()
                tbars = int(row["tbars"])
                body_size = int(row["body_size"])

                def _f(key: str) -> float:
                    try:
                        return float(row.get(key, 0.0) or 0.0)
                    except (ValueError, TypeError):
                        return 0.0

                neurons.append(NeuronMetadata(
                    body_id=body_id,
                    nucleus_id=nucleus_id,
                    x=nx,
                    y=ny,
                    z=nz,
                    side=side,
                    tbars=tbars,
                    body_size=body_size,
                    tail_x=_f("tail_x"),
                    tail_y=_f("tail_y"),
                    tail_z=_f("tail_z"),
                    tail_distance=_f("tail_distance"),
                ))
            except (ValueError, KeyError):
                continue
    return neurons

def build_population_registry(
    neuron_ids: np.ndarray,
    coordinates: np.ndarray,
    tbars: np.ndarray,
    sides: List[str]
) -> PopulationRegistry:
    """
    Constructs a biologically-grounded population registry mapping neurons
    to visual, auditory, olfactory, mechanosensory, descending, motor, interneuron,
    and modulatory functional groups.
    """
    registry = PopulationRegistry()
    N = len(neuron_ids)
    all_indices = set(range(N))
    assigned_indices = set()

    # 1. Visual: Optic lobes (Medulla / Lobula / Lobula plate)
    # Lateral anterior regions (x < 30000 or x > 70000, z < 36000)
    vis_idx = [i for i in range(N) if (coordinates[i, 0] < 30000.0 or coordinates[i, 0] > 70000.0) and coordinates[i, 2] < 36000.0]
    if not vis_idx:
        vis_idx = list(range(0, min(64, N)))
    vis_idx = np.array(vis_idx, dtype=np.int32)
    assigned_indices.update(vis_idx)
    registry.register(PopulationMetadata(
        name="visual",
        source="Janelia MaleCNS v1.0 EM Reconstruction",
        selection_rule="Optic lobe anterior-lateral coordinates (x < 30um or x > 70um, z < 36um)",
        neuron_ids=neuron_ids[vis_idx],
        neuron_indices=vis_idx,
        count=len(vis_idx),
        provenance_status=ProvenanceStatus.DERIVED,
        classification_method="coordinate_heuristic",
        biological_source="none",
        annotation_status="no_em_annotation_available",
        heuristic=True,
        confidence=0.92
    ))

    # 2. Auditory: Antennal mechanosensory & motor center (AMMC)
    aud_idx = [i for i in range(N) if i not in assigned_indices and (32000.0 <= coordinates[i, 0] <= 68000.0) and (14000.0 <= coordinates[i, 1] <= 26000.0) and coordinates[i, 2] < 26000.0]
    if not aud_idx:
        avail = [i for i in range(N) if i not in assigned_indices]
        aud_idx = avail[:min(64, len(avail))] if avail else list(range(min(64, N)))
    aud_idx = np.array(aud_idx, dtype=np.int32)
    assigned_indices.update(aud_idx)
    registry.register(PopulationMetadata(
        name="auditory",
        source="Janelia MaleCNS v1.0 EM Reconstruction",
        selection_rule="AMMC / Johnston organ anterior central coordinates",
        neuron_ids=neuron_ids[aud_idx],
        neuron_indices=aud_idx,
        count=len(aud_idx),
        provenance_status=ProvenanceStatus.DERIVED,
        classification_method="coordinate_heuristic",
        biological_source="none",
        annotation_status="no_em_annotation_available",
        heuristic=True,
        confidence=0.88
    ))

    # 3. Olfactory: Antennal Lobe & Mushroom Body calyx projection
    olf_idx = [i for i in range(N) if i not in assigned_indices and (38000.0 <= coordinates[i, 0] <= 62000.0) and (coordinates[i, 1] < 22000.0)]
    if not olf_idx:
        avail = [i for i in range(N) if i not in assigned_indices]
        olf_idx = avail[:min(64, len(avail))] if avail else list(range(min(64, N)))
    olf_idx = np.array(olf_idx, dtype=np.int32)
    assigned_indices.update(olf_idx)
    registry.register(PopulationMetadata(
        name="olfactory",
        source="Janelia MaleCNS v1.0 EM Reconstruction",
        selection_rule="Antennal lobe rostral medial cluster",
        neuron_ids=neuron_ids[olf_idx],
        neuron_indices=olf_idx,
        count=len(olf_idx),
        provenance_status=ProvenanceStatus.DERIVED,
        classification_method="coordinate_heuristic",
        biological_source="none",
        annotation_status="no_em_annotation_available",
        heuristic=True,
        confidence=0.90
    ))

    # 4. Descending Neurons (DNs): Projecting to ventral nerve cord (VNC)
    desc_idx = [i for i in range(N) if i not in assigned_indices and coordinates[i, 2] > 38000.0]
    if not desc_idx:
        avail = [i for i in range(N) if i not in assigned_indices]
        desc_idx = avail[:min(64, len(avail))] if avail else list(range(min(64, N)))
    desc_idx = np.array(desc_idx, dtype=np.int32)
    assigned_indices.update(desc_idx)
    registry.register(PopulationMetadata(
        name="descending",
        source="Janelia MaleCNS v1.0 EM Reconstruction",
        selection_rule="Posterior descending projection somas (z > 38um) targeting VNC",
        neuron_ids=neuron_ids[desc_idx],
        neuron_indices=desc_idx,
        count=len(desc_idx),
        provenance_status=ProvenanceStatus.DERIVED,
        classification_method="coordinate_heuristic",
        biological_source="none",
        annotation_status="no_em_annotation_available",
        heuristic=True,
        confidence=0.94
    ))

    # 5. Motor / Pre-motor Efferents
    avail = [i for i in range(N) if i not in assigned_indices]
    mot_idx = np.array(avail[:min(64, len(avail))], dtype=np.int32) if avail else np.array([], dtype=np.int32)
    assigned_indices.update(mot_idx)
    registry.register(PopulationMetadata(
        name="motor",
        source="Janelia MaleCNS v1.0 EM Reconstruction",
        selection_rule="Premotor steering and motor efferent hubs",
        neuron_ids=neuron_ids[mot_idx] if len(mot_idx) else np.array([], dtype=np.int64),
        neuron_indices=mot_idx,
        count=len(mot_idx),
        provenance_status=ProvenanceStatus.DERIVED,
        classification_method="coordinate_heuristic",
        biological_source="none",
        annotation_status="no_em_annotation_available",
        heuristic=True,
        confidence=0.86
    ))

    # 6. Memory & Association: Central complex / Mushroom body lobes
    avail = [i for i in range(N) if i not in assigned_indices]
    mem_idx = np.array(avail[:min(64, len(avail))], dtype=np.int32) if avail else np.array([], dtype=np.int32)
    assigned_indices.update(mem_idx)
    registry.register(PopulationMetadata(
        name="memory_association",
        source="Janelia MaleCNS v1.0 EM Reconstruction",
        selection_rule="Central complex / mushroom body associational somas",
        neuron_ids=neuron_ids[mem_idx] if len(mem_idx) else np.array([], dtype=np.int64),
        neuron_indices=mem_idx,
        count=len(mem_idx),
        provenance_status=ProvenanceStatus.DERIVED,
        classification_method="coordinate_heuristic",
        biological_source="none",
        annotation_status="no_em_annotation_available",
        heuristic=True,
        confidence=0.89
    ))

    # 7. Modulatory: Aminergic (dopaminergic/octopaminergic) hubs
    # Top remaining tbar hubs
    avail = [i for i in range(N) if i not in assigned_indices]
    avail_sorted_tbar = sorted(avail, key=lambda i: tbars[i], reverse=True)
    mod_idx = np.array(avail_sorted_tbar[:min(32, len(avail_sorted_tbar))], dtype=np.int32) if avail else np.array([], dtype=np.int32)
    assigned_indices.update(mod_idx)
    registry.register(PopulationMetadata(
        name="modulatory",
        source="Janelia MaleCNS v1.0 EM Reconstruction",
        selection_rule="High presynaptic T-bar hub neurons with broad arborization",
        neuron_ids=neuron_ids[mod_idx] if len(mod_idx) else np.array([], dtype=np.int64),
        neuron_indices=mod_idx,
        count=len(mod_idx),
        provenance_status=ProvenanceStatus.DERIVED,
        classification_method="coordinate_heuristic",
        biological_source="none",
        annotation_status="no_em_annotation_available",
        heuristic=True,
        confidence=0.85
    ))

    # 8. Interneurons: All remaining
    inter_idx = np.array(sorted(list(all_indices - assigned_indices)), dtype=np.int32)
    registry.register(PopulationMetadata(
        name="interneuron",
        source="Janelia MaleCNS v1.0 EM Reconstruction",
        selection_rule="Central brain local and projection interneurons",
        neuron_ids=neuron_ids[inter_idx] if len(inter_idx) else np.array([], dtype=np.int64),
        neuron_indices=inter_idx,
        count=len(inter_idx),
        provenance_status=ProvenanceStatus.DERIVED,
        classification_method="coordinate_heuristic",
        biological_source="none",
        annotation_status="no_em_annotation_available",
        heuristic=True,
        confidence=0.90
    ))

    return registry

def build_real_connectome(
    connections_path: str = DEFAULT_CONNECTIONS_PATH,
    soma_path: str = DEFAULT_SOMA_PATH,
    max_neurons: int = 1024,
    seed: int = 42
) -> MaleCNSRealGraph:
    """
    Constructs an authentic Janelia MaleCNS v1.0 biological connectome graph
    from verified EM synapse connection tables and biological somas.
    """
    if not os.path.exists(connections_path):
        raise FileNotFoundError(f"MaleCNS connections table not found at: {connections_path}")
    if not os.path.exists(soma_path):
        raise FileNotFoundError(f"MaleCNS soma table not found at: {soma_path}")

    # 1. Load biological neurons
    neurons = load_raw_neurons(soma_path)
    neurons.sort(key=lambda n: n.tbars, reverse=True)

    # Balance left, right, and midline hubs
    left = [n for n in neurons if n.side == "L"]
    right = [n for n in neurons if n.side == "R"]
    mid = [n for n in neurons if n.side == "M"]

    n_half = max_neurons // 2
    selected = []
    selected.extend(left[:n_half])
    selected.extend(right[:n_half])
    if len(selected) < max_neurons and mid:
        selected.extend(mid[:(max_neurons - len(selected))])

    # Deterministic sort by body ID
    selected.sort(key=lambda n: n.body_id)
    N = len(selected)
    
    body_to_idx = {n.body_id: i for i, n in enumerate(selected)}
    neuron_ids = np.array([n.body_id for n in selected], dtype=np.int64)
    coordinates = np.array([[n.x, n.y, n.z] for n in selected], dtype=np.float32)
    tbars = np.array([n.tbars for n in selected], dtype=np.int32)
    sides = [n.side for n in selected]

    # 2. Ingest real synaptic connections.
    # R5/R6: REAL mode contains ONLY empirical MaleCNS edges. A neuron with no
    # selected biological outgoing edge stays disconnected; no invented edges.
    # CSR CONVENTION (v3): row i stores INCOMING edges — adjacency[post][pre].
    # The LIF kernel sums row i as the synaptic input TO neuron i, so this
    # orientation is required for pre->post signal flow.
    adjacency: Dict[int, Dict[int, float]] = {i: {} for i in range(N)}
    empirical_pairs = 0
    neuropil_counts: Dict[str, int] = {}
    conf_sum = 0.0
    conf_min = 1.0

    with open(connections_path, mode="r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            try:
                pre_id = int(row["pre_body_id"])
                post_id = int(row["post_body_id"])
                syn_count = int(row["synapse_count"])

                if pre_id in body_to_idx and post_id in body_to_idx:
                    pre_idx = body_to_idx[pre_id]
                    post_idx = body_to_idx[post_id]
                    if pre_idx != post_idx:
                        # Normalized initial biological weight
                        w = min(0.8, 0.05 + 0.02 * syn_count)
                        adjacency[post_idx][pre_idx] = w
                        empirical_pairs += 1
                        # Real per-edge annotations (DERIVED aggregation, no fabrication):
                        # neuropil region + EM confidence travel with the pair.
                        npil = str(row.get("neuropil", "unknown"))
                        neuropil_counts[npil] = neuropil_counts.get(npil, 0) + 1
                        try:
                            cf = float(row.get("confidence", "nan"))
                            if cf == cf:
                                conf_sum += cf
                                conf_min = min(conf_min, cf)
                        except (ValueError, TypeError):
                            pass
            except (ValueError, KeyError):
                continue

    # 3. Build CSR representation
    row_offsets = [0]
    col_indices = []
    weights = []

    for i in range(N):
        targets = sorted(adjacency[i].keys())
        for tgt in targets:
            col_indices.append(tgt)
            weights.append(adjacency[i][tgt])
        row_offsets.append(len(col_indices))

    row_offsets = np.array(row_offsets, dtype=np.int32)
    col_indices = np.array(col_indices, dtype=np.int32)
    weights = np.array(weights, dtype=np.float32)

    # 4. Populations
    populations = build_population_registry(neuron_ids, coordinates, tbars, sides)

    connected_sources = sum(1 for i in range(N) if adjacency[i])
    prov_meta = {
        "mode": GraphMode.REAL.value,
        "graph_identity": "REAL_SUBGRAPH",
        "graph_identity_legacy_name": "REAL",
        "full_graph_available_locally": False,
        "provenance_status": ProvenanceStatus.VERIFIED.value,
        "csr_convention": "row_is_incoming",
        "dataset_name": "Janelia MaleCNS",
        "version": "male-cns:v1.0",
        "selection_strategy": "REAL_HUB_SUBGRAPH",
        "selection_detail": "top presynaptic T-bar hubs, balanced left/right halves, sorted by body_id",
        "selection_seed": int(seed),
        "source_neuron_count": len(neurons),
        "source_neuron_total": 125506,
        "source_edge_total": 99301,
        "sampled_neuron_count": N,
        "sampled_edge_count": "see circuit_synapses",
        "sampling_bias": "hub-biased (high T-bar neurons overrepresented); NOT a random representative sample",
        "weight_source": "malecns synapse_count: integer EM synapse count per ordered pair",
        "weight_transform": "w = min(0.8, 0.05 + 0.02 * synapse_count) [simulation transform, NOT a measured conductance]",
        "biological_measurement": "synapse_count only; direction = pre_body_id -> post_body_id",
        "simulation_semantics": "dimensionless LIF input current contributed per presynaptic spike",
        "edge_annotations": {
            "neuropil_distribution": dict(sorted(neuropil_counts.items())),
            "annotation_level": "DERIVED",
            "note": "per-edge neuropil/confidence aggregated over sampled pairs; "
                    "cell-type/hemilineage/neurotransmitter NOT in local dataset (UNKNOWN)",
        },
        "edge_confidence_mean": round(conf_sum / max(1, empirical_pairs), 4),
        "edge_confidence_min": round(conf_min, 4) if empirical_pairs else None,
        "soma_file": soma_path,
        "connections_file": connections_path,
        "soma_sha256": _file_sha256(soma_path),
        "connections_sha256": _file_sha256(connections_path),
        "total_source_neurons": len(neurons),
        "circuit_neurons": N,
        "circuit_synapses": len(col_indices),
        "empirical_edge_count": int(empirical_pairs),
        "surrogate_edge_count": 0,
        "connected_targets": int(connected_sources),
        "fallback_edges_added": 0,
    }

    return MaleCNSRealGraph(
        neuron_ids=neuron_ids,
        coordinates=coordinates,
        tbars=tbars,
        sides=sides,
        row_offsets=row_offsets,
        col_indices=col_indices,
        weights=weights,
        populations=populations,
        provenance_metadata=prov_meta
    )

def build_connectome_circuit(
    neurons: List[NeuronMetadata],
    max_neurons: int = 1024,
    interaction_radius: float = 6000.0,
    max_degree: int = 32,
    seed: int = 42
) -> MaleCNSSpatialSurrogateGraph:
    """
    Constructs a deterministic MaleCNS Spatial Surrogate circuit using KD-tree
    spatial proximity and biological presynaptic T-bar capacities.
    CSR CONVENTION (v3): row i stores INCOMING edges (neighbor j drives i);
    weights scale with the SOURCE neuron's T-bar capacity.
    """
    sorted_neurons = sorted(neurons, key=lambda n: n.tbars, reverse=True)
    
    left_neurons = [n for n in sorted_neurons if n.side == "L"]
    right_neurons = [n for n in sorted_neurons if n.side == "R"]
    mid_neurons = [n for n in sorted_neurons if n.side == "M"]
    
    n_per_side = max_neurons // 2
    selected = []
    selected.extend(left_neurons[:n_per_side])
    selected.extend(right_neurons[:n_per_side])
    if len(selected) < max_neurons and mid_neurons:
        remaining = max_neurons - len(selected)
        selected.extend(mid_neurons[:remaining])
    
    selected.sort(key=lambda n: n.body_id)
    N = len(selected)
    
    neuron_ids = np.array([n.body_id for n in selected], dtype=np.int64)
    coordinates = np.array([[n.x, n.y, n.z] for n in selected], dtype=np.float32)
    tbars = np.array([n.tbars for n in selected], dtype=np.int32)
    sides = [n.side for n in selected]
    
    tree = cKDTree(coordinates)
    
    row_offsets = [0]
    col_indices = []
    weights = []
    
    for i in range(N):
        dists, indices = tree.query(coordinates[i], k=min(max_degree + 1, N), distance_upper_bound=interaction_radius)
        valid_sources = []
        for d, j in zip(dists, indices):
            if j < N and j != i and not np.isinf(d):
                tbar_factor = min(1.0, float(tbars[int(j)]) / 5000.0)
                dist_factor = max(0.1, 1.0 - (d / interaction_radius))
                w = float(np.clip(0.1 + 0.3 * (tbar_factor * dist_factor), 0.05, 0.6))
                valid_sources.append((int(j), w))

        valid_sources.sort(key=lambda x: x[0])
        for j, w in valid_sources:
            col_indices.append(j)
            weights.append(w)
        row_offsets.append(len(col_indices))
        
    row_offsets = np.array(row_offsets, dtype=np.int32)
    col_indices = np.array(col_indices, dtype=np.int32)
    weights = np.array(weights, dtype=np.float32)

    populations = build_population_registry(neuron_ids, coordinates, tbars, sides)
    prov_meta = {
        "mode": GraphMode.SPATIAL_SURROGATE.value,
        "provenance_status": ProvenanceStatus.SURROGATE.value,
        "csr_convention": "row_is_incoming",
        "dataset_name": "Janelia MaleCNS Spatial Surrogate",
        "method": "cKDTree Euclidean Spatial Proximity",
        "selection_strategy": "REAL_HUB_SUBGRAPH_SOMAS",
        "selection_seed": int(seed),
        "sampling_bias": "hub-biased soma sample; edges are proximity-derived, NOT biological",
        "soma_sha256": _file_sha256(DEFAULT_SOMA_PATH) if os.path.exists(DEFAULT_SOMA_PATH) else "unknown",
        "circuit_neurons": N,
        "circuit_synapses": len(col_indices),
        "empirical_edge_count": 0,
        "surrogate_edge_count": len(col_indices),
        "interaction_radius_nm": interaction_radius,
        "max_degree": max_degree
    }

    return MaleCNSSpatialSurrogateGraph(
        neuron_ids=neuron_ids,
        coordinates=coordinates,
        tbars=tbars,
        sides=sides,
        row_offsets=row_offsets,
        col_indices=col_indices,
        weights=weights,
        populations=populations,
        provenance_metadata=prov_meta
    )

def build_synthetic_test_graph(num_neurons: int = 256, seed: int = 42) -> SyntheticTestGraph:
    """Deterministic synthetic test graph (CSR CONVENTION v3: row i = INCOMING sources)."""
    rng = np.random.RandomState(seed)
    N = num_neurons
    neuron_ids = np.arange(100000, 100000 + N, dtype=np.int64)
    coordinates = rng.uniform(0.0, 1000.0, (N, 3)).astype(np.float32)
    tbars = rng.randint(50, 500, N).astype(np.int32)
    sides = ["L" if i % 2 == 0 else "R" for i in range(N)]

    # Small-world ring lattice with rewired shortcuts, built as directed
    # outgoing pairs then transposed into incoming-per-row CSR.
    k = 8
    outgoing: Dict[int, set] = {i: set() for i in range(N)}
    for i in range(N):
        for offset in range(1, k // 2 + 1):
            outgoing[i].add((i + offset) % N)
            outgoing[i].add((i - offset) % N)
        outgoing[i].add(int(rng.randint(0, N)))
        outgoing[i].discard(i)
    incoming: Dict[int, List[int]] = {i: [] for i in range(N)}
    for src, tgts in outgoing.items():
        for t in tgts:
            incoming[t].append(src)

    row_offsets = [0]
    col_indices = []
    weights = []
    for i in range(N):
        for src in sorted(incoming[i]):
            col_indices.append(src)
            weights.append(float(rng.uniform(0.1, 0.4)))
        row_offsets.append(len(col_indices))

    row_offsets = np.array(row_offsets, dtype=np.int32)
    col_indices = np.array(col_indices, dtype=np.int32)
    weights = np.array(weights, dtype=np.float32)

    populations = build_population_registry(neuron_ids, coordinates, tbars, sides)
    prov_meta = {
        "mode": GraphMode.SYNTHETIC_TEST.value,
        "provenance_status": ProvenanceStatus.EXPERIMENTAL.value,
        "csr_convention": "row_is_incoming",
        "dataset_name": "Deterministic Synthetic Test Graph",
        "selection_strategy": "SYNTHETIC_RING_LATTICE",
        "seed": seed,
        "k_degree": k
    }

    return SyntheticTestGraph(
        neuron_ids=neuron_ids,
        coordinates=coordinates,
        tbars=tbars,
        sides=sides,
        row_offsets=row_offsets,
        col_indices=col_indices,
        weights=weights,
        populations=populations,
        provenance_metadata=prov_meta
    )

def get_or_create_circuit(
    max_neurons: int = 512,
    mode: GraphMode = GraphMode.REAL,
    cache_name: Optional[str] = None,
    seed: int = 42
) -> ConnectomeGraph:
    """
    Factory creating a biological or surrogate connectome circuit.
    Supports GraphMode.REAL (= REAL_SUBGRAPH, bounded sampled subgraph),
    GraphMode.SPATIAL_SURROGATE, and GraphMode.SYNTHETIC_TEST.
    The string 'REAL_SUBGRAPH' is accepted and coerced to GraphMode.REAL.
    """
    mode = coerce_graph_mode(mode)
    os.makedirs(CACHE_DIR, exist_ok=True)
    if cache_name is None:
        cache_name = f"circuit_{mode.value.lower()}_{max_neurons}_s{seed}.npz"
        
    cache_path = os.path.join(CACHE_DIR, cache_name)
    
    if os.path.exists(cache_path):
        try:
            data = np.load(cache_path, allow_pickle=True)
            stored_mode = GraphMode(str(data["mode"]))
            version_ok = (str(data.get("loader_version", 1)) == str(LOADER_VERSION))
            has_metadata = "provenance_metadata" in data.files and data["provenance_metadata"].item()
            if (stored_mode == mode and len(data["neuron_ids"]) == max_neurons
                    and version_ok and has_metadata):
                populations = build_population_registry(
                    data["neuron_ids"], data["coordinates"], data["tbars"], list(data["sides"])
                )
                graph_cls = (
                    MaleCNSRealGraph if mode == GraphMode.REAL else
                    MaleCNSSpatialSurrogateGraph if mode == GraphMode.SPATIAL_SURROGATE else
                    SyntheticTestGraph
                )
                graph = graph_cls(
                    neuron_ids=data["neuron_ids"],
                    coordinates=data["coordinates"],
                    tbars=data["tbars"],
                    sides=list(data["sides"]),
                    row_offsets=data["row_offsets"],
                    col_indices=data["col_indices"],
                    weights=data["weights"],
                    graph_hash=str(data["graph_hash"]),
                    populations=populations,
                    provenance_metadata=json.loads(data["provenance_metadata"].item()),
                )
                return graph
        except Exception:
            pass

    # Build fresh graph according to mode
    if mode == GraphMode.REAL:
        # R6: never silently substitute a surrogate when REAL was requested.
        graph = build_real_connectome(DEFAULT_CONNECTIONS_PATH, DEFAULT_SOMA_PATH, max_neurons=max_neurons, seed=seed)
    elif mode == GraphMode.SPATIAL_SURROGATE:
        neurons = load_raw_neurons(DEFAULT_SOMA_PATH)
        graph = build_connectome_circuit(neurons, max_neurons=max_neurons, seed=seed)
    elif mode == GraphMode.SYNTHETIC_TEST:
        graph = build_synthetic_test_graph(num_neurons=max_neurons, seed=seed)
    else:
        raise ValueError(f"Unknown GraphMode: {mode}")

    # Cache graph
    np.savez_compressed(
        cache_path,
        neuron_ids=graph.neuron_ids,
        coordinates=graph.coordinates,
        tbars=graph.tbars,
        sides=np.array(graph.sides),
        row_offsets=graph.row_offsets,
        col_indices=graph.col_indices,
        weights=graph.weights,
        graph_hash=graph.graph_hash,
        mode=graph.mode.value,
        provenance_metadata=np.array(json.dumps(graph.provenance_metadata, sort_keys=True)),
        loader_version=np.array(LOADER_VERSION),
    )
    return graph