#!/usr/bin/env python3
"""Visualise how strongly metadata drives the learned embedding space.

This script mirrors the functionality of ``task1/plot_mod_tsne.py`` but groups
spectrograms by their SNR folder name (e.g. ``SNR0dB``) instead of modulation.
It is useful for checking whether the self-supervised LWM backbone mostly
captures channel/SNR differences rather than modulation characteristics.

Pass ``--label-field modulation`` to reuse the same sampled spectrograms while
colouring and scoring them by their modulation folder instead of SNR. Use
``--label-field mobility`` to highlight link-level mobility categories when
present in the dataset tree. Saved figures automatically include the detected
communication profile (e.g. LTE/WiFi/5G) and label mode in the filename when
those suffixes are not already present.

Usage example:

```bash
python task1/plot_snr_tsne.py \
    --data-root spectrograms/city_1_losangeles/LTE \
    --snrs SNR-5dB,SNR0dB,SNR10dB,SNR15dB,SNR20dB,SNR25dB \
    --save-path task1/snr_separation_plot_latest.png
```
Shortcut presets:

```bash
python task1/plot_snr_tsne.py --WiFi --report-metrics
```
"""

from __future__ import annotations

import argparse
import glob
import pickle
import random
import re
from pathlib import Path
from collections import Counter, defaultdict
from typing import Dict, Iterable, List, Tuple

import matplotlib.pyplot as plt
import numpy as np
import torch
from sklearn.manifold import TSNE
from sklearn.metrics import silhouette_score
from sklearn.model_selection import StratifiedKFold
from sklearn.neighbors import KNeighborsClassifier
from sklearn.preprocessing import StandardScaler

from pretraining.pretrained_model import lwm as lwm_model
from utils import load_spectrogram_data  # support .mat and .pkl uniformly


DEFAULT_DATA_ROOT = "spectrograms/city_1_losangeles/LTE"
DEFAULT_MODELS_ROOT = "models/LTE_models"

PROFILE_PRESETS: Dict[str, Dict[str, str]] = {
    "LTE": {
        "data_root": DEFAULT_DATA_ROOT,
        "models_root": DEFAULT_MODELS_ROOT,
    },
    "WiFi": {
        "data_root": "spectrograms/city_1_losangeles/WiFi",
        "models_root": "models/WiFi_models",
    },
    "5G": {
        "data_root": "spectrograms/city_1_losangeles/5G",
        "models_root": "models/5G_models",
    },
}


def normalize_per_sample(specs: np.ndarray, eps: float = 1e-6) -> np.ndarray:
    means = specs.mean(axis=(1, 2), keepdims=True)
    stds = specs.std(axis=(1, 2), keepdims=True)
    stds = np.maximum(stds, eps)
    return ((specs - means) / stds).astype(np.float32, copy=False)


def normalize_dataset(specs: np.ndarray, eps: float = 1e-6) -> np.ndarray:
    mean = float(specs.mean())
    std = float(specs.std())
    std = max(std, eps)
    return ((specs - mean) / std).astype(np.float32, copy=False)


# ---------------------------------------------------------------------------
# Utility helpers
# ---------------------------------------------------------------------------

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--data-root",
        default=DEFAULT_DATA_ROOT,
        help="Root directory containing modulation folders (default: %(default)s)",
    )
    parser.add_argument(
        "--modulation",
        default="all",
        help="Modulation folder to load (default: %(default)s)",
    )
    parser.add_argument(
        "--snrs",
        default="SNR-5dB,SNR0dB,SNR5dB,SNR10dB,SNR15dB,SNR20dB,SNR25dB",
        help=(
            "Comma-separated list of SNR folder names to include. Pass 'all' "
            "to include every SNR discovered under the modulation (default: %(default)s)"
        ),
    )
    parser.add_argument(
        "--mobility",
        nargs="+",
        default=["all"],
        help=(
            "Mobility folder(s) to filter on. Pass 'all' to include every mobility "
            "(default: %(default)s). Multiple values can be provided either as a "
            "space-separated list (e.g. '--mobility vehicular pedestrian') or a "
            "comma-separated string."
        ),
    )
    parser.add_argument(
        "--fft-folder",
        default="all",
        help=(
            "FFT size folder name to use. Pass 'all' to include every FFT variant "
            "(default: %(default)s)"
        ),
    )
    parser.add_argument(
        "--samples-per-snr",
        type=int,
        default=500,
        help="Maximum number of samples to draw for each SNR label",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for sampling and t-SNE",
    )
    parser.add_argument(
        "--pooling",
        choices=("mean", "cls"),
        default="mean",
        help="How to collapse token embeddings into a single vector",
    )
    parser.add_argument(
        "--save-path",
        default="task1/snr_separation_plot_latest.png",
        help="Location to save the generated figure (default: %(default)s)",
    )
    parser.add_argument(
        "--checkpoint",
        default=None,
        help="Optional explicit checkpoint path; overrides automatic latest selection",
    )
    parser.add_argument(
        "--models-root",
        default=DEFAULT_MODELS_ROOT,
        help=(
            "Directory containing checkpoints. When --checkpoint is not given, "
            "the latest/best checkpoint inside this directory will be used "
            "(default: %(default)s)"
        ),
    )
    preset_group = parser.add_mutually_exclusive_group()
    preset_group.add_argument(
        "--profile",
        dest="profile",
        choices=tuple(PROFILE_PRESETS.keys()),
        help=(
            "Convenience preset that sets --data-root and --models-root when they "
            "are left at their defaults"
        ),
    )
    preset_group.add_argument(
        "--LTE",
        dest="profile",
        action="store_const",
        const="LTE",
        help="Shortcut for --profile LTE",
    )
    preset_group.add_argument(
        "--WiFi",
        dest="profile",
        action="store_const",
        const="WiFi",
        help="Shortcut for --profile WiFi",
    )
    preset_group.add_argument(
        "--5G",
        dest="profile",
        action="store_const",
        const="5G",
        help="Shortcut for --profile 5G",
    )
    parser.add_argument(
        "--report-metrics",
        action="store_true",
        help="Print clustering metrics (silhouette, 5-fold kNN accuracy)",
    )
    parser.add_argument(
        "--metrics-only",
        action="store_true",
        help="Exit after reporting metrics without running t-SNE or saving figures",
    )
    parser.add_argument(
        "--sampling-mode",
        choices=("first", "reservoir"),
        default="first",
        help="How to down-sample each class (default: first)",
    )
    parser.add_argument(
        "--complex-mode",
        choices=("auto", "magnitude", "interleaved"),
        default="auto",
        help=(
            "How to handle complex spectrograms: 'magnitude' (abs), 'interleaved' (real/imag interleaved along width), "
            "or 'auto' (prefer interleaved when complex). Real-valued inputs are unaffected."
        ),
    )
    parser.add_argument(
        "--label-field",
        choices=("snr", "modulation", "mobility"),
        default="snr",
        help="Choose which label to visualise and score (default: %(default)s)",
    )
    parser.add_argument(
        "--normalization",
        choices=("per-sample", "dataset"),
        default="per-sample",
        help="Normalisation strategy applied before embedding extraction",
    )
    return parser.parse_args()


def find_latest_checkpoint(models_root: Path) -> Path:
    """Return a checkpoint path under ``models_root``.

    Works with either a parent directory that contains multiple run folders,
    or directly with a single run directory containing ``*.pth`` files.
    Chooses the checkpoint with the lowest parsed validation value when
    available, else falls back to most-recent modification time.
    """

    if not models_root.exists():
        raise FileNotFoundError(f"Models root not found: {models_root}")

    if models_root.is_file():
        raise FileNotFoundError(f"Expected a directory, got file: {models_root}")

    # If the provided directory itself contains checkpoints, use it directly.
    checkpoints = list(models_root.glob("*.pth"))
    if not checkpoints:
        # Otherwise, look for subdirectories that contain checkpoints and ignore others (e.g., tensorboard)
        run_dirs = [p for p in models_root.iterdir() if p.is_dir()]
        candidate_runs = [d for d in run_dirs if any(d.glob("*.pth"))]
        if not candidate_runs:
            raise FileNotFoundError(
                f"No checkpoints found under {models_root} (no .pth files in this dir or its run subdirs)"
            )
        latest_run = max(candidate_runs, key=lambda p: p.stat().st_mtime)
        checkpoints = list(latest_run.glob("*.pth"))

    def parse_val_metric(path: Path) -> float | None:
        match = re.search(r"_val([0-9]+(?:\.[0-9]+)?)", path.name)
        if match:
            try:
                return float(match.group(1))
            except ValueError:
                return None
        return None

    parsed = [(parse_val_metric(p), p) for p in checkpoints]
    valid = [item for item in parsed if item[0] is not None]
    if valid:
        valid.sort(key=lambda item: item[0])
        return valid[0][1]

    # Fallback to most recent modification time
    return max(checkpoints, key=lambda p: p.stat().st_mtime)


def parse_snr_list(snr_argument: str | None) -> set[str] | None:
    if snr_argument is None or snr_argument.lower() == "all":
        return None
    values = [item.strip() for item in snr_argument.split(",") if item.strip()]
    return set(values)


def list_snr_samples(
    data_root: Path,
    modulation: str,
    allowed_snrs: set[str] | None,
    mobility_filter: set[str] | None,
    fft_folder: str,
    max_per_class: int,
    rng: random.Random,
    mode: str,
    complex_mode: str,
) -> Dict[str, List[Tuple[np.ndarray, str, str]]]:
    """Collect spectrogram samples grouped by SNR label.

    Supports both legacy PKL layout with a trailing 'spectrograms/' folder and
    MATLAB .mat bundles saved directly under the mobility folder.

    Returns: mapping from SNR label to list of tuples: (spec, modulation, mobility)
    """

    class_samples: Dict[str, List[Tuple[np.ndarray, str, str]]] = defaultdict(list)
    seen_counts: Dict[str, int] = defaultdict(int)

    # Search patterns:
    #  - PKL under .../spectrograms/*.pkl
    #  - MAT under .../spectrogram_*.mat
    patterns = [
        str(data_root / "**" / "spectrograms" / "*.pkl"),
        str(data_root / "**" / "spectrogram_*.mat"),
    ]

    mobility_set = {"static", "pedestrian", "vehicular"}

    def extract_tokens(rel_parts: Tuple[str, ...]) -> Tuple[str, str, str, str] | None:
        # Heuristic extraction to support both layouts
        # modulation: first path segment below data_root
        if not rel_parts:
            return None
        modulation_folder = rel_parts[0]

        # snr: first segment like SNR(-?)NdB
        snr_folder = next((p for p in rel_parts if re.match(r"^SNR-?\d+dB$", p)), None)
        if snr_folder is None:
            return None

        # mobility: one of known labels
        mobility_folder = next((p for p in rel_parts if p.lower() in mobility_set), None)
        if mobility_folder is None:
            return None

        # fft/window folder if present (PKL layout), else fallback for MAT
        fft_folder_name = next((p for p in rel_parts if p.startswith("win") or p.startswith("fft")), "fft_unknown")

        return modulation_folder, snr_folder, mobility_folder, fft_folder_name

    for pattern in patterns:
        for path_str in glob.iglob(pattern, recursive=True):
            path = Path(path_str)
            try:
                rel_parts = path.relative_to(data_root).parts
            except ValueError:
                continue

            tokens = extract_tokens(rel_parts)
            if tokens is None:
                continue
            modulation_folder, snr_folder, mobility_folder, fft_folder_name = tokens

            # Apply filters
            if modulation.lower() != "all" and modulation_folder != modulation:
                continue
            if allowed_snrs is not None and snr_folder not in allowed_snrs:
                continue
            if mobility_filter is not None and mobility_folder.lower() not in mobility_filter:
                continue
            if fft_folder != "all" and fft_folder_name != fft_folder:
                continue

            class_label = snr_folder
            if mode == "first" and len(class_samples[class_label]) >= max_per_class:
                continue

            # Load spectrogram data (supports .pkl and .mat)
            try:
                arr = load_spectrogram_data(str(path))
            except Exception as exc:  # pragma: no cover - I/O heavy
                print(f"[WARN] Failed to load {path}: {exc}")
                continue

            if not isinstance(arr, np.ndarray) or arr.size == 0:
                continue

            # If loaded spectrograms are complex, convert according to mode
            if np.iscomplexobj(arr):
                if complex_mode == "magnitude":
                    arr = np.abs(arr)
                else:
                    # Interleave real/imag parts along the width dimension
                    if arr.ndim == 4 and arr.shape[1] == 1:
                        arr = arr[:, 0]
                    if arr.ndim == 3:
                        real = arr.real.astype(np.float32, copy=False)
                        imag = arr.imag.astype(np.float32, copy=False)
                        n, h, w = real.shape
                        inter = np.empty((n, h, w * 2), dtype=np.float32)
                        inter[:, :, 0::2] = real
                        inter[:, :, 1::2] = imag
                        arr = inter
                    else:
                        # Fallback to magnitude for unsupported shapes
                        arr = np.abs(arr)

            # Normalize shapes:
            #  - (N, H, W)
            #  - (N, C, H, W) -> collapse channels via mean
            if arr.ndim == 4:
                # (N, C, H, W) -> (N, H, W)
                if arr.shape[1] > 1:
                    specs = arr.mean(axis=1)
                else:
                    specs = arr[:, 0]
            elif arr.ndim == 3:
                specs = arr
            elif arr.ndim == 2:
                specs = arr[None, ...]
            else:
                print(f"[WARN] Unexpected spectrogram shape in {path}: {arr.shape}")
                continue

            for spec in specs:
                sample = np.asarray(spec, dtype=np.float32)
                bucket = class_samples[class_label]

                if len(bucket) < max_per_class:
                    bucket.append((sample, modulation_folder, mobility_folder))
                    seen_counts[class_label] += 1
                elif mode == "reservoir":
                    seen_counts[class_label] += 1
                    j = rng.randint(0, seen_counts[class_label] - 1)
                    if j < max_per_class:
                        bucket[j] = (sample, modulation_folder, mobility_folder)
                else:  # mode == "first" and already full
                    break

    return class_samples


def sample_balanced_dataset(
    class_samples: Dict[str, List[Tuple[np.ndarray, str, str]]],
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, List[str]]:
    """Stack the sampled spectrograms alongside SNR, modulation, and mobility labels."""

    features: List[np.ndarray] = []
    snr_labels: List[str] = []
    modulation_labels: List[str] = []
    mobility_labels: List[str] = []
    class_names = sorted(class_samples.keys())

    for class_name in class_names:
        samples = class_samples[class_name]
        if not samples:
            continue
        for sample, modulation_label, mobility_label in samples:
            features.append(sample)
            snr_labels.append(class_name)
            modulation_labels.append(modulation_label)
            mobility_labels.append(mobility_label)

    if not features:
        raise RuntimeError("No spectrogram samples collected for the specified filters")

    stacked = np.stack(features)  # [N, 128, 128]
    return (
        stacked,
        np.array(snr_labels),
        np.array(modulation_labels),
        np.array(mobility_labels),
        class_names,
    )


def unfold_patches_square(x: torch.Tensor, patch_size: int = 4) -> torch.Tensor:
    # Input shape: [B, H, W]; extracts (patch_size x patch_size) patches
    patches_h = x.unfold(1, patch_size, patch_size)
    patches = patches_h.unfold(2, patch_size, patch_size)
    return patches.contiguous().view(x.shape[0], -1, patch_size * patch_size)


def unfold_patches_rect(x: torch.Tensor, patch_rows: int = 4, patch_cols: int = 8) -> torch.Tensor:
    # Input shape: [B, H, W]; extracts (patch_rows x patch_cols) patches (for interleaved complex)
    patches_h = x.unfold(1, patch_rows, patch_rows)
    patches = patches_h.unfold(2, patch_cols, patch_cols)
    return patches.contiguous().view(x.shape[0], -1, patch_rows * patch_cols)


def extract_tokens(spec: np.ndarray, device: torch.device, interleaved: bool) -> torch.Tensor:
    tensor = torch.from_numpy(spec).unsqueeze(0).to(device)
    if interleaved:
        # Rectangular patches 4x8 to cover 4x4 complex bins (real+imag)
        return unfold_patches_rect(tensor, 4, 8)  # [1, 1024, 32]
    else:
        return unfold_patches_square(tensor, 4)   # [1, 1024, 16]


def pool_embeddings(
    tokens: torch.Tensor,
    model: torch.nn.Module,
    pooling: str,
) -> np.ndarray:
    # Append CLS token (value 0.2) before passing through the transformer.
    cls_token = torch.full((tokens.size(0), 1, tokens.size(-1)), 0.2, device=tokens.device)
    inputs = torch.cat([cls_token, tokens], dim=1)  # [B, 1025, 16]

    with torch.no_grad():
        outputs = model(inputs)  # [B, 1025, 128]

    if pooling == "cls":
        pooled = outputs[:, 0]
    else:  # mean pooling across patch tokens (exclude CLS)
        pooled = outputs[:, 1:].mean(dim=1)

    return pooled.detach().cpu().numpy()


def sort_snr_labels(labels: List[str]) -> List[str]:
    """Sort SNR labels by numeric value instead of lexicographic order."""
    def extract_snr_value(label: str) -> float:
        """Extract numeric SNR value from label like 'SNR-5dB' -> -5.0"""
        import re
        match = re.search(r'SNR(-?\d+)dB', label)
        if match:
            return float(match.group(1))
        else:
            return float('inf')  # Put non-SNR labels at the end
    
    return sorted(labels, key=extract_snr_value)


def run_tsne(x: np.ndarray, labels: np.ndarray, title: str, ax: plt.Axes) -> None:
    scaler = StandardScaler()
    x_scaled = scaler.fit_transform(x)
    # Guard against NaN/Inf or extreme values that can break SVD/TSNE
    x_scaled = np.nan_to_num(x_scaled, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
    x_scaled = np.clip(x_scaled, -1e6, 1e6)
    x_scaled = x_scaled.astype(np.float32, copy=False)
    # Use a safe perplexity relative to sample count (sklearn requirement: < n_samples).
    max_perplexity = max(5, min(30, len(x_scaled) // 10))
    perplexity = min(max_perplexity, len(x_scaled) - 1)
    perplexity = max(perplexity, 5)

    tsne = TSNE(
        n_components=2,
        perplexity=perplexity,
        random_state=42,
        init="random",
        learning_rate="auto",
    )
    try:
        embedding = tsne.fit_transform(x_scaled)
    except Exception as e:
        # Fallback to PCA if TSNE/SVD fails
        print(f"[WARN] t-SNE failed ({e}); falling back to PCA.")
        pca = PCA(n_components=2, svd_solver="full", random_state=42)
        embedding = pca.fit_transform(x_scaled)

    class_names = sort_snr_labels(list(np.unique(labels)))
    colors = plt.cm.Set3(np.linspace(0, 1, len(class_names)))
    for color, class_name in zip(colors, class_names):
        mask = labels == class_name
        ax.scatter(embedding[mask, 0], embedding[mask, 1], c=[color], s=18, alpha=0.7, label=class_name)

    # ax.set_title(title, fontsize=14, fontweight="bold")  # Title removed for paper
    ax.set_xlabel("t-SNE Component 1", fontsize=16)
    ax.set_ylabel("t-SNE Component 2", fontsize=16)
    ax.tick_params(labelsize=14)  # Increase tick label size
    ax.grid(True, alpha=0.3)
    ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", fontsize=12)


def compute_metrics(name: str, features: np.ndarray, labels: np.ndarray) -> None:
    if len(np.unique(labels)) < 2:
        print(f"[METRIC] {name}: skipped (only one class present)")
        return

    scaler = StandardScaler()
    features_scaled = scaler.fit_transform(features)

    silhouette = silhouette_score(features_scaled, labels)

    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    scores: List[float] = []
    for train_idx, test_idx in skf.split(features_scaled, labels):
        clf = KNeighborsClassifier(n_neighbors=5)
        clf.fit(features_scaled[train_idx], labels[train_idx])
        scores.append(clf.score(features_scaled[test_idx], labels[test_idx]))

    mean_acc = float(np.mean(scores))
    std_acc = float(np.std(scores))
    print(
        f"[METRIC] {name}: silhouette={silhouette:.3f}, "
        f"5-NN accuracy={mean_acc:.3f} ± {std_acc:.3f}"
    )


# ---------------------------------------------------------------------------
# Main execution
# ---------------------------------------------------------------------------


def main() -> None:
    args = parse_args()

    if args.profile:
        preset = PROFILE_PRESETS.get(args.profile)
        if not preset:
            raise ValueError(f"Unknown profile requested: {args.profile}")
        if args.data_root == DEFAULT_DATA_ROOT:
            args.data_root = preset["data_root"]
        if args.models_root == DEFAULT_MODELS_ROOT:
            args.models_root = preset["models_root"]

    if args.profile:
        print(f"[INFO] Profile preset active: {args.profile}")

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    data_root = Path(args.data_root)
    if not data_root.exists():
        raise FileNotFoundError(f"Data root not found: {data_root}")

    allowed_snrs = parse_snr_list(args.snrs)

    mobility_filter: set[str] | None = None
    if args.mobility:
        mobility_values: List[str] = []
        for value in args.mobility:
            mobility_values.extend([item.strip() for item in value.split(",") if item.strip()])
        mobility_values = [value for value in mobility_values if value]
        if mobility_values and not (len(mobility_values) == 1 and mobility_values[0].lower() == "all"):
            mobility_filter = {value.lower() for value in mobility_values}
            print(
                "[INFO] Mobility filter active: "
                + ", ".join(sorted(mobility_filter))
            )

    class_samples = list_snr_samples(
        data_root,
        args.modulation,
        allowed_snrs,
        mobility_filter,
        args.fft_folder,
        args.samples_per_snr,
        random,
        args.sampling_mode,
        args.complex_mode,
    )
    samples, snr_labels, modulation_labels, mobility_labels, _ = sample_balanced_dataset(class_samples)

    if args.label_field == "snr":
        labels = snr_labels
        label_name = "SNR"
        label_display = "SNR"
    elif args.label_field == "modulation":
        labels = modulation_labels
        label_name = "modulation"
        label_display = "Modulation"
    else:  # mobility
        labels = mobility_labels
        label_name = "mobility"
        label_display = "Mobility"

    unique_labels = np.unique(labels)
    print(
        f"[INFO] Loaded {samples.shape[0]} spectrograms across {len(unique_labels)} {label_name} buckets"
    )
    class_counts = Counter(labels)
    print(f"[INFO] Samples per {label_name}:")
    for name, count in sorted(class_counts.items()):
        print(f"  {name}: {count}")

    if args.label_field != "snr":
        snr_counts = Counter(snr_labels)
        print("[INFO] SNR distribution (sampling classes):")
        for name, count in sorted(snr_counts.items()):
            print(f"  {name}: {count}")
    if args.label_field == "mobility":
        modulation_counts = Counter(modulation_labels)
        print("[INFO] Modulation distribution:")
        for name, count in sorted(modulation_counts.items()):
            print(f"  {name}: {count}")

    normalization_mode = args.normalization
    if normalization_mode == "per-sample":
        normalized_samples = normalize_per_sample(samples)
    else:
        normalized_samples = normalize_dataset(samples)
    print(f"[INFO] Normalisation mode: {normalization_mode}")

    # Flatten spectrograms (after optional normalization) for the raw t-SNE view.
    raw_vectors = normalized_samples.reshape(normalized_samples.shape[0], -1)

    # Prepare LWM model and embeddings for the right subplot.
    if args.checkpoint:
        checkpoint_path = Path(args.checkpoint)
        if not checkpoint_path.exists():
            raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
    else:
        checkpoint_path = find_latest_checkpoint(Path(args.models_root))
    print(f"[INFO] Using checkpoint: {checkpoint_path}")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[INFO] Using device: {device}")
    print(f"[INFO] Pooling strategy: {args.pooling}")
    # Determine complex handling strategy for model/patching
    use_interleaved = False
    if args.complex_mode == "interleaved":
        use_interleaved = True
    elif args.complex_mode == "auto":
        # Heuristic: if any sample contains width > 128, assume interleaved (e.g., 128x256)
        sample_shape = tuple(normalized_samples.shape[1:])
        if len(sample_shape) == 2 and sample_shape[1] > 128:
            use_interleaved = True

    element_length = 32 if use_interleaved else 16

    model = lwm_model(element_length=element_length, d_model=128, n_layers=12, max_len=1025, n_heads=8, dropout=0.1)
    state_dict = torch.load(checkpoint_path, map_location=device)
    if any(k.startswith("module.") for k in state_dict):
        state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
    try:
        model.load_state_dict(state_dict, strict=False)
    except RuntimeError as e:
        msg = str(e)
        # Fallback: checkpoint expects element_length=16 (magnitude), but we constructed 32 (interleaved)
        mismatch16 = "[128, 16]" in msg or "[16]" in msg
        mismatch32 = "[128, 32]" in msg or "[32]" in msg
        if mismatch16 and not mismatch32:
            print("[WARN] Checkpoint expects token dimension 16. Falling back to magnitude embedding.")
            use_interleaved = False
            element_length = 16
            # Recreate model and reload
            model = lwm_model(element_length=element_length, d_model=128, n_layers=12, max_len=1025, n_heads=8, dropout=0.1)
            model.load_state_dict(state_dict, strict=False)
        else:
            raise
    model = model.to(device).eval()

    def collapse_interleaved_to_magnitude(spec: np.ndarray) -> np.ndarray:
        # spec: [H, 2W] with interleaved real/imag along width -> [H, W] magnitude
        h, w2 = spec.shape
        if w2 % 2 != 0:
            return spec  # cannot collapse; return as-is
        real = spec[:, 0::2]
        imag = spec[:, 1::2]
        return np.sqrt(np.maximum(real * real + imag * imag, 0.0, dtype=np.float32))

    # If we fell back to magnitude (use_interleaved False) but inputs are interleaved, collapse for embeddings only
    embed_inputs = normalized_samples
    if not use_interleaved and normalized_samples.shape[2] > 128:
        collapsed = []
        for spec in normalized_samples:
            collapsed.append(collapse_interleaved_to_magnitude(spec))
        embed_inputs = np.stack(collapsed).astype(np.float32, copy=False)

    embeddings: List[np.ndarray] = []
    for spec in embed_inputs:
        tokens = extract_tokens(spec, device, interleaved=use_interleaved)
        embedding = pool_embeddings(tokens, model, args.pooling)
        embeddings.append(embedding.squeeze(0))

    embeddings_np = np.vstack(embeddings)
    print(f"[INFO] Generated embeddings with shape {embeddings_np.shape}")

    if args.report_metrics:
        compute_metrics("Raw spectrogram", raw_vectors, labels)
        pool_label = "LWM mean" if args.pooling == "mean" else "LWM CLS"
        compute_metrics(pool_label, embeddings_np, labels)
        if args.metrics_only:
            return

    # Plot results (two subplots matching the original figure format).
    fig, axes = plt.subplots(1, 2, figsize=(18, 7))
    raw_title = f"Raw Spectrogram t-SNE (by {label_display})"
    pooling_label = "Mean Pool" if args.pooling == "mean" else "CLS Token"
    embedding_title = f"LWM Embedding t-SNE ({pooling_label}, by {label_display})"
    run_tsne(raw_vectors, labels, raw_title, axes[0])
    run_tsne(embeddings_np, labels, embedding_title, axes[1])

    fig.tight_layout()
    save_path = Path(args.save_path)

    communication_tag: str | None = None
    if args.profile:
        communication_tag = args.profile
    else:
        root_name = Path(args.data_root).name
        if root_name:
            communication_tag = root_name

    def ensure_suffix(stem: str, suffix: str) -> str:
        return stem if stem.endswith(suffix) else f"{stem}_{suffix}"

    updated_stem = save_path.stem
    if communication_tag:
        updated_stem = ensure_suffix(updated_stem, communication_tag)
    if args.label_field != "snr":
        label_suffix = f"by_{args.label_field}"
        updated_stem = ensure_suffix(updated_stem, label_suffix)

    if updated_stem != save_path.stem:
        save_path = save_path.with_name(f"{updated_stem}{save_path.suffix}")
    save_path.parent.mkdir(parents=True, exist_ok=True)
    plt.savefig(save_path, dpi=600, bbox_inches="tight")
    print(f"[INFO] Figure saved to {save_path}")
    
    # Also save PDF version for paper (vector format, no resolution limit)
    pdf_path = save_path.with_suffix('.pdf')
    plt.savefig(pdf_path, format='pdf', bbox_inches="tight")
    print(f"[INFO] PDF version saved to {pdf_path}")


if __name__ == "__main__":
    main()