src/temporal_embedding.py

"""Temporal Embedding: Visualize how brain activity evolves within a recording.
One file per intent type, compute features per window, PCA to 2D, draw trajectory
with arrows colored by phase (rest->initiation->sustained).
Saves results/intent_evolution.png."""
import numpy as np
from pathlib import Path
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch

from preprocess import bandpass_filter, extract_active_segment, normalize_channels, segment_windows
from features import extract_psd_features, extract_stat_features, extract_cross_channel_features

PROJECT_ROOT = Path(__file__).resolve().parent.parent
RESULTS_DIR = PROJECT_ROOT / "results"

# One representative file per intent type
INTENT_FILES = {
    "Left Fist":       "2562e7bd-14.npz",
    "Right Fist":      "0b2dbd41-34.npz",
    "Both Fists":      "4787dfb9-10.npz",
    "Tongue Tapping":  "0b2dbd41-16.npz",
    "Relax":           "2161ecb6-12.npz",
}

INTENT_COLORS = {
    "Left Fist":       "#4285f4",
    "Right Fist":      "#ea4335",
    "Both Fists":      "#34a853",
    "Tongue Tapping":  "#fbbc04",
    "Relax":           "#9aa0a6",
}

INTENT_ACTIONS = {
    "Left Fist":       "LEFT",
    "Right Fist":      "RIGHT",
    "Both Fists":      "FORWARD",
    "Tongue Tapping":  "BACKWARD",
    "Relax":           "STOP",
}

PHASE_MARKERS = {
    "rest":        "o",
    "initiation":  "^",
    "sustained":   "s",
}


def extract_windows_and_features(npz_path):
    """Extract per-window features from a single file, return features + window indices."""
    arr = np.load(str(npz_path), allow_pickle=True)
    eeg = arr["feature_eeg"]
    info = arr["label"].item()

    filtered = bandpass_filter(eeg)
    if np.any(np.isnan(filtered)) or np.any(np.isinf(filtered)):
        return None, None
    active = extract_active_segment(filtered, info["duration"])
    normed = normalize_channels(active)
    windows = segment_windows(normed, 500, 250)

    feats = []
    for w in windows:
        feat = np.concatenate([
            extract_psd_features(w),
            extract_stat_features(w),
            extract_cross_channel_features(w),
        ])
        feats.append(feat)

    return np.array(feats), info


def assign_phases(n_windows, label):
    """Assign phase labels to each window based on position and label."""
    phases = []
    for i in range(n_windows):
        if label == "Relax":
            phases.append("rest")
        elif i < 2:
            phases.append("rest")
        elif i == 2:
            phases.append("initiation")
        else:
            phases.append("sustained")
    return phases


def main():
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    data_dir = PROJECT_ROOT / "data"

    print("=" * 70)
    print("  Temporal Embedding: Intent Evolution Visualization")
    print("=" * 70)

    # Collect all features for PCA fitting
    all_features = []
    intent_data = {}

    for label, fname in INTENT_FILES.items():
        fpath = data_dir / fname
        feats, info = extract_windows_and_features(fpath)
        if feats is None:
            print(f"  Skipping {fname} (filter error)")
            continue

        phases = assign_phases(len(feats), label)
        intent_data[label] = {"features": feats, "phases": phases, "info": info}
        all_features.append(feats)
        print(f"  {label:17s}: {len(feats)} windows from {fname}")

    # Fit PCA on all windows jointly
    X_all = np.vstack(all_features)
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X_all)
    pca = PCA(n_components=2)
    X_2d = pca.fit_transform(X_scaled)

    print(f"\n  PCA explained variance: {pca.explained_variance_ratio_[0]:.1%}, {pca.explained_variance_ratio_[1]:.1%}")

    # Split back into per-intent
    idx = 0
    for label in intent_data:
        n = len(intent_data[label]["features"])
        intent_data[label]["pca"] = X_2d[idx:idx + n]
        idx += n

    # Plot
    fig, ax = plt.subplots(figsize=(12, 9))
    fig.patch.set_facecolor("#1a1a2e")
    ax.set_facecolor("#16213e")

    phase_colors = {
        "rest":        "#666666",
        "initiation":  "#ffffff",
        "sustained":   None,  # use intent color
    }

    for label, data in intent_data.items():
        pts = data["pca"]
        phases = data["phases"]
        color = INTENT_COLORS[label]
        action = INTENT_ACTIONS[label]

        # Draw trajectory line
        ax.plot(pts[:, 0], pts[:, 1], "-", color=color, alpha=0.3, linewidth=1.5)

        # Draw arrows along trajectory
        for i in range(len(pts) - 1):
            dx = pts[i+1, 0] - pts[i, 0]
            dy = pts[i+1, 1] - pts[i, 1]
            ax.annotate("", xy=(pts[i+1, 0], pts[i+1, 1]),
                        xytext=(pts[i, 0], pts[i, 1]),
                        arrowprops=dict(arrowstyle="->", color=color, alpha=0.5, lw=1.2))

        # Draw points colored by phase
        for i, (pt, phase) in enumerate(zip(pts, phases)):
            marker = PHASE_MARKERS[phase]
            if phase == "sustained":
                fc = color
            elif phase == "initiation":
                fc = "white"
            else:
                fc = "#666666"
            ec = color
            size = 80 if phase == "initiation" else 40
            ax.scatter(pt[0], pt[1], marker=marker, s=size, c=fc,
                      edgecolors=ec, linewidths=1.5, zorder=5)

        # Label start and end
        ax.annotate(f"{action}\nstart", xy=(pts[0, 0], pts[0, 1]),
                   fontsize=8, color=color, fontweight="bold",
                   ha="center", va="bottom",
                   xytext=(0, 8), textcoords="offset points")
        ax.annotate(f"end", xy=(pts[-1, 0], pts[-1, 1]),
                   fontsize=8, color=color,
                   ha="center", va="top",
                   xytext=(0, -8), textcoords="offset points")

    # Legend for intents
    for label, color in INTENT_COLORS.items():
        action = INTENT_ACTIONS[label]
        ax.plot([], [], "o-", color=color, markersize=6, label=f"{label} -> {action}")

    # Legend for phases
    ax.scatter([], [], marker="o", s=40, c="#666666", edgecolors="white",
              label="Rest phase")
    ax.scatter([], [], marker="^", s=80, c="white", edgecolors="white",
              label="Initiation")
    ax.scatter([], [], marker="s", s=40, c="#aaaaaa", edgecolors="white",
              label="Sustained")

    legend = ax.legend(loc="upper left", fontsize=10, facecolor="#16213e",
                      edgecolor="#444", labelcolor="white", ncol=2)

    ax.set_xlabel(f"PC1 ({pca.explained_variance_ratio_[0]:.1%} variance)",
                 fontsize=12, color="white")
    ax.set_ylabel(f"PC2 ({pca.explained_variance_ratio_[1]:.1%} variance)",
                 fontsize=12, color="white")
    ax.set_title("Intent Evolution: How Brain Activity Changes Within a Recording",
                fontsize=14, fontweight="bold", color="white", pad=15)

    ax.tick_params(colors="white")
    for spine in ax.spines.values():
        spine.set_color("#444")

    # Subtitle
    fig.text(0.5, 0.01,
             "Each trajectory shows one EEG recording projected to feature space. "
             "Arrows show temporal progression. Phases: rest (circle) -> initiation (triangle) -> sustained (square).",
             ha="center", fontsize=9, color="#aaaaaa")

    plt.tight_layout(rect=[0, 0.03, 1, 1])
    out_path = str(RESULTS_DIR / "intent_evolution.png")
    plt.savefig(out_path, dpi=150, facecolor=fig.get_facecolor(), bbox_inches="tight")
    plt.close()
    print(f"\n  Saved to {out_path}")


if __name__ == "__main__":
    main()