PSO/test/deep_pso_methods.py

"""
MNIST Deep PSO Methods Study: Latent Space Subspaces, Adaptive-Moment PSO, and Ensembling

Protocol: MNIST-PSO-RAW-V5 1.0.0
- Official raw MNIST (60,000 train / 10,000 test).
- Split 60k train first into 50k search and 10k validation using deterministic stratified sampling (seed 20260902).
- Mean and std fit on 50k search subset ONLY; applied to search, validation, and test.
- Nested stratified ordering: 2k subset inside 10k subset inside 50k search set.
- Base CompactCNN (9,098 parameters).
- Latent subspace transform via deterministic sparse signed hash mapping for d in [290, 1024, 4096, full].
- Exact base model at particle 0; remaining particles in antithetic pairs.
- Device-resident latent adaptive-moment PSO (c0=c1=1.49618, w=0.7298, blend=0.06, step=0.5, beta1=0.9, beta2=0.999).
- Lexicographical CE loss primary, accuracy tie-break selection.
- Objective transition: complete pbest re-evaluation and gbest rebuild.
- Validation-only pilot selection and elite selection.
- Official 10k test set evaluated exactly once per final reported endpoint.
"""

import argparse
import csv
import datetime
import hashlib
import json
import math
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
from sklearn.model_selection import train_test_split

sys.path.insert(0, str(Path(__file__).resolve().parent))

from benchmark_suite import (
    calc_stats,
    compute_data_fingerprint,
    compute_model_fingerprint,
    get_hardware_provenance,
    resolve_execution_device,
    save_json_atomic,
    sync_device,
)
from pso import __version__ as pso_version

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


# =====================================================================
# 1. Architecture: CompactCNN (9,098 Parameters)
# =====================================================================

class CompactCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 8, kernel_size=3, padding=1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(8, 16, kernel_size=3, padding=1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2, 2)
        self.flatten = nn.Flatten()
        self.fc = nn.Linear(784, 10)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.dim() == 2 and x.shape[1] == 784:
            x = x.view(-1, 1, 28, 28)
        out = self.pool1(self.relu1(self.conv1(x)))
        out = self.pool2(self.relu2(self.conv2(out)))
        out = self.flatten(out)
        return self.fc(out)


def make_compact_cnn(seed: int = 41) -> nn.Module:
    torch.manual_seed(seed)
    return CompactCNN()


# =====================================================================
# 2. Data Split, Normalization, & Stratified Subsets
# =====================================================================

def prepare_mnist_v5_data(
    split_seed: int = 20260902,
    cache_dir: Optional[Path] = None,
) -> Tuple[
    torch.Tensor, torch.Tensor,
    torch.Tensor, torch.Tensor,
    torch.Tensor, torch.Tensor,
    Dict[int, torch.Tensor],
    str, Dict[str, Any]
]:
    from torchvision.datasets import MNIST

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

    raw_train = MNIST(root=str(cache_dir), train=True, download=True)
    raw_test = MNIST(root=str(cache_dir), train=False, download=True)

    x_train_raw = raw_train.data.float() / 255.0  # (60000, 28, 28)
    y_train_raw = raw_train.targets.long()
    x_test_raw = raw_test.data.float() / 255.0    # (10000, 28, 28)
    y_test_raw = raw_test.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)
    x_test_norm = ((x_test_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 test
    data_fp = compute_data_fingerprint(x_search_norm, x_test_norm, y_search, y_test_raw)

    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": 10000,
        "split_seed": split_seed,
        "split_fingerprint": split_fp,
    }

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


def build_nested_stratified_subsets(
    y_search: torch.Tensor,
    subset_sizes: List[int],
    subset_seed: int = 20260902,
) -> Dict[int, torch.Tensor]:
    """
    Builds nested stratified index tensors: I_2k subset of I_10k subset of I_50k.
    Uses Hamilton / Largest-Remainder Method for exact subset sizes and nesting.
    """
    rng = np.random.RandomState(subset_seed)
    y_np = y_search.numpy()
    total_samples = len(y_np)
    unique_classes, counts = np.unique(y_np, return_counts=True)

    class_indices = {}
    for c in unique_classes:
        c_idxs = np.where(y_np == c)[0]
        rng.shuffle(c_idxs)
        class_indices[c] = c_idxs

    ordered_subset_sizes = sorted(subset_sizes)
    nested_subsets: Dict[int, torch.Tensor] = {}
    selected_per_class: Dict[int, List[int]] = {c: [] for c in unique_classes}

    for size in ordered_subset_sizes:
        if size == total_samples:
            nested_subsets[size] = torch.arange(total_samples, dtype=torch.long)
            continue

        exact_quotas = [size * (counts[i] / total_samples) for i in range(len(unique_classes))]
        floor_quotas = [int(np.floor(q)) for q in exact_quotas]
        remainders = [exact_quotas[i] - floor_quotas[i] for i in range(len(unique_classes))]

        needed_extra = size - sum(floor_quotas)
        ranked_indices = np.argsort(remainders)[::-1]
        target_counts = list(floor_quotas)
        for i in range(needed_extra):
            target_counts[ranked_indices[i]] += 1

        target_subset_idxs = []
        for i, c in enumerate(unique_classes):
            target_c_count = target_counts[i]
            current_list = selected_per_class[c]
            needed = target_c_count - len(current_list)
            if needed > 0:
                available = class_indices[c]
                added = list(available[len(current_list):len(current_list) + needed])
                current_list.extend(added)
            target_subset_idxs.extend(current_list[:target_c_count])

        subset_tensor = torch.tensor(sorted(target_subset_idxs), dtype=torch.long)
        nested_subsets[size] = subset_tensor

    return nested_subsets


# =====================================================================
# 3. Latent Subspace Transform & Antithetic Swarm Construction
# =====================================================================

class LatentTransform:
    def __init__(
        self,
        base_model: nn.Module,
        latent_dim: Union[int, str],
        device: torch.device,
    ):
        self.device = device
        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)

        # Per-tensor scale calculation: positive scale per parameter tensor
        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)
        self.base_vec = torch.cat([p.view(-1) for p in self.base_params]).to(device)

        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:
            # Deterministic sparse signed hash mapping
            j_indices = np.arange(self.total_dim, dtype=np.int64)
            h1 = ((j_indices + 1) * 2654435761) % (2**32)
            k_indices = h1 % self.latent_dim
            h2 = ((j_indices + 1) * 1597334677) % (2**32)
            signs = np.where((h2 % 2) == 0, 1.0, -1.0)

            # Count normalization to maintain unit variance
            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, init_radius: float = 0.5) -> torch.Tensor:
        """
        Initializes particle positions in latent space Z (N, d).
        Particle 0 is exact base vector (z = 0).
        For even swarm sizes N, particles 1..N-2 form (N-2)//2 exact pairs, and particle N-1 is zero.
        """
        rng = torch.Generator(device="cpu")
        rng.manual_seed(seed)

        Z = torch.zeros((swarm_size, self.latent_dim), dtype=torch.float32)
        # Particle 0 stays exact 0

        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

        return Z.to(self.device)


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

def evaluate_latent_batch(
    Z: torch.Tensor,
    transform: LatentTransform,
    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) without host roundtrips.
    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_latent_pso(
    transform: LatentTransform,
    base_model: nn.Module,
    x_search: torch.Tensor,
    y_search: torch.Tensor,
    nested_subsets: Dict[int, torch.Tensor],
    schedule_str: str,
    epochs: int,
    swarm_size: int,
    seed: int,
    device: torch.device,
) -> Dict[str, Any]:
    """
    Runs device-resident Latent Adaptive-Moment PSO with progressive schedule.
    Performs full pbest re-evaluation and gbest rebuild ONLY on objective transitions (stages > 0).
    """
    sync_device(device)
    start_time = time.time()

    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 = 3.0

    latent_dim = transform.latent_dim
    Z = transform.init_swarm(swarm_size=swarm_size, seed=seed)
    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)

    # State tracking: P scores start at inf / 0. DO NOT evaluate before stage 0!
    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

    total_queries = 0
    total_sample_evaluations = 0
    transition_reevaluation_counts = 0

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

    model = make_compact_cnn(seed=41).to(device)

    stage_histories = []
    t_step = 0
    epoch_counter = 0

    for stage_idx, (size, stage_epochs) in enumerate(schedule_stages):
        subset_indices = nested_subsets[size]
        # Move stage subset to device ONCE per stage
        x_sub_dev = x_search[subset_indices].to(device)
        y_sub_dev = y_search[subset_indices].to(device)

        # Objective transition check: ONLY for stages >= 1!
        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())

        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 on-device pbest updates (lexicographical loss primary, acc tie-break)
            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)

            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 on-device
            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())

            stage_histories.append({
                "epoch": epoch_counter,
                "stage": stage_idx,
                "subset_size": size,
                "gbest_loss": round(gbest_loss, 6),
                "gbest_acc": round(gbest_acc, 4),
            })

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

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

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

            M_hat = M / (1.0 - beta1 ** t_step)
            V_sq_hat = V_sq / (1.0 - beta2 ** t_step)

            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

            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)
    wall_time = time.time() - start_time

    return {
        "gbest_z": gbest_z,
        "gbest_loss": gbest_loss,
        "gbest_acc": gbest_acc,
        "final_P": P,
        "wall_time_sec": round(wall_time, 4),
        "total_queries": total_queries,
        "total_sample_evaluations": total_sample_evaluations,
        "transition_reevaluation_counts": transition_reevaluation_counts,
        "stage_histories": stage_histories,
    }


# =====================================================================
# 5. Evaluation Metrics (NLL, Brier, ECE, Margin, Disagreement)
# =====================================================================

def evaluate_probabilistic_metrics(
    prob_matrix: torch.Tensor,
    y_true: torch.Tensor,
) -> Dict[str, float]:
    """
    Computes accuracy, NLL, Brier score, 15-bin ECE, and probability margin.
    """
    N, C = prob_matrix.shape
    probs = prob_matrix.cpu().numpy()
    labels = y_true.cpu().numpy()

    preds = probs.argmax(axis=1)
    acc = float((preds == labels).mean()) * 100.0

    eps = 1e-12
    clipped_probs = np.clip(probs, eps, 1.0 - eps)
    nll = -float(np.log(clipped_probs[np.arange(N), labels]).mean())

    y_onehot = np.zeros((N, C), dtype=np.float32)
    y_onehot[np.arange(N), labels] = 1.0
    brier = float(np.mean(np.sum((probs - y_onehot) ** 2, axis=1)))

    n_bins = 15
    bin_boundaries = np.linspace(0.0, 1.0, n_bins + 1)
    confidences = probs.max(axis=1)
    ece = 0.0

    for i in range(n_bins):
        bin_lower = bin_boundaries[i]
        bin_upper = bin_boundaries[i + 1]
        in_bin = (confidences > bin_lower) & (confidences <= bin_upper) if i > 0 else (confidences >= bin_lower) & (confidences <= bin_upper)
        prop_in_bin = in_bin.mean()
        if prop_in_bin > 0:
            accuracy_in_bin = (preds[in_bin] == labels[in_bin]).mean()
            avg_confidence_in_bin = confidences[in_bin].mean()
            ece += np.abs(accuracy_in_bin - avg_confidence_in_bin) * prop_in_bin

    sorted_probs = np.sort(probs, axis=1)[:, ::-1]
    margins = sorted_probs[:, 0] - sorted_probs[:, 1]
    margin_mean = float(margins.mean())

    return {
        "accuracy": round(acc, 4),
        "nll": round(nll, 6),
        "brier": round(brier, 6),
        "ece": round(float(ece), 6),
        "margin": round(margin_mean, 6),
    }


def compute_pairwise_disagreement(model_preds_list: List[np.ndarray]) -> float:
    num_models = len(model_preds_list)
    if num_models < 2:
        return 0.0

    disagreements = []
    for i in range(num_models):
        for j in range(i + 1, num_models):
            dis = float((model_preds_list[i] != model_preds_list[j]).mean())
            disagreements.append(dis)

    return round(float(np.mean(disagreements)), 6)


def select_diverse_candidates(
    candidates: List[Dict[str, Any]],
    *,
    max_size: int = 5,
    accuracy_window: float = 2.0,
) -> List[Dict[str, Any]]:
    if not candidates:
        raise ValueError("candidates must not be empty")
    if max_size <= 0:
        raise ValueError("max_size must be positive")
    if accuracy_window < 0.0:
        raise ValueError("accuracy_window must be nonnegative")

    ranked = sorted(candidates, key=lambda c: (c["val_loss"], -c["val_acc"]))
    accuracy_threshold = ranked[0]["val_acc"] - accuracy_window
    eligible = [c for c in ranked if c["val_acc"] >= accuracy_threshold]

    selected = [eligible[0]]
    selected_keys = {(eligible[0]["seed"], eligible[0]["particle_idx"])}
    predictions = {
        (candidate["seed"], candidate["particle_idx"]):
        candidate["val_probs"].argmax(dim=1).cpu().numpy()
        for candidate in eligible
    }

    while len(selected) < min(max_size, len(eligible)):
        best_next = None
        best_key = None
        max_disagreement = -1.0
        best_val_loss = float("inf")

        for candidate in eligible:
            candidate_key = (candidate["seed"], candidate["particle_idx"])
            if candidate_key in selected_keys:
                continue
            candidate_predictions = predictions[candidate_key]
            mean_disagreement = float(np.mean([
                (candidate_predictions - predictions[
                    (member["seed"], member["particle_idx"])
                ] != 0).mean()
                for member in selected
            ]))
            if (
                mean_disagreement > max_disagreement + 1e-7
                or (
                    abs(mean_disagreement - max_disagreement) <= 1e-7
                    and candidate["val_loss"] < best_val_loss
                )
            ):
                best_next = candidate
                best_key = candidate_key
                max_disagreement = mean_disagreement
                best_val_loss = candidate["val_loss"]

        if best_next is None or best_key is None:
            break
        selected.append(best_next)
        selected_keys.add(best_key)

    return selected


# =====================================================================
# 6. Full Experiment Pipeline & CLI Runner
# =====================================================================

def validate_cli_args(args: argparse.Namespace):
    if args.pilot_epochs <= 0 or args.confirmation_epochs <= 0:
        raise ValueError("Pilot and confirmation epochs must be positive integers.")
    if args.pilot_particles <= 0 or args.confirmation_particles <= 0:
        raise ValueError("Pilot and confirmation particles must be positive integers.")
    if any(s < 0 for s in args.seeds):
        raise ValueError("Seeds must be non-negative integers.")
    if len(args.seeds) != len(set(args.seeds)):
        raise ValueError("Confirmation seeds must be unique.")

    max_d = 9098
    for d_str in args.dimensions:
        if d_str.lower() != "full":
            try:
                d_val = int(d_str)
                if d_val <= 0 or d_val > max_d:
                    raise ValueError(f"Latent dimension {d_val} must be in range [1, {max_d}].")
            except ValueError:
                raise ValueError(f"Invalid dimension specifier: {d_str}")

    if args.confirmation_schedule:
        stages = args.confirmation_schedule.split(",")
        total_sched_epochs = 0
        known_sizes = {2000, 10000, 50000}
        for s in stages:
            sz_str, ep_str = s.split(":")
            sz, ep = int(sz_str), int(ep_str)
            if sz not in known_sizes:
                raise ValueError(f"Unknown schedule subset size {sz}; known sizes are {known_sizes}.")
            if ep <= 0:
                raise ValueError(f"Schedule epochs must be positive; got {ep}.")
            total_sched_epochs += ep
        if total_sched_epochs != args.confirmation_epochs:
            raise ValueError(
                f"Schedule epoch sum ({total_sched_epochs}) must equal confirmation_epochs ({args.confirmation_epochs})."
            )


def run_deep_pso_study(args: argparse.Namespace) -> Dict[str, Any]:
    validate_cli_args(args)

    device = resolve_execution_device(args.device)
    print(f"=== MNIST Deep PSO Methods Study (Protocol {PROTOCOL_VERSION}) ===")
    print(f"Device: {device}")

    # Data preparation
    (
        x_search, y_search,
        x_val, y_val,
        x_test, y_test,
        nested_subsets,
        data_fp, provenance
    ) = prepare_mnist_v5_data(split_seed=args.split_seed)

    base_model = make_compact_cnn(seed=41).to(device)
    base_fp = compute_model_fingerprint(base_model)
    hardware_prov = get_hardware_provenance(device)

    failure_record: Optional[str] = None

    # Pilot Phase: Validation-only selection of latent dimension
    print("\n--- Pilot Phase (Dimension Selection on 10k Validation Set) ---")
    pilot_results = []
    best_pilot_dim = None
    best_pilot_val_loss = float("inf")
    best_pilot_val_acc = 0.0

    dimensions = args.dimensions

    for dim_str in dimensions:
        print(f"Running Pilot: dim={dim_str}, seed={args.pilot_seed}, particles={args.pilot_particles}, epochs={args.pilot_epochs}")
        transform = LatentTransform(base_model, latent_dim=dim_str, device=device)
        res = run_latent_pso(
            transform=transform,
            base_model=base_model,
            x_search=x_search,
            y_search=y_search,
            nested_subsets=nested_subsets,
            schedule_str=f"2000:{args.pilot_epochs}",
            epochs=args.pilot_epochs,
            swarm_size=args.pilot_particles,
            seed=args.pilot_seed,
            device=device,
        )

        model = make_compact_cnn(seed=41).to(device)
        transform.load_vector_to_model(transform.decode(res["gbest_z"].unsqueeze(0)).squeeze(0), model)

        probs_val = get_model_probabilities(model, x_val, device)
        val_metrics = evaluate_probabilistic_metrics(probs_val, y_val)

        pilot_record = {
            "dimension": str(dim_str),
            "val_loss": val_metrics["nll"],
            "val_acc": val_metrics["accuracy"],
            "wall_time_sec": res["wall_time_sec"],
            "queries": res["total_queries"],
            "sample_evaluations": res["total_sample_evaluations"],
        }
        pilot_results.append(pilot_record)
        print(f"Pilot dim={dim_str} -> Val Loss: {val_metrics['nll']:.6f}, Val Acc: {val_metrics['accuracy']:.2f}%")

        if (val_metrics["nll"] < best_pilot_val_loss - 1e-7) or (abs(val_metrics["nll"] - best_pilot_val_loss) <= 1e-7 and val_metrics["accuracy"] > best_pilot_val_acc):
            best_pilot_val_loss = val_metrics["nll"]
            best_pilot_val_acc = val_metrics["accuracy"]
            best_pilot_dim = dim_str

    print(f"\nSelected Pilot Dimension: {best_pilot_dim} (Val Loss: {best_pilot_val_loss:.6f}, Val Acc: {best_pilot_val_acc:.2f}%)")

    # Confirmation Phase: Run multi-seed progressive PSO on selected dimension
    print(f"\n--- Confirmation Phase (Dimension={best_pilot_dim}, Seeds={args.seeds}) ---")
    confirmation_runs = []
    val_candidate_pool = []

    conf_transform = LatentTransform(base_model, latent_dim=best_pilot_dim, device=device)

    for c_seed in args.seeds:
        print(f"Running Confirmation: seed={c_seed}, particles={args.confirmation_particles}, epochs={args.confirmation_epochs}, schedule={args.confirmation_schedule}")
        res = run_latent_pso(
            transform=conf_transform,
            base_model=base_model,
            x_search=x_search,
            y_search=y_search,
            nested_subsets=nested_subsets,
            schedule_str=args.confirmation_schedule,
            epochs=args.confirmation_epochs,
            swarm_size=args.confirmation_particles,
            seed=c_seed,
            device=device,
        )

        # Evaluate EVERY particle's pbest on the 10k VALIDATION set
        P_final = res["final_P"]
        run_best_val_loss = float("inf")
        run_best_val_acc = 0.0

        for p_idx in range(len(P_final)):
            p_z = P_final[p_idx]
            cand_model = make_compact_cnn(seed=41).to(device)
            conf_transform.load_vector_to_model(conf_transform.decode(p_z.unsqueeze(0)).squeeze(0), cand_model)
            probs_val = get_model_probabilities(cand_model, x_val, device)
            val_metrics = evaluate_probabilistic_metrics(probs_val, y_val)

            cand_rec = {
                "seed": c_seed,
                "particle_idx": p_idx,
                "val_loss": val_metrics["nll"],
                "val_acc": val_metrics["accuracy"],
                "val_probs": probs_val,
                "latent_z": p_z,
            }
            val_candidate_pool.append(cand_rec)

            if (val_metrics["nll"] < run_best_val_loss - 1e-7) or (abs(val_metrics["nll"] - run_best_val_loss) <= 1e-7 and val_metrics["accuracy"] > run_best_val_acc):
                run_best_val_loss = val_metrics["nll"]
                run_best_val_acc = val_metrics["accuracy"]

        conf_record = {
            "seed": c_seed,
            "val_loss": run_best_val_loss,
            "val_acc": run_best_val_acc,
            "wall_time_sec": res["wall_time_sec"],
            "queries": res["total_queries"],
            "sample_evaluations": res["total_sample_evaluations"],
            "transition_reevaluations": res["transition_reevaluation_counts"],
            "stage_histories": res["stage_histories"],
        }
        confirmation_runs.append(conf_record)

        print(f"Confirmation seed={c_seed} -> Best Val Loss: {run_best_val_loss:.6f}, Best Val Acc: {run_best_val_acc:.2f}%")

    # Selection on Validation ONLY:
    # 1. Single Final Model (lowest val NLL, then highest val Acc)
    val_candidate_pool.sort(key=lambda c: (c["val_loss"], -c["val_acc"]))
    best_single_candidate = val_candidate_pool[0]

    # 2. Predeclared validation-performing, prediction-diverse Top-5 Ensemble
    top_ensemble_candidates = select_diverse_candidates(val_candidate_pool)

    # Official 10k Test Set Evaluation (EXACTLY ONCE PER ENDPOINT)
    print("\n--- Final Official 10k Test Evaluation ---")

    # Single Final Model Test Evaluation
    single_model = make_compact_cnn(seed=41).to(device)
    conf_transform.load_vector_to_model(conf_transform.decode(best_single_candidate["latent_z"].unsqueeze(0)).squeeze(0), single_model)

    single_test_probs = get_model_probabilities(single_model, x_test, device)
    single_test_metrics = evaluate_probabilistic_metrics(single_test_probs, y_test)
    single_model_fp = compute_model_fingerprint(single_model)

    print(f"Final Single Model (Seed {best_single_candidate['seed']}, Part {best_single_candidate['particle_idx']}) -> Test Acc: {single_test_metrics['accuracy']:.2f}%, Test NLL: {single_test_metrics['nll']:.6f}")

    # Top Ensemble Test Evaluation
    ensemble_test_probs_list = []
    ensemble_preds_list = []

    for cand in top_ensemble_candidates:
        cand_model = make_compact_cnn(seed=41).to(device)
        conf_transform.load_vector_to_model(conf_transform.decode(cand["latent_z"].unsqueeze(0)).squeeze(0), cand_model)
        t_probs = get_model_probabilities(cand_model, x_test, device)
        ensemble_test_probs_list.append(t_probs)
        ensemble_preds_list.append(t_probs.argmax(dim=1).cpu().numpy())

    ensemble_mean_probs = torch.stack(ensemble_test_probs_list).mean(dim=0)
    ensemble_test_metrics = evaluate_probabilistic_metrics(ensemble_mean_probs, y_test)
    ensemble_disagreement = compute_pairwise_disagreement(ensemble_preds_list)

    print(f"Top-{len(top_ensemble_candidates)} Ensemble -> Test Acc: {ensemble_test_metrics['accuracy']:.2f}%, Test NLL: {ensemble_test_metrics['nll']:.6f}, Disagreement: {ensemble_disagreement:.6f}")
    ensemble_val_probs = torch.stack(
        [cand["val_probs"] for cand in top_ensemble_candidates]
    ).mean(dim=0)
    ensemble_val_metrics = evaluate_probabilistic_metrics(
        ensemble_val_probs, y_val
    )
    ensemble_val_disagreement = compute_pairwise_disagreement(
        [
            cand["val_probs"].argmax(dim=1).cpu().numpy()
            for cand in top_ensemble_candidates
        ]
    )


    total_wall_time = sum(c["wall_time_sec"] for c in confirmation_runs) + sum(p["wall_time_sec"] for p in pilot_results)
    total_queries_all = sum(c["queries"] for c in confirmation_runs) + sum(p["queries"] for p in pilot_results)
    total_samples_all = sum(c["sample_evaluations"] for c in confirmation_runs) + sum(p["sample_evaluations"] for p in pilot_results)

    final_payload = {
        "protocol_version": PROTOCOL_VERSION,
        "pso_version": pso_version,
        "timestamp": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "completed": True,
        "hardware_provenance": hardware_prov,
        "data_provenance": provenance,
        "data_fingerprint": data_fp,
        "base_model_fingerprint": base_fp,
        "configuration": {
            "base_model_seed": 41,
            "split_seed": args.split_seed,
            "pilot_seed": args.pilot_seed,
            "pilot_dimensions": [str(dim) for dim in args.dimensions],
            "pilot_particles": args.pilot_particles,
            "pilot_epochs": args.pilot_epochs,
            "confirmation_seeds": args.seeds,
            "confirmation_particles": args.confirmation_particles,
            "confirmation_epochs": args.confirmation_epochs,
            "confirmation_schedule": args.confirmation_schedule,
            "fitness_objective": "cross_entropy_loss_primary_accuracy_tiebreak",
            "parameterization": {
                "layer_scale": "per_parameter_tensor_std_floor_1e-4",
                "subspace": "deterministic_sparse_signed_hash_count_normalized",
                "initialization": "exact_base_plus_antithetic",
                "initial_radius": 0.5,
                "reflective_bound": 3.0,
            },
            "movement": {
                "name": "latent_adaptive_moment_pso",
                "c0": 1.49618,
                "c1": 1.49618,
                "w": 0.7298,
                "moment_blend": 0.06,
                "moment_step_size": 0.5,
                "moment_beta1": 0.9,
                "moment_beta2": 0.999,
            },
            "objective_transition": "reevaluate_all_pbests_then_rebuild_gbest",
            "validation_selection": {
                "single": "lowest_nll_then_highest_accuracy",
                "ensemble": "within_2_accuracy_points_then_greedy_disagreement",
                "ensemble_size": 5,
            },
            "fitness_batch_size": 1000,
            "ece_bins": 15,
        },
        "pilot_phase": {
            "selected_dimension": str(best_pilot_dim),
            "results": pilot_results,
        },
        "confirmation_phase": {
            "runs": [
                {
                    "seed": r["seed"],
                    "val_loss": r["val_loss"],
                    "val_acc": r["val_acc"],
                    "wall_time_sec": r["wall_time_sec"],
                    "queries": r["queries"],
                    "sample_evaluations": r["sample_evaluations"],
                    "transition_reevaluations": r["transition_reevaluations"],
                    "stage_histories": r["stage_histories"],
                }
                for r in confirmation_runs
            ],
            "validation_summary": {
                "best_pbest_nll": calc_stats(
                    [run["val_loss"] for run in confirmation_runs]
                ),
                "best_pbest_accuracy": calc_stats(
                    [run["val_acc"] for run in confirmation_runs]
                ),
                "wall_time_sec": calc_stats(
                    [run["wall_time_sec"] for run in confirmation_runs]
                ),
            },
        },
        "final_endpoints": {
            "single_model": {
                "selected_seed": best_single_candidate["seed"],
                "selected_particle_idx": best_single_candidate["particle_idx"],
                "model_fingerprint": single_model_fp,
                "val_loss": best_single_candidate["val_loss"],
                "val_acc": best_single_candidate["val_acc"],
                "selection_rule": "lowest_validation_nll_then_highest_accuracy",
                "test_accuracy": single_test_metrics["accuracy"],
                "test_nll": single_test_metrics["nll"],
                "test_brier": single_test_metrics["brier"],
                "test_ece": single_test_metrics["ece"],
                "test_margin": single_test_metrics["margin"],
            },
            "ensemble": {
                "ensemble_size": len(top_ensemble_candidates),
                "members": [
                    {
                        "seed": cand["seed"],
                        "particle_idx": cand["particle_idx"],
                        "val_loss": cand["val_loss"],
                        "val_acc": cand["val_acc"],
                    }
                    for cand in top_ensemble_candidates
                ],
                "validation_accuracy": ensemble_val_metrics["accuracy"],
                "validation_nll": ensemble_val_metrics["nll"],
                "validation_brier": ensemble_val_metrics["brier"],
                "validation_ece": ensemble_val_metrics["ece"],
                "validation_margin": ensemble_val_metrics["margin"],
                "validation_pairwise_disagreement": ensemble_val_disagreement,
                "test_accuracy": ensemble_test_metrics["accuracy"],
                "test_nll": ensemble_test_metrics["nll"],
                "test_brier": ensemble_test_metrics["brier"],
                "test_ece": ensemble_test_metrics["ece"],
                "test_margin": ensemble_test_metrics["margin"],
                "pairwise_disagreement": ensemble_disagreement,
            },
        },
        "accounting": {
            "total_wall_time_sec": round(total_wall_time, 4),
            "total_queries": total_queries_all,
            "total_sample_evaluations": total_samples_all,
            "scope": (
                "pilot_and_confirmation_training_objectives_including_"
                "transition_reevaluations; excludes validation and test"
            ),
        },
        "failure_record": failure_record,
    }

    # Save atomic JSON
    json_path = Path(args.json_path)
    save_json_atomic(final_payload, json_path)
    print(f"Saved atomic JSON to {json_path}")

    # Save CSV
    csv_path = Path(args.csv_path)
    save_csv_summary(final_payload, csv_path)
    print(f"Saved CSV summary to {csv_path}")

    # Save PNG plot
    plot_path = Path(args.plot_path)
    generate_study_plots(final_payload, plot_path)
    print(f"Saved PNG plot to {plot_path}")

    return final_payload


def get_model_probabilities(model: nn.Module, x_data: torch.Tensor, device: torch.device, batch_size: int = 1000) -> torch.Tensor:
    model.eval()
    prob_list = []
    num_samples = len(x_data)
    with torch.inference_mode():
        for b_start in range(0, num_samples, batch_size):
            xb = x_data[b_start:b_start + batch_size].to(device)
            logits = model(xb)
            probs = torch.softmax(logits, dim=1)
            prob_list.append(probs)
    return torch.cat(prob_list, dim=0)


def save_csv_summary(payload: Dict[str, Any], csv_path: Path):
    csv_path.parent.mkdir(parents=True, exist_ok=True)
    with open(csv_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["section", "metric", "value"])
        writer.writerow(["protocol", "version", payload["protocol_version"]])

        pilot = payload["pilot_phase"]
        writer.writerow(["pilot", "selected_dimension", pilot["selected_dimension"]])

        single = payload["final_endpoints"]["single_model"]
        writer.writerow(["single_model", "test_accuracy", single["test_accuracy"]])
        writer.writerow(["single_model", "test_nll", single["test_nll"]])
        writer.writerow(["single_model", "test_brier", single["test_brier"]])
        writer.writerow(["single_model", "test_ece", single["test_ece"]])

        ens = payload["final_endpoints"]["ensemble"]
        writer.writerow(["ensemble", "test_accuracy", ens["test_accuracy"]])
        writer.writerow(["ensemble", "test_nll", ens["test_nll"]])
        writer.writerow(["ensemble", "pairwise_disagreement", ens["pairwise_disagreement"]])


def generate_study_plots(payload: Dict[str, Any], plot_path: Path):
    plot_path.parent.mkdir(parents=True, exist_ok=True)
    fig, axes = plt.subplots(1, 3, figsize=(18, 5))

    # Panel 1: Pilot Dimension Selection
    pilot_results = payload["pilot_phase"]["results"]
    dims = [r["dimension"] for r in pilot_results]
    val_losses = [r["val_loss"] for r in pilot_results]
    val_accs = [r["val_acc"] for r in pilot_results]

    ax1 = axes[0]
    ax1_twin = ax1.twinx()
    b1 = ax1.bar(np.arange(len(dims)) - 0.2, val_losses, width=0.4, color="#56B4E9", label="Val NLL")
    b2 = ax1_twin.bar(np.arange(len(dims)) + 0.2, val_accs, width=0.4, color="#009E73", label="Val Acc (%)")
    ax1.set_xticks(range(len(dims)))
    ax1.set_xticklabels(dims)
    ax1.set_xlabel("Latent Dimension")
    ax1.set_ylabel("Validation NLL")
    ax1_twin.set_ylabel("Validation Accuracy (%)")
    ax1.set_title("Pilot Dimension Selection")

    # Panel 2: Confirmation Training Histories
    ax2 = axes[1]
    conf_runs = payload["confirmation_phase"]["runs"]
    for run in conf_runs:
        hist = run["stage_histories"]
        epochs = [h["epoch"] for h in hist]
        losses = [h["gbest_loss"] for h in hist]
        ax2.plot(epochs, losses, label=f"Seed {run['seed']}")
    ax2.set_xlabel("Epoch")
    ax2.set_ylabel("Gbest CE Loss")
    ax2.set_title("Confirmation Stage Training Histories")
    ax2.legend()
    ax2.grid(True, linestyle="--", alpha=0.5)

    # Panel 3: Final Endpoint Comparison
    ax3 = axes[2]
    single_acc = payload["final_endpoints"]["single_model"]["test_accuracy"]
    ens_acc = payload["final_endpoints"]["ensemble"]["test_accuracy"]
    single_nll = payload["final_endpoints"]["single_model"]["test_nll"]
    ens_nll = payload["final_endpoints"]["ensemble"]["test_nll"]

    x_labels = ["Single Model", "Top-5 Ensemble"]
    accs = [single_acc, ens_acc]
    nlls = [single_nll, ens_nll]

    ax3_twin = ax3.twinx()
    ax3.bar(np.arange(2) - 0.15, accs, width=0.3, color="#CC79A7", label="Test Acc (%)")
    ax3_twin.bar(np.arange(2) + 0.15, nlls, width=0.3, color="#D55E00", label="Test NLL")
    ax3.set_xticks(range(2))
    ax3.set_xticklabels(x_labels)
    ax3.set_ylabel("Test Accuracy (%)")
    ax3_twin.set_ylabel("Test NLL")
    ax3.set_title("Final Official Test Endpoints")

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


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="MNIST Deep PSO Methods Study (V5)")
    parser.add_argument("--pilot-epochs", type=int, default=160)
    parser.add_argument("--pilot-particles", type=int, default=30)
    parser.add_argument("--pilot-seed", type=int, default=91)
    parser.add_argument("--split-seed", type=int, default=20260902)
    parser.add_argument("--confirmation-epochs", type=int, default=600)
    parser.add_argument("--confirmation-particles", type=int, default=60)
    parser.add_argument("--confirmation-schedule", type=str, default="2000:420,10000:135,50000:45")
    parser.add_argument("--seeds", nargs="+", type=int, default=[101, 102, 103])
    parser.add_argument("--dimensions", nargs="+", type=str, default=["290", "1024", "4096", "full"])
    parser.add_argument("--device", type=str, default=None)
    parser.add_argument("--json-path", type=str, default="benchmark_results/pso_v5_deep_methods.json")
    parser.add_argument("--csv-path", type=str, default="benchmark_results/pso_v5_deep_methods.csv")
    parser.add_argument("--plot-path", type=str, default="history_plt/pso_v5_deep_methods.png")
    return parser


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