PSO/test/deep_pso_v6.py

"""
MNIST PSO V6 Study - Phase A & B: Geometry Ablation & Root-Cause Isolation.

Protocol Version: MNIST-PSO-RAW-V6 1.0.0
"""

from __future__ import annotations

import argparse
import copy
from dataclasses import dataclass, asdict
import hashlib
import math
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import torch
import torch.nn as nn
from sklearn.model_selection import train_test_split

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

# Ensure test directory and repo root are in Python path
REPO_ROOT = Path(__file__).resolve().parent.parent if Path(__file__).resolve().parent.name != "PSO" else Path(__file__).resolve().parent
TEST_DIR = REPO_ROOT / "test"
if str(TEST_DIR) not in sys.path:
    sys.path.insert(0, str(TEST_DIR))
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from benchmark_suite import (
    calc_stats,
    compute_model_fingerprint,
    get_hardware_provenance,
    resolve_execution_device,
    save_json_atomic,
    sync_device,
)
from deep_pso_methods import (
    make_compact_cnn,
    build_nested_stratified_subsets,
    evaluate_probabilistic_metrics,
    get_model_probabilities,
)
from pso import __version__ as pso_version

PROTOCOL_VERSION = "MNIST-PSO-RAW-V6 1.0.0"


# =====================================================================
# 1. Dataset Preparation: Train-Only Search/Validation (No Test Split)
# =====================================================================

def prepare_mnist_v6_data(
    split_seed: int = 20260902,
    cache_dir: Optional[Path] = None,
) -> Tuple[
    torch.Tensor, torch.Tensor,
    torch.Tensor, torch.Tensor,
    Dict[int, torch.Tensor],
    str, Dict[str, Any]
]:
    """
    Train-only MNIST data preparation using exclusively MNIST(train=True).
    Never constructs MNIST(train=False).
    Preserves exact V5 split seed (20260902), search-only normalization,
    and nested 2k/10k/50k index stratification.
    """
    from torchvision.datasets import MNIST

    if cache_dir is None:
        cache_dir = Path("result/cache")
    cache_dir.mkdir(parents=True, exist_ok=True)

    # Strictly train=True. Never construct train=False.
    raw_train = MNIST(root=str(cache_dir), train=True, download=True)

    x_train_raw = raw_train.data.float() / 255.0  # (60000, 28, 28)
    y_train_raw = raw_train.targets.long()

    # Stratified split: 50,000 search set and 10,000 validation set
    indices = np.arange(len(y_train_raw))
    search_idx, val_idx = train_test_split(
        indices,
        train_size=50000,
        test_size=10000,
        stratify=y_train_raw.numpy(),
        random_state=split_seed,
    )

    x_search_raw = x_train_raw[search_idx]
    y_search = y_train_raw[search_idx]
    x_val_raw = x_train_raw[val_idx]
    y_val = y_train_raw[val_idx]

    # Fit mean and std on 50k search subset ONLY
    mean_val = float(x_search_raw.mean())
    std_val = float(x_search_raw.std())

    x_search_norm = ((x_search_raw - mean_val) / std_val).unsqueeze(1)  # (50000, 1, 28, 28)
    x_val_norm = ((x_val_raw - mean_val) / std_val).unsqueeze(1)        # (10000, 1, 28, 28)

    # Nested stratified subsets inside 50k search set: 2k inside 10k inside 50k
    nested_subsets = build_nested_stratified_subsets(
        y_search=y_search,
        subset_sizes=[2000, 10000, 50000],
        subset_seed=split_seed,
    )

    # Data fingerprint over search and validation splits (no test split)
    h = hashlib.sha256()
    for t in (x_search_norm, x_val_norm, y_search, y_val):
        h.update(t.detach().cpu().numpy().tobytes())
    data_fp = h.hexdigest()[:16]

    split_h = hashlib.sha256()
    split_h.update(search_idx.tobytes())
    split_h.update(val_idx.tobytes())
    split_fp = split_h.hexdigest()[:16]

    provenance = {
        "input_shape": [1, 28, 28],
        "pca": False,
        "raw_inputs": True,
        "normalization_scope": "search_train_50000_only",
        "train_mean": round(mean_val, 6),
        "train_std": round(std_val, 6),
        "search_samples": 50000,
        "val_samples": 10000,
        "test_samples": 0,
        "official_test_evaluations": 0,
        "split_seed": split_seed,
        "split_fingerprint": split_fp,
    }

    return (
        x_search_norm, y_search,
        x_val_norm, y_val,
        nested_subsets,
        data_fp, provenance
    )


# =====================================================================
# 2. Immutable Geometry Configuration & Protocol Table (G0 - G8)
# =====================================================================

@dataclass(frozen=True)
class V6GeometryConfig:
    """
    Immutable geometry configuration specifying coordinate scales,
    initial position / velocity distributions, mutation, and bounds.
    """
    config_id: str
    scale_type: str = "per_tensor_sd"       # "per_tensor_sd", "global_rms", "identity", "optimizer_default"
    init_position_mode: str = "antithetic"  # "antithetic", "independent"
    position_radius: float = 0.5            # initial position radius
    initial_velocity_radius: float = 0.0    # 0.0 for zero launch velocity, >0 for U(-r, r)
    mutation_prob: float = 0.0              # 0.0 or 0.02
    reset_velocity_radius: float = 0.02     # velocity radius sampled upon mutation
    reflective_bound: float = 3.0           # reflective box half-width (e.g. 3.0 or 6.0)
    projection_seed: Optional[int] = None   # seed for sparse signed-hash subspace projection
    latent_dim: Union[int, str] = "full"    # "full" or integer dimension (e.g. 290, 1024)
    description: str = ""


def get_v6_geometry_table() -> Dict[str, V6GeometryConfig]:
    """
    Returns the complete, approved Phase B geometry configuration table (G0 - G8).
    """
    return {
        "G0": V6GeometryConfig(
            config_id="G0",
            scale_type="per_tensor_sd",
            init_position_mode="antithetic",
            position_radius=0.5,
            initial_velocity_radius=0.0,
            mutation_prob=0.0,
            reflective_bound=3.0,
            description="exact V5 control",
        ),
        "G1": V6GeometryConfig(
            config_id="G1",
            scale_type="global_rms",
            init_position_mode="antithetic",
            position_radius=0.5,
            initial_velocity_radius=0.0,
            mutation_prob=0.0,
            reflective_bound=3.0,
            description="isolate anisotropic per-tensor scaling",
        ),
        "G2": V6GeometryConfig(
            config_id="G2",
            scale_type="per_tensor_sd",
            init_position_mode="antithetic",
            position_radius=0.5,
            initial_velocity_radius=0.5,
            mutation_prob=0.0,
            reflective_bound=3.0,
            description="isolate nonzero launch velocity",
        ),
        "G3": V6GeometryConfig(
            config_id="G3",
            scale_type="per_tensor_sd",
            init_position_mode="antithetic",
            position_radius=0.5,
            initial_velocity_radius=0.0,
            mutation_prob=0.02,
            reset_velocity_radius=0.02,
            reflective_bound=3.0,
            description="isolate mutation",
        ),
        "G4": V6GeometryConfig(
            config_id="G4",
            scale_type="per_tensor_sd",
            init_position_mode="antithetic",
            position_radius=0.5,
            initial_velocity_radius=0.5,
            mutation_prob=0.02,
            reset_velocity_radius=0.02,
            reflective_bound=3.0,
            description="velocity x mutation interaction",
        ),
        "G5": V6GeometryConfig(
            config_id="G5",
            scale_type="per_tensor_sd",
            init_position_mode="antithetic",
            position_radius=0.5,
            initial_velocity_radius=0.5,
            mutation_prob=0.02,
            reset_velocity_radius=0.02,
            reflective_bound=6.0,
            description="test sufficient bound expansion",
        ),
        "G6": V6GeometryConfig(
            config_id="G6",
            scale_type="per_tensor_sd",
            init_position_mode="antithetic",
            position_radius=1.5,
            initial_velocity_radius=0.5,
            mutation_prob=0.02,
            reset_velocity_radius=0.02,
            reflective_bound=6.0,
            description="test broader normalized initialization",
        ),
        "G7": V6GeometryConfig(
            config_id="G7",
            scale_type="per_tensor_sd",
            init_position_mode="independent",
            position_radius=0.5,
            initial_velocity_radius=0.5,
            mutation_prob=0.0,
            reflective_bound=3.0,
            description="isolate antithetic position coupling against G2",
        ),
        "G8": V6GeometryConfig(
            config_id="G8",
            scale_type="optimizer_default",
            init_position_mode="independent",
            position_radius=0.05,
            initial_velocity_radius=0.05,
            mutation_prob=0.02,
            reset_velocity_radius=0.02,
            reflective_bound=3.0,
            description="retained semantic control (public Optimizer)",
        ),
    }


def compute_equalized_subspace_radius(
    latent_dim: int,
    total_dim: int = 9098,
    base_radius: float = 0.5,
) -> float:
    """
    Computes equalized subspace initialization/bound radius for Phase C.
    Scales base_radius by sqrt(total_dim / latent_dim) to hold decoded per-parameter RMS constant.
    """
    if latent_dim >= total_dim:
        return base_radius
    return float(base_radius * math.sqrt(total_dim / latent_dim))


# =====================================================================
# 3. Deterministic Latent Transform & Swarm Construction
# =====================================================================

class V6LatentTransform:
    def __init__(
        self,
        base_model: nn.Module,
        geom_config: V6GeometryConfig,
        device: torch.device,
    ):
        self.device = device
        self.geom_config = geom_config
        self.base_params = [p.detach().clone().to(device) for p in base_model.parameters()]
        self.param_shapes = [p.shape for p in self.base_params]
        self.param_numels = [p.numel() for p in self.base_params]
        self.total_dim = sum(self.param_numels)
        self.base_vec = torch.cat([p.view(-1) for p in self.base_params]).to(device)

        scale_type = geom_config.scale_type.lower()
        if scale_type == "per_tensor_sd":
            tensor_scales = []
            for p in self.base_params:
                std_val = float(p.std())
                scale = max(std_val, 1e-4)
                scale_tensor = torch.full_like(p, scale)
                tensor_scales.append(scale_tensor.view(-1))
            self.scale_vec = torch.cat(tensor_scales).to(device)
        elif scale_type == "global_rms":
            rms_val = float(torch.sqrt(torch.mean(self.base_vec ** 2)))
            scale = max(rms_val, 1e-4)
            self.scale_vec = torch.full_like(self.base_vec, scale, device=device)
        elif scale_type == "identity":
            self.scale_vec = torch.ones_like(self.base_vec, device=device)
        else:
            # Default fallback for optimizer or custom
            self.scale_vec = torch.ones_like(self.base_vec, device=device)

        latent_dim = geom_config.latent_dim
        if isinstance(latent_dim, str) and latent_dim.lower() == "full":
            self.latent_dim = self.total_dim
            self.is_full = True
        else:
            self.latent_dim = int(latent_dim)
            self.is_full = (self.latent_dim == self.total_dim)

        if not self.is_full:
            j_indices = np.arange(self.total_dim, dtype=np.int64)
            seed_offset = geom_config.projection_seed if geom_config.projection_seed is not None else 0
            h1 = ((j_indices + 1 + seed_offset) * 2654435761) % (2**32)
            k_indices = h1 % self.latent_dim
            h2 = ((j_indices + 1 + seed_offset) * 1597334677) % (2**32)
            signs = np.where((h2 % 2) == 0, 1.0, -1.0)

            bin_counts = np.bincount(k_indices, minlength=self.latent_dim)
            count_per_j = bin_counts[k_indices]
            scale_per_j = 1.0 / np.sqrt(np.maximum(count_per_j, 1))
            combined_weights = signs * scale_per_j

            self.k_indices = torch.tensor(k_indices, dtype=torch.long, device=device)
            self.weights = torch.tensor(combined_weights, dtype=torch.float32, device=device)

    def decode(self, Z: torch.Tensor) -> torch.Tensor:
        """
        Transforms latent batch Z (N, d) into full parameter batch (N, D).
        theta = base_vec + scale_vec * delta
        """
        if self.is_full:
            delta = Z
        else:
            delta = Z[:, self.k_indices] * self.weights
        return self.base_vec + self.scale_vec * delta

    def load_vector_to_model(self, theta_vec: torch.Tensor, model: nn.Module):
        """Loads a single parameter vector into model parameters in-place."""
        offset = 0
        with torch.no_grad():
            for p, shape, numel in zip(model.parameters(), self.param_shapes, self.param_numels):
                p.copy_(theta_vec[offset:offset + numel].view(shape))
                offset += numel

    def init_swarm(self, swarm_size: int, seed: int) -> torch.Tensor:
        """
        Initializes particle positions in latent space Z (N, d).
        Particle 0 is exact base vector (z = 0).
        Uses a separate CPU generator to preserve G0 parity with V5.
        """
        init_radius = self.geom_config.position_radius
        mode = self.geom_config.init_position_mode.lower()

        rng = torch.Generator(device="cpu")
        rng.manual_seed(seed)

        Z = torch.zeros((swarm_size, self.latent_dim), dtype=torch.float32)

        if mode == "antithetic":
            max_pair_idx = swarm_size - 1 if (swarm_size % 2 != 0) else swarm_size - 2
            idx = 1
            while idx < max_pair_idx:
                sample = (torch.rand(self.latent_dim, generator=rng) * 2.0 - 1.0) * init_radius
                Z[idx] = sample
                Z[idx + 1] = -sample
                idx += 2
        else:  # independent
            for idx in range(1, swarm_size):
                Z[idx] = (torch.rand(self.latent_dim, generator=rng) * 2.0 - 1.0) * init_radius

        return Z.to(self.device)


# =====================================================================
# 4. Device-Resident Configurable V6 Latent Adaptive-Moment PSO Engine
# =====================================================================

def evaluate_latent_batch(
    Z: torch.Tensor,
    transform: V6LatentTransform,
    model: nn.Module,
    x_sub_dev: torch.Tensor,
    y_sub_dev: torch.Tensor,
    batch_size: int = 1000,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Evaluates latent batch Z (N, d) on device-resident (x_sub_dev, y_sub_dev).
    Returns (losses, accuracies) tensors of shape (N,).
    """
    N = Z.shape[0]
    device = Z.device
    losses = torch.zeros(N, dtype=torch.float32, device=device)
    accuracies = torch.zeros(N, dtype=torch.float32, device=device)

    loss_fn = nn.CrossEntropyLoss(reduction="sum")
    num_samples = len(y_sub_dev)

    model.eval()
    with torch.inference_mode():
        for i in range(N):
            theta_vec = transform.decode(Z[i:i+1]).squeeze(0)
            transform.load_vector_to_model(theta_vec, model)

            total_loss = torch.tensor(0.0, device=device)
            correct = torch.tensor(0, dtype=torch.long, device=device)

            for b_start in range(0, num_samples, batch_size):
                xb = x_sub_dev[b_start:b_start + batch_size]
                yb = y_sub_dev[b_start:b_start + batch_size]
                logits = model(xb)
                batch_loss = loss_fn(logits, yb)
                total_loss += batch_loss
                preds = logits.argmax(dim=1)
                correct += (preds == yb).sum()

            losses[i] = total_loss / num_samples
            accuracies[i] = (correct.float() / num_samples) * 100.0

    return losses, accuracies


def run_v6_pso(
    transform: V6LatentTransform,
    base_model: nn.Module,
    x_search: torch.Tensor,
    y_search: torch.Tensor,
    x_val: torch.Tensor,
    y_val: torch.Tensor,
    nested_subsets: Dict[int, torch.Tensor],
    schedule_str: str,
    epochs: int,
    swarm_size: int,
    seed: int,
    device: torch.device,
    geom_config: V6GeometryConfig,
    val_check_interval: int = 10,
    transition_reset_policy: str = "none",
) -> Dict[str, Any]:
    """
    Device-resident Latent Adaptive-Moment PSO for V6 study.
    Supports G0-G7 configurations, nonzero launch velocity, mutation moment reset,
    reflective bounds, state-neutral validation checkpoints, and full telemetry.
    """
    sync_device(device)
    start_time = time.time()
    validation_wall_time = 0.0

    schedule_stages = []
    if schedule_str:
        parts = schedule_str.split(",")
        for p in parts:
            sz_str, ep_str = p.split(":")
            schedule_stages.append((int(sz_str), int(ep_str)))

    if not schedule_stages:
        schedule_stages = [(50000, epochs)]

    c0 = c1 = 1.49618
    w = 0.7298
    blend = 0.06
    step = 0.5
    beta1 = 0.9
    beta2 = 0.999
    reflective_bound = geom_config.reflective_bound

    latent_dim = transform.latent_dim
    Z = transform.init_swarm(swarm_size=swarm_size, seed=seed)

    # Initial launch velocity
    if geom_config.initial_velocity_radius > 0.0:
        vel_rng = torch.Generator(device="cpu")
        vel_rng.manual_seed(seed + 1000)
        r_v = geom_config.initial_velocity_radius
        V_cpu = (torch.rand((swarm_size, latent_dim), generator=vel_rng) * 2.0 - 1.0) * r_v
        V_cpu[0] = 0.0  # Particle 0 launch velocity remains zero
        V = V_cpu.to(device)
    else:
        V = torch.zeros((swarm_size, latent_dim), dtype=torch.float32, device=device)

    M = torch.zeros((swarm_size, latent_dim), dtype=torch.float32, device=device)
    V_sq = torch.zeros((swarm_size, latent_dim), dtype=torch.float32, device=device)
    moment_steps = torch.zeros(swarm_size, dtype=torch.int64, device=device)

    P = Z.clone()
    P_loss = torch.full((swarm_size,), float("inf"), dtype=torch.float32, device=device)
    P_acc = torch.zeros((swarm_size,), dtype=torch.float32, device=device)

    gbest_z = Z[0].clone()
    gbest_loss = float("inf")
    gbest_acc = 0.0

    # Counters & Telemetry
    total_queries = 0
    total_sample_evaluations = 0
    transition_reevaluation_counts = 0
    validation_evaluations = 0
    pbest_update_counts = 0
    boundary_hits = 0
    last_improvement_epoch = 0
    mutation_events = 0

    # Random generators: move_rng handles standard velocity draws; mut_rng handles mutation draws
    move_rng = torch.Generator(device=device)
    move_rng.manual_seed(seed)
    mut_rng = torch.Generator(device=device)
    mut_rng.manual_seed(seed + 2000)

    model = copy.deepcopy(base_model).to(device)
    x_val_dev = x_val.to(device)
    y_val_dev = y_val.to(device)

    stage_histories = []
    epoch_counter = 0

    for stage_idx, (size, stage_epochs) in enumerate(schedule_stages):
        subset_indices = nested_subsets[size]
        x_sub_dev = x_search[subset_indices].to(device)
        y_sub_dev = y_search[subset_indices].to(device)

        # Objective transition check
        if stage_idx > 0:
            re_losses, re_accs = evaluate_latent_batch(
                P, transform, model, x_sub_dev, y_sub_dev
            )
            P_loss = re_losses
            P_acc = re_accs

            total_queries += swarm_size
            total_sample_evaluations += swarm_size * size
            transition_reevaluation_counts += swarm_size

            min_loss_val = P_loss.min()
            candidates_mask = (P_loss <= min_loss_val + 1e-7)
            best_p_idx = int(torch.where(candidates_mask, P_acc, torch.tensor(-1.0, device=device)).argmax().item())

            gbest_z = P[best_p_idx].clone()
            gbest_loss = float(P_loss[best_p_idx].item())
            gbest_acc = float(P_acc[best_p_idx].item())

            # Transition reset policy
            if transition_reset_policy in ("reset_vm", "reset_all"):
                V.zero_()
                M.zero_()
                V_sq.zero_()
                moment_steps.zero_()

        for ep in range(1, stage_epochs + 1):
            epoch_counter += 1

            # 1. Evaluate current swarm positions
            curr_losses, curr_accs = evaluate_latent_batch(
                Z, transform, model, x_sub_dev, y_sub_dev
            )
            total_queries += swarm_size
            total_sample_evaluations += swarm_size * size

            # 2. Vectorized pbest updates
            better_loss = curr_losses < P_loss - 1e-7
            equal_loss = torch.abs(curr_losses - P_loss) <= 1e-7
            better_acc = curr_accs > P_acc
            update_mask = better_loss | (equal_loss & better_acc)

            n_updated = int(update_mask.sum().item())
            pbest_update_counts += n_updated
            if n_updated > 0:
                last_improvement_epoch = epoch_counter

            P[update_mask] = Z[update_mask]
            P_loss[update_mask] = curr_losses[update_mask]
            P_acc[update_mask] = curr_accs[update_mask]

            # Rebuild gbest from P each epoch
            min_loss_val = P_loss.min()
            candidates_mask = (P_loss <= min_loss_val + 1e-7)
            best_p_idx = int(torch.where(candidates_mask, P_acc, torch.tensor(-1.0, device=device)).argmax().item())

            gbest_z = P[best_p_idx].clone()
            gbest_loss = float(P_loss[best_p_idx].item())
            gbest_acc = float(P_acc[best_p_idx].item())

            # Validation Checkpoint (State-Neutral)
            val_loss = None
            val_acc = None
            if val_check_interval > 0 and (epoch_counter % val_check_interval == 0 or epoch_counter == epochs):
                sync_device(device)
                validation_start = time.time()
                with torch.inference_mode():
                    v_losses, v_accs = evaluate_latent_batch(
                        gbest_z.unsqueeze(0), transform, model, x_val_dev, y_val_dev
                    )
                    val_loss = float(v_losses[0].item())
                    val_acc = float(v_accs[0].item())
                    validation_evaluations += 1
                sync_device(device)
                validation_wall_time += time.time() - validation_start

            stage_histories.append({
                "epoch": epoch_counter,
                "stage": stage_idx,
                "subset_size": size,
                "gbest_loss": round(gbest_loss, 6),
                "gbest_acc": round(gbest_acc, 4),
                "val_loss": round(val_loss, 6) if val_loss is not None else None,
                "val_acc": round(val_acc, 4) if val_acc is not None else None,
            })

            # 3. Movement Step
            r1 = torch.rand((swarm_size, latent_dim), generator=move_rng, device=device)
            r2 = torch.rand((swarm_size, latent_dim), generator=move_rng, device=device)

            V_raw = w * V + c0 * r1 * (P - Z) + c1 * r2 * (gbest_z.unsqueeze(0) - Z)

            # Mutation check
            if geom_config.mutation_prob > 0.0:
                mut_draws = torch.rand(swarm_size, generator=mut_rng, device=device)
                mut_mask = mut_draws < geom_config.mutation_prob
                if mut_mask.any():
                    n_mut = int(mut_mask.sum().item())
                    r_mut = geom_config.reset_velocity_radius
                    mut_v = (torch.rand((n_mut, latent_dim), generator=mut_rng, device=device) * 2.0 - 1.0) * r_mut
                    V_raw[mut_mask] = mut_v
                    # Clear moment state for mutated particles
                    M[mut_mask] = 0.0
                    V_sq[mut_mask] = 0.0
                    moment_steps[mut_mask] = 0
                    mutation_events += n_mut

            moment_steps += 1
            M = beta1 * M + (1 - beta1) * V_raw
            V_sq = beta2 * V_sq + (1 - beta2) * (V_raw ** 2)

            step_values = moment_steps.to(dtype=M.dtype).unsqueeze(1)
            M_hat = M / (1.0 - torch.pow(beta1, step_values))
            V_sq_hat = V_sq / (1.0 - torch.pow(beta2, step_values))

            dir_moment = M_hat / (torch.sqrt(V_sq_hat) + 1e-8)
            historical_scale = torch.sqrt(torch.mean(V_sq_hat, dim=1, keepdim=True))
            V_moment = step * dir_moment * historical_scale
            V_new = (1.0 - blend) * V_raw + blend * V_moment

            Z_new = Z + V_new

            pos_mask = Z_new > reflective_bound
            neg_mask = Z_new < -reflective_bound
            hit_mask = pos_mask | neg_mask
            boundary_hits += int(hit_mask.sum().item())

            Z_new[pos_mask] = 2.0 * reflective_bound - Z_new[pos_mask]
            V_new[pos_mask] = -V_new[pos_mask]

            Z_new[neg_mask] = -2.0 * reflective_bound - Z_new[neg_mask]
            V_new[neg_mask] = -V_new[neg_mask]

            Z_new = torch.clamp(Z_new, -reflective_bound, reflective_bound)

            Z = Z_new
            V = V_new

            if epoch_counter >= epochs:
                break
        if epoch_counter >= epochs:
            break

    sync_device(device)
    validation_start = time.time()
    # Select the best final pbest on validation, matching the G8 control.
    pbest_val_losses, pbest_val_accs = evaluate_latent_batch(
        P, transform, model, x_val_dev, y_val_dev
    )
    validation_evaluations += swarm_size
    min_val_loss = pbest_val_losses.min()
    val_candidates = pbest_val_losses <= min_val_loss + 1e-7
    val_best_idx = int(
        torch.where(
            val_candidates,
            pbest_val_accs,
            torch.tensor(-1.0, device=device),
        ).argmax().item()
    )

    val_selected_model = copy.deepcopy(base_model).to(device)
    theta_selected = transform.decode(P[val_best_idx].unsqueeze(0)).squeeze(0)
    transform.load_vector_to_model(theta_selected, val_selected_model)
    val_probabilities = get_model_probabilities(val_selected_model, x_val_dev, device)
    val_metrics = evaluate_probabilistic_metrics(val_probabilities, y_val_dev)
    validation_evaluations += 1
    sync_device(device)
    validation_wall_time += time.time() - validation_start

    gbest_val_loss = float(pbest_val_losses[best_p_idx].item())
    gbest_val_acc = float(pbest_val_accs[best_p_idx].item())

    wall_time = time.time() - start_time
    optimization_wall_time = max(0.0, wall_time - validation_wall_time)

    velocity_rms = float(torch.sqrt(torch.mean(V ** 2)).item())
    position_radius = float(torch.norm(Z - Z.mean(dim=0), dim=1).mean().item())
    total_coords = swarm_size * latent_dim * max(epoch_counter, 1)
    boundary_occupancy = round(boundary_hits / max(total_coords, 1), 6)

    return {
        "config_id": geom_config.config_id,
        "gbest_z": gbest_z,
        "gbest_loss": round(gbest_loss, 6),
        "gbest_acc": round(gbest_acc, 4),
        "gbest_val_loss": round(gbest_val_loss, 6),
        "gbest_val_acc": round(gbest_val_acc, 4),
        "val_selected_particle_idx": val_best_idx,
        "val_selected_loss": round(val_metrics["nll"], 6),
        "val_selected_acc": round(val_metrics["accuracy"], 4),
        "val_metrics": val_metrics,
        "final_P": P,
        "wall_time_sec": round(wall_time, 4),
        "optimization_wall_time_sec": round(optimization_wall_time, 4),
        "validation_wall_time_sec": round(validation_wall_time, 4),
        "total_queries": total_queries,
        "total_sample_evaluations": total_sample_evaluations,
        "transition_reevaluation_counts": transition_reevaluation_counts,
        "validation_evaluations": validation_evaluations,
        "official_test_evaluations": 0,
        "mutation_events": mutation_events,
        "final_moment_steps": moment_steps.detach().cpu().tolist(),
        "pbest_update_counts": pbest_update_counts,
        "boundary_hits": boundary_hits,
        "boundary_occupancy": boundary_occupancy,
        "last_improvement_epoch": last_improvement_epoch,
        "velocity_rms": round(velocity_rms, 6),
        "position_radius": round(position_radius, 6),
        "stage_histories": stage_histories,
    }


# =====================================================================
# 5. Public Optimizer G8 Semantic Control Engine
# =====================================================================

def run_g8_optimizer(
    base_model: nn.Module,
    x_2k: torch.Tensor,
    y_2k: torch.Tensor,
    x_val: torch.Tensor,
    y_val: torch.Tensor,
    epochs: int,
    swarm_size: int,
    seed: int,
    device: torch.device,
) -> Dict[str, Any]:
    """
    Runs public Optimizer for G8 semantic control over the exact 2k search tensors.
    Evaluates every particle's pbest on validation to select the endpoint.
    Retains training gbest metrics separately.
    Official test evaluations count is explicitly zero.
    """
    from pso.optimizer import Optimizer
    sync_device(device)
    start_t = time.time()

    model = copy.deepcopy(base_model).to(device)
    loss_fn = nn.CrossEntropyLoss()
    validation_loss_fn = nn.CrossEntropyLoss(reduction="sum")

    opt = Optimizer(
        model=model,
        loss=loss_fn,
        task="multiclass",
        method="adaptive_moment",
        evaluation="full",
        n_particles=swarm_size,
        c0=1.49618,
        c1=1.49618,
        w_min=0.7298,
        w_max=0.7298,
        particle_min=-3.0,
        particle_max=3.0,
        boundary_strategy="reflect",
        velocity_limit_ratio=0.025,
        mutation_swarm=0.02,
        initialization="model_noise",
        initial_position_noise=0.05,
        moment_blend=0.06,
        moment_step_size=0.5,
        moment_beta1=0.9,
        seed=seed,
        device=device,
    )

    x_2k_dev = x_2k.to(device)
    y_2k_dev = y_2k.to(device)

    best_score = opt.fit(x_2k_dev, y_2k_dev, epochs=epochs, renewal="loss")
    sync_device(device)
    optimization_wall_time = time.time() - start_t
    validation_start = time.time()

    # Evaluate every particle pbest vector on the validation set
    val_losses = []
    val_accs = []
    pbest_models = []

    x_val_dev = x_val.to(device)
    y_val_dev = y_val.to(device)
    num_val = len(y_val_dev)

    for p in opt.particles:
        p_model = copy.deepcopy(base_model).to(device)
        offset = 0
        with torch.no_grad():
            for param in p_model.parameters():
                n = param.numel()
                param.copy_(p.personal_best_weights[offset:offset + n].view(param.shape))
                offset += n

        p_model.eval()
        with torch.inference_mode():
            total_loss = 0.0
            correct = 0
            for b_start in range(0, num_val, 1000):
                xb = x_val_dev[b_start:b_start + 1000]
                yb = y_val_dev[b_start:b_start + 1000]
                logits = p_model(xb)
                total_loss += float(validation_loss_fn(logits, yb).item())
                correct += int((logits.argmax(dim=1) == yb).sum().item())
            val_loss = total_loss / num_val
            val_acc = (correct / num_val) * 100.0

            val_losses.append(val_loss)
            val_accs.append(val_acc)
            pbest_models.append(p_model)

    # Rank particles on validation NLL ascending, val accuracy descending
    min_val_loss = min(val_losses)
    val_candidates = [
        index
        for index, loss in enumerate(val_losses)
        if loss <= min_val_loss + 1e-7
    ]
    best_idx = max(val_candidates, key=lambda index: val_accs[index])

    val_selected_model = pbest_models[best_idx]
    val_probabilities = get_model_probabilities(val_selected_model, x_val_dev, device)
    val_metrics = evaluate_probabilistic_metrics(val_probabilities, y_val_dev)
    sync_device(device)
    validation_wall_time = time.time() - validation_start
    wall_t = optimization_wall_time + validation_wall_time

    total_queries = swarm_size * epochs
    total_sample_evaluations = total_queries * len(y_2k)
    validation_evaluations = swarm_size + 1

    return {
        "config_id": "G8",
        "gbest_loss": round(float(best_score[0]), 6),
        "gbest_acc": round(float(best_score[1]) * 100.0, 4),
        "val_selected_loss": round(val_metrics["nll"], 6),
        "val_selected_acc": round(val_metrics["accuracy"], 4),
        "val_metrics": val_metrics,
        "wall_time_sec": round(wall_t, 4),
        "optimization_wall_time_sec": round(optimization_wall_time, 4),
        "validation_wall_time_sec": round(validation_wall_time, 4),
        "total_queries": total_queries,
        "total_sample_evaluations": total_sample_evaluations,
        "transition_reevaluation_counts": 0,
        "validation_evaluations": validation_evaluations,
        "official_test_evaluations": 0,
        "pbest_update_counts": 0,
        "boundary_hits": 0,
        "boundary_occupancy": 0.0,
        "last_improvement_epoch": epochs,
        "velocity_rms": 0.0,
        "position_radius": 0.0,
        "stage_histories": [],
    }


# =====================================================================
# 6. Selection & Paired Factor Analysis
# =====================================================================

def compute_paired_factor_deltas(
    screen_results: Dict[str, Dict[str, Any]]
) -> Dict[str, Dict[str, float]]:
    """
    Computes paired factor deltas for Phase B screen.
    A factor is provisionally material if NLL improves by >= 0.05 or Acc improves by >= 2.0pp.
    """
    pairs = [
        ("delta_scale_G1_vs_G0", "G1", "G0", "global RMS scale vs per-tensor SD"),
        ("delta_vel_G2_vs_G0", "G2", "G0", "launch velocity U(-0.5,0.5) vs 0"),
        ("delta_mut_G3_vs_G0", "G3", "G0", "mutation 0.02 vs 0"),
        ("delta_vel_mut_G4_vs_G2", "G4", "G2", "mutation interaction given velocity"),
        ("delta_bound_G5_vs_G4", "G5", "G4", "bound box 6 vs 3"),
        ("delta_radius_G6_vs_G5", "G6", "G5", "initial position radius 1.5 vs 0.5"),
        ("delta_init_G7_vs_G2", "G7", "G2", "independent vs antithetic init"),
    ]

    deltas = {}
    for key, c_test, c_ref, desc in pairs:
        if c_test in screen_results and c_ref in screen_results:
            r_test = screen_results[c_test]
            r_ref = screen_results[c_ref]
            nll_diff = round(r_test["val_selected_loss"] - r_ref["val_selected_loss"], 6)
            acc_diff = round(r_test["val_selected_acc"] - r_ref["val_selected_acc"], 4)
            is_material = (nll_diff <= -0.05) or (acc_diff >= 2.0)
            deltas[key] = {
                "test_config": c_test,
                "ref_config": c_ref,
                "nll_diff": nll_diff,
                "acc_diff": acc_diff,
                "material": is_material,
                "description": desc,
            }

    # Bundle comparison: G8 vs best normalized config
    norm_configs = [c for c in screen_results if c != "G8"]
    if norm_configs:
        best_norm_id = min(norm_configs, key=lambda k: (screen_results[k]["val_selected_loss"], -screen_results[k]["val_selected_acc"]))
        r_g8 = screen_results["G8"]
        r_best_norm = screen_results[best_norm_id]
        nll_diff = round(r_g8["val_selected_loss"] - r_best_norm["val_selected_loss"], 6)
        acc_diff = round(r_g8["val_selected_acc"] - r_best_norm["val_selected_acc"], 4)
        deltas["delta_bundle_G8_vs_best_norm"] = {
            "test_config": "G8",
            "ref_config": best_norm_id,
            "nll_diff": nll_diff,
            "acc_diff": acc_diff,
            "material": abs(nll_diff) >= 0.05 or abs(acc_diff) >= 2.0,
            "description": f"Optimizer G8 control vs best normalized ({best_norm_id})",
        }

    return deltas


def select_confirmation_configs(
    screen_results: Dict[str, Dict[str, Any]]
) -> List[str]:
    """
    Deterministically selects G0, G1, G8 plus top 2 eligible normalized configs from G2-G7.
    """
    mandatory = ["G0", "G1", "G8"]
    eligible = ["G2", "G3", "G4", "G5", "G6", "G7"]

    # Filter available eligible configs
    valid_eligible = [c for c in eligible if c in screen_results]
    valid_eligible.sort(
        key=lambda c: (screen_results[c]["val_selected_loss"], -screen_results[c]["val_selected_acc"])
    )

    top2_other = valid_eligible[:2]
    selected = mandatory + top2_other
    return selected


def evaluate_root_cause_statuses(
    confirm_aggregates: Dict[str, Dict[str, Any]]
) -> Dict[str, Dict[str, Any]]:
    """Classify each predeclared geometry hypothesis from confirmed mean metrics."""
    statuses: Dict[str, Dict[str, Any]] = {}
    g8_stats = confirm_aggregates.get("G8")

    if g8_stats is None:
        statuses["regression_recovered"] = {
            "status": "unresolved",
            "recovered_configs": [],
            "description": "G8 control was not confirmed",
        }
    else:
        g8_acc = g8_stats["val_selected_acc"]["mean"]
        g8_nll = g8_stats["val_selected_nll"]["mean"]
        recovered = []
        for cid, stats in confirm_aggregates.items():
            if cid == "G8":
                continue
            acc = stats["val_selected_acc"]["mean"]
            nll = stats["val_selected_nll"]["mean"]
            if acc >= g8_acc - 1.0 and nll <= g8_nll + 0.05:
                recovered.append(cid)
        statuses["regression_recovered"] = {
            "status": "supported" if recovered else "rejected",
            "recovered_configs": recovered,
            "g8_mean_acc": g8_acc,
            "g8_mean_nll": g8_nll,
            "description": "Normalized geometry is within 1.0pp accuracy and 0.05 NLL of G8",
        }

    hypotheses = [
        ("anisotropic_per_tensor_scaling", "G1", "G0", "global RMS scale versus per-tensor SD"),
        ("nonzero_launch_velocity", "G2", "G0", "nonzero launch velocity versus zero"),
        ("mutation", "G3", "G0", "mutation 0.02 versus none"),
        ("velocity_mutation_interaction", "G4", "G2", "mutation given nonzero velocity"),
        ("bound_expansion", "G5", "G4", "normalized bound 6 versus 3"),
        ("broader_initialization", "G6", "G5", "initial radius 1.5 versus 0.5"),
        ("independent_initialization", "G7", "G2", "independent versus antithetic positions"),
    ]
    for name, test_id, ref_id, description in hypotheses:
        if test_id not in confirm_aggregates or ref_id not in confirm_aggregates:
            statuses[name] = {
                "status": "unresolved",
                "test_config": test_id,
                "ref_config": ref_id,
                "description": description,
            }
            continue
        test_stats = confirm_aggregates[test_id]
        ref_stats = confirm_aggregates[ref_id]
        nll_diff = (
            test_stats["val_selected_nll"]["mean"]
            - ref_stats["val_selected_nll"]["mean"]
        )
        acc_diff = (
            test_stats["val_selected_acc"]["mean"]
            - ref_stats["val_selected_acc"]["mean"]
        )
        supported = nll_diff <= -0.05 or acc_diff >= 2.0
        statuses[name] = {
            "status": "supported" if supported else "rejected",
            "test_config": test_id,
            "ref_config": ref_id,
            "mean_nll_diff": round(nll_diff, 6),
            "mean_acc_diff": round(acc_diff, 4),
            "description": description,
        }

    return statuses


def artifact_safe_run(result: Dict[str, Any]) -> Dict[str, Any]:
    """Drop engine-only state and convert remaining tensor values for JSON."""
    safe: Dict[str, Any] = {}
    for key, value in result.items():
        if key in {"gbest_z", "final_P"}:
            continue
        if torch.is_tensor(value):
            value = value.detach().cpu().item() if value.numel() == 1 else value.detach().cpu().tolist()
        safe[key] = value
    return safe


# =====================================================================
# 7. Experiment Runner Pipeline (Screen, Confirm, All)
# =====================================================================

def run_phase_b_screen(
    x_search: torch.Tensor,
    y_search: torch.Tensor,
    x_val: torch.Tensor,
    y_val: torch.Tensor,
    nested_subsets: Dict[int, torch.Tensor],
    device: torch.device,
    seed: int = 91,
    swarm_size: int = 30,
    epochs: int = 160,
) -> Dict[str, Any]:
    """Runs Phase B screen (G0 - G8 at seed 91)."""
    table = get_v6_geometry_table()
    screen_results = {}

    for cid in ["G0", "G1", "G2", "G3", "G4", "G5", "G6", "G7"]:
        cfg = table[cid]
        base_model = make_compact_cnn(seed=41).to(device)
        transform = V6LatentTransform(base_model, cfg, device)
        res = run_v6_pso(
            transform=transform,
            base_model=base_model,
            x_search=x_search,
            y_search=y_search,
            x_val=x_val,
            y_val=y_val,
            nested_subsets=nested_subsets,
            schedule_str=f"2000:{epochs}",
            epochs=epochs,
            swarm_size=swarm_size,
            seed=seed,
            device=device,
            geom_config=cfg,
        )
        res["seed"] = seed
        res["geometry_config"] = asdict(cfg)
        screen_results[cid] = artifact_safe_run(res)

    # Run G8
    base_model_g8 = make_compact_cnn(seed=41).to(device)
    x_2k = x_search[nested_subsets[2000]]
    y_2k = y_search[nested_subsets[2000]]
    g8_res = run_g8_optimizer(
        base_model=base_model_g8,
        x_2k=x_2k,
        y_2k=y_2k,
        x_val=x_val,
        y_val=y_val,
        epochs=epochs,
        swarm_size=swarm_size,
        seed=seed,
        device=device,
    )
    g8_res["seed"] = seed
    g8_res["geometry_config"] = asdict(table["G8"])
    screen_results["G8"] = artifact_safe_run(g8_res)

    # Paired factor deltas & confirmation selection
    factor_deltas = compute_paired_factor_deltas(screen_results)
    selected_for_confirm = select_confirmation_configs(screen_results)

    return {
        "phase": "screen",
        "seed": seed,
        "swarm_size": swarm_size,
        "epochs": epochs,
        "screen_results": screen_results,
        "factor_deltas": factor_deltas,
        "selected_for_confirm": selected_for_confirm,
    }


def run_phase_b_confirm(
    selected_configs: List[str],
    x_search: torch.Tensor,
    y_search: torch.Tensor,
    x_val: torch.Tensor,
    y_val: torch.Tensor,
    nested_subsets: Dict[int, torch.Tensor],
    device: torch.device,
    seeds: List[int] = [101, 102, 103],
    swarm_size: int = 60,
    epochs: int = 420,
) -> Dict[str, Any]:
    """Runs Phase B confirmation over selected configurations across seeds 101-103."""
    table = get_v6_geometry_table()
    confirm_runs = {cid: [] for cid in selected_configs}

    x_2k = x_search[nested_subsets[2000]]
    y_2k = y_search[nested_subsets[2000]]

    for cid in selected_configs:
        for seed in seeds:
            if cid == "G8":
                base_model = make_compact_cnn(seed=41).to(device)
                res = run_g8_optimizer(
                    base_model=base_model,
                    x_2k=x_2k,
                    y_2k=y_2k,
                    x_val=x_val,
                    y_val=y_val,
                    epochs=epochs,
                    swarm_size=swarm_size,
                    seed=seed,
                    device=device,
                )
            else:
                cfg = table[cid]
                base_model = make_compact_cnn(seed=41).to(device)
                transform = V6LatentTransform(base_model, cfg, device)
                res = run_v6_pso(
                    transform=transform,
                    base_model=base_model,
                    x_search=x_search,
                    y_search=y_search,
                    x_val=x_val,
                    y_val=y_val,
                    nested_subsets=nested_subsets,
                    schedule_str=f"2000:{epochs}",
                    epochs=epochs,
                    swarm_size=swarm_size,
                    seed=seed,
                    device=device,
                    geom_config=cfg,
                )
            res["seed"] = seed
            res["geometry_config"] = asdict(table[cid])
            confirm_runs[cid].append(artifact_safe_run(res))

    # Compute aggregate stats per config across seeds
    confirm_aggregates = {}
    for cid, runs in confirm_runs.items():
        accs = [r["val_selected_acc"] for r in runs]
        nlls = [r["val_selected_loss"] for r in runs]
        briers = [r["val_metrics"]["brier"] for r in runs if "val_metrics" in r and "brier" in r["val_metrics"]]
        eces = [r["val_metrics"]["ece"] for r in runs if "val_metrics" in r and "ece" in r["val_metrics"]]

        confirm_aggregates[cid] = {
            "val_selected_acc": calc_stats(accs),
            "val_selected_nll": calc_stats(nlls),
            "brier": calc_stats(briers) if briers else {},
            "ece": calc_stats(eces) if eces else {},
            "num_seeds": len(runs),
        }

    root_cause_statuses = evaluate_root_cause_statuses(confirm_aggregates)

    return {
        "phase": "confirm",
        "seeds": seeds,
        "swarm_size": swarm_size,
        "epochs": epochs,
        "selected_configs": selected_configs,
        "confirm_runs": confirm_runs,
        "confirm_aggregates": confirm_aggregates,
        "root_cause_statuses": root_cause_statuses,
    }


# =====================================================================
# 8. CSV & Plot Artifact Writers
# =====================================================================

def save_csv_summary_v6(payload: Dict[str, Any], csv_path: Path):
    """Saves concise CSV summary of V6 study results."""
    import csv
    csv_path.parent.mkdir(parents=True, exist_ok=True)

    with open(csv_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["protocol_version", payload.get("protocol_version", PROTOCOL_VERSION)])
        writer.writerow([])

        if "screen_payload" in payload and "screen_results" in payload["screen_payload"]:
            writer.writerow(["--- Phase B Screen Results ---"])
            writer.writerow(["config_id", "description", "val_nll", "val_acc_%", "queries", "sample_evals", "wall_time_sec"])
            table = get_v6_geometry_table()
            s_results = payload["screen_payload"]["screen_results"]
            for cid in sorted(s_results.keys()):
                r = s_results[cid]
                desc = table[cid].description if cid in table else ""
                writer.writerow([
                    cid, desc,
                    r["val_selected_loss"],
                    r["val_selected_acc"],
                    r["total_queries"],
                    r["total_sample_evaluations"],
                    r["wall_time_sec"],
                ])
            writer.writerow([])

        if "confirm_payload" in payload and "confirm_aggregates" in payload["confirm_payload"]:
            writer.writerow(["--- Phase B Confirmation Aggregates ---"])
            writer.writerow(["config_id", "mean_val_acc_%", "std_val_acc", "mean_val_nll", "std_val_nll", "num_seeds"])
            c_aggs = payload["confirm_payload"]["confirm_aggregates"]
            for cid in sorted(c_aggs.keys()):
                agg = c_aggs[cid]
                writer.writerow([
                    cid,
                    agg["val_selected_acc"]["mean"],
                    agg["val_selected_acc"]["std"],
                    agg["val_selected_nll"]["mean"],
                    agg["val_selected_nll"]["std"],
                    agg["num_seeds"],
                ])


def generate_study_plots_v6(payload: Dict[str, Any], plot_path: Path):
    """Generates validation trajectory & comparison figures for V6 study."""
    plot_path.parent.mkdir(parents=True, exist_ok=True)
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    # Left subplot: Validation Trajectories from Screen
    ax_traj = axes[0]
    if "screen_payload" in payload and "screen_results" in payload["screen_payload"]:
        s_results = payload["screen_payload"]["screen_results"]
        for cid, res in s_results.items():
            if "stage_histories" in res and res["stage_histories"]:
                epochs = [h["epoch"] for h in res["stage_histories"] if h.get("val_loss") is not None]
                nlls = [h["val_loss"] for h in res["stage_histories"] if h.get("val_loss") is not None]
                if epochs and nlls:
                    ax_traj.plot(epochs, nlls, label=cid, alpha=0.8)
            elif cid == "G8":
                ax_traj.scatter(
                    [payload["screen_payload"]["epochs"]],
                    [res["val_selected_loss"]],
                    marker="x",
                    s=60,
                    label="G8 final",
                )

    ax_traj.set_title("Screen Validation NLL Trajectory")
    ax_traj.set_xlabel("Epoch")
    ax_traj.set_ylabel("Validation NLL")
    ax_traj.grid(True, linestyle="--", alpha=0.5)
    ax_traj.legend(fontsize=8, loc="upper right")

    # Right subplot: Confirmation Mean Validation Accuracy Bar Chart
    ax_bar = axes[1]
    if "confirm_payload" in payload and "confirm_aggregates" in payload["confirm_payload"]:
        c_aggs = payload["confirm_payload"]["confirm_aggregates"]
        cids = sorted(c_aggs.keys())
        means = [c_aggs[c]["val_selected_acc"]["mean"] for c in cids]
        stds = [c_aggs[c]["val_selected_acc"]["std"] for c in cids]

        colors = ["skyblue" if c != "G8" else "coral" for c in cids]
        ax_bar.bar(cids, means, yerr=stds, capsize=5, color=colors, alpha=0.85)
        ax_bar.set_ylabel("Mean Validation Accuracy (%)")
        ax_bar.set_title("Confirmation Accuracy (3 Seeds)")
        ax_bar.set_ylim(0, 100)
        for i, (m, s) in enumerate(zip(means, stds)):
            ax_bar.text(i, m + s + 1.0, f"{m:.1f}%", ha="center", va="bottom", fontsize=8)

    plt.tight_layout()
    fig.savefig(plot_path, dpi=200)
    plt.close(fig)


# =====================================================================
# 9. Main CLI Entrypoint
# =====================================================================

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="MNIST Deep PSO V6 Root-Cause Study (Phase A & B)")
    parser.add_argument("--phase", type=str, choices=["screen", "confirm", "all"], default="all", help="Phase to execute")
    parser.add_argument("--device", type=str, default=None, help="Execution device (e.g. mps, cuda, cpu)")
    parser.add_argument("--screen-artifact", type=str, default="benchmark_results/pso_v6_phase_b_screen.json", help="Path to screen artifact for confirm phase")
    parser.add_argument("--out-dir", type=str, default="benchmark_results", help="Output directory for results")
    parser.add_argument("--plot-dir", type=str, default="history_plt", help="Output directory for plots")
    parser.add_argument("--override-particles", type=int, default=None, help="Override particle count for testing/smokes")
    parser.add_argument("--override-epochs", type=int, default=None, help="Override epoch count for testing/smokes")
    return parser


def run_deep_pso_v6_study(args: argparse.Namespace) -> Dict[str, Any]:
    device = resolve_execution_device(args.device)
    out_dir = Path(args.out_dir)
    plot_dir = Path(args.plot_dir)

    out_dir.mkdir(parents=True, exist_ok=True)
    plot_dir.mkdir(parents=True, exist_ok=True)

    # 1. Prepare data strictly without test set
    x_search, y_search, x_val, y_val, nested_subsets, data_fp, provenance = prepare_mnist_v6_data()
    hw_prov = get_hardware_provenance(device)
    geometry_table = get_v6_geometry_table()
    base_model = make_compact_cnn(seed=41)

    final_payload: Dict[str, Any] = {
        "protocol_version": PROTOCOL_VERSION,
        "pso_version": pso_version,
        "official_test_data_loaded": False,
        "official_test_evaluations": 0,
        "data_fingerprint": data_fp,
        "base_model_seed": 41,
        "base_model_fingerprint": compute_model_fingerprint(base_model),
        "geometry_configs": {
            config_id: asdict(config)
            for config_id, config in geometry_table.items()
        },
        "selection_rule": "lowest_validation_nll_then_highest_accuracy",
        "provenance": provenance,
        "hardware_provenance": hw_prov,
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
    }

    screen_payload = None
    confirm_payload = None

    screen_particles = args.override_particles if args.override_particles is not None else 30
    screen_epochs = args.override_epochs if args.override_epochs is not None else 160

    confirm_particles = args.override_particles if args.override_particles is not None else 60
    confirm_epochs = args.override_epochs if args.override_epochs is not None else 420

    # 2. Execute Screen Phase if requested or in 'all'
    if args.phase in ("screen", "all"):
        screen_payload = run_phase_b_screen(
            x_search=x_search,
            y_search=y_search,
            x_val=x_val,
            y_val=y_val,
            nested_subsets=nested_subsets,
            device=device,
            seed=91,
            swarm_size=screen_particles,
            epochs=screen_epochs,
        )
        final_payload["screen_payload"] = screen_payload
        save_json_atomic(final_payload, out_dir / "pso_v6_phase_b_screen.json")

    # 3. Execute Confirm Phase if requested or in 'all'
    if args.phase in ("confirm", "all"):
        if screen_payload is not None:
            selected_configs = screen_payload["selected_for_confirm"]
        else:
            screen_art_path = Path(args.screen_artifact)
            if screen_art_path.exists():
                import json
                with open(screen_art_path, "r") as f:
                    art_data = json.load(f)
                selected_configs = art_data.get("screen_payload", {}).get("selected_for_confirm", ["G0", "G1", "G8"])
            else:
                selected_configs = ["G0", "G1", "G8", "G2", "G4"]

        confirm_payload = run_phase_b_confirm(
            selected_configs=selected_configs,
            x_search=x_search,
            y_search=y_search,
            x_val=x_val,
            y_val=y_val,
            nested_subsets=nested_subsets,
            device=device,
            seeds=[101, 102, 103],
            swarm_size=confirm_particles,
            epochs=confirm_epochs,
        )
        final_payload["confirm_payload"] = confirm_payload

    persisted_runs: List[Dict[str, Any]] = []
    if screen_payload is not None:
        persisted_runs.extend(screen_payload["screen_results"].values())
    if confirm_payload is not None:
        for runs in confirm_payload["confirm_runs"].values():
            persisted_runs.extend(runs)
    final_payload["resource_totals"] = {
        "candidate_objective_queries": sum(r["total_queries"] for r in persisted_runs),
        "candidate_sample_evaluations": sum(
            r["total_sample_evaluations"] for r in persisted_runs
        ),
        "validation_model_evaluations": sum(
            r["validation_evaluations"] for r in persisted_runs
        ),
        "summed_optimization_wall_time_sec": round(
            sum(r["optimization_wall_time_sec"] for r in persisted_runs), 4
        ),
        "summed_validation_wall_time_sec": round(
            sum(r["validation_wall_time_sec"] for r in persisted_runs), 4
        ),
        "official_test_evaluations": 0,
    }

    # Save final atomic JSON, CSV, and plots
    json_path = out_dir / "pso_v6_phase_b.json"
    csv_path = out_dir / "pso_v6_phase_b.csv"
    plot_path = plot_dir / "pso_v6_phase_b.png"

    save_json_atomic(final_payload, json_path)
    save_csv_summary_v6(final_payload, csv_path)
    generate_study_plots_v6(final_payload, plot_path)

    return final_payload


if __name__ == "__main__":
    parser = build_parser()
    args = parser.parse_args()
    run_deep_pso_v6_study(args)