Upload 3 files

.gitattributes

@@ -1 +1 @@
-*.pth filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -textpic.png filter=lfs diff=lfs merge=lfs -text

pkas_cal_trainer_gemini.py

@@ -0,0 +1,540 @@
+"""
+Calcium-Bridged Temporal EEG Decoder (V2 - Extended Time Window)
+Integrates Phase-Calcium-Latent constraint satisfaction dynamics with EEG temporal windows.
+
+Core Concept: Each EEG time window is processed by a constraint solver whose 
+calcium/W state carries over to initialize the next window, modeling how the brain 
+sequentially satisfies perceptual constraints.
+
+V2 Update: The time window has been extended to 550ms based on ERP analysis from the 
+Alljoined1 paper, adding a 'CognitiveEvaluation' stage to capture late-stage
+semantic and working memory signals.
+"""
+
+import os
+import json
+import tkinter as tk
+from tkinter import ttk, filedialog, messagebox
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+import torch.nn.functional as F
+import numpy as np
+import threading
+import queue
+from pathlib import Path
+from collections import defaultdict
+import matplotlib.pyplot as plt
+from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
+
+try:
+    from datasets import load_dataset
+    torch.backends.cudnn.benchmark = True
+except ImportError as e:
+    print(f"Missing dependency: {e}")
+    exit()
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+EEG_SAMPLE_RATE = 512
+BATCH_SIZE = 64
+
+# --- CRITICAL CHANGE V2 ---
+# Extended temporal windows to capture later cognitive processing (P300/N400/P600)
+# This aligns the model's "attention span" with the neuroscience data.
+TIME_WINDOWS = [
+    (50, 150, "EarlyVisual"),      # Low-level visual constraints (P100)
+    (150, 250, "MidFeature"),      # Mid-level binding (N170/P200)
+    (250, 350, "LateSemantic"),    # High-level semantics (P300/N400 start)
+    (350, 550, "CognitiveEvaluation") # Deeper context, memory, final check (P300/P600)
+]
+
+TARGET_CATEGORIES = {
+    'elephant': 22, 'giraffe': 25, 'bear': 23, 'zebra': 24,
+    'cow': 21, 'sheep': 20, 'horse': 19, 'dog': 18, 'cat': 17, 'bird': 16,
+    'airplane': 5, 'train': 7, 'boat': 9, 'bus': 6, 'truck': 8,
+    'motorcycle': 4, 'bicycle': 2, 'car': 3,
+    'traffic light': 10, 'fire hydrant': 11, 'stop sign': 13,
+    'parking meter': 14, 'bench': 15,
+}
+
+CATEGORY_NAMES = {v: k for k, v in TARGET_CATEGORIES.items()}
+
+class CalciumAttentionModule(nn.Module):
+    """
+    Phase-Calcium-Latent dynamics for one time window.
+    Models constraint satisfaction via neuromorphic oscillator dynamics.
+    """
+    def __init__(self, n_features, d_model=256):
+        super().__init__()
+        self.n_features = n_features
+        self.d_model = d_model
+        
+        # Phase dynamics (Kuramoto-like)
+        self.phase_proj = nn.Linear(n_features, d_model)
+        
+        # Calcium dynamics (gating/attention)
+        self.ca_gate = nn.Sequential(
+            nn.Linear(d_model, d_model // 2),
+            nn.Sigmoid()
+        )
+        
+        # Latent coupling matrix (W) - learned constraint structure
+        self.W = nn.Parameter(torch.randn(d_model, d_model) * 0.01)
+        
+        # Layer norm for stability
+        self.norm = nn.LayerNorm(d_model)
+        
+    def forward(self, x, prev_ca=None, prev_W=None):
+        """
+        x: Input features [batch, n_features]
+        prev_ca: Previous window's calcium state [batch, d_model]
+        prev_W: Previous window's coupling matrix [d_model, d_model]
+        
+        Returns: features, calcium_state, W_matrix
+        """
+        batch_size = x.size(0)
+        
+        # Phase projection
+        phi = self.phase_proj(x)  # [batch, d_model]
+        
+        # Initialize or carry over calcium
+        if prev_ca is None:
+            ca = torch.zeros(batch_size, self.d_model, device=x.device)
+        else:
+            ca = prev_ca.clone()
+        
+        # Initialize or carry over W (coupling structure)
+        W = self.W if prev_W is None else prev_W
+        
+        # Calcium accumulation (coherence-based)
+        # High when features are aligned (low when conflicting)
+        coherence = torch.abs(torch.cos(phi[:, :, None] - phi[:, None, :]))
+        ca_update = torch.mean(coherence, dim=2)  # [batch, d_model]
+        ca = ca * 0.95 + ca_update * 0.05  # Temporal integration
+        
+        # Calcium-gated attention
+        ca_gate = self.ca_gate(ca)  # [batch, d_model//2]
+        
+        # Apply constraint coupling (W matrix)
+        # This is where "mutual constraint satisfaction" happens
+        coupled = torch.matmul(phi, W)  # [batch, d_model]
+        
+        # Gate the coupling by calcium (only attend where calcium is high)
+        ca_gate_full = torch.cat([ca_gate, ca_gate], dim=1)  # Expand to d_model
+        features = coupled * ca_gate_full
+        
+        # Normalize
+        features = self.norm(features + phi)  # Residual connection
+        
+        return features, ca, W
+
+
+class TemporalConstraintEEGModel(nn.Module):
+    """
+    Sequential constraint satisfaction across EEG time windows.
+    Each window is a constraint solver whose state primes the next.
+    (Dynamically sized based on TIME_WINDOWS constant)
+    """
+    def __init__(self, n_channels=64, num_classes=len(TARGET_CATEGORIES)):
+        super().__init__()
+        self.n_channels = n_channels
+        
+        # CNN feature extractors for each time window
+        self.window_encoders = nn.ModuleList([
+            self._build_cnn_encoder() for _ in TIME_WINDOWS
+        ])
+        
+        # Calcium-attention modules for each window
+        self.ca_modules = nn.ModuleList([
+            CalciumAttentionModule(256, d_model=256) for _ in TIME_WINDOWS
+        ])
+        
+        # --- CRITICAL CHANGE V2 ---
+        # The input layer is now automatically larger (256 * 4) because len(TIME_WINDOWS) is 4.
+        self.classifier = nn.Sequential(
+            nn.Linear(256 * len(TIME_WINDOWS), 512),
+            nn.BatchNorm1d(512),
+            nn.GELU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, num_classes)
+        )
+        
+    def _build_cnn_encoder(self):
+        """Simple CNN for one time window"""
+        return nn.Sequential(
+            nn.Conv1d(self.n_channels, 128, kernel_size=15, padding=7),
+            nn.BatchNorm1d(128),
+            nn.ELU(),
+            nn.MaxPool1d(2),
+            nn.Conv1d(128, 256, kernel_size=7, padding=3),
+            nn.BatchNorm1d(256),
+            nn.ELU(),
+            nn.AdaptiveAvgPool1d(1)
+        )
+    
+    def forward(self, eeg_windows):
+        """
+        eeg_windows: List of tensors [batch, channels, timepoints] for each window
+        
+        Returns: logits, calcium_states (for visualization/analysis)
+        """
+        batch_size = eeg_windows[0].size(0)
+        
+        # Process windows sequentially with calcium/W carryover
+        window_features = []
+        ca_state = None
+        W_state = None
+        ca_history = []
+        
+        for i, (encoder, ca_module, eeg_window) in enumerate(
+            zip(self.window_encoders, self.ca_modules, eeg_windows)
+        ):
+            # Extract CNN features
+            cnn_features = encoder(eeg_window).squeeze(-1)  # [batch, 256]
+            
+            # Apply constraint satisfaction dynamics
+            features, ca_state, W_state = ca_module(
+                cnn_features, 
+                prev_ca=ca_state, 
+                prev_W=W_state
+            )
+            
+            window_features.append(features)
+            ca_history.append(ca_state.detach().cpu().numpy())
+        
+        # Concatenate all window features
+        combined = torch.cat(window_features, dim=1)
+        
+        # Final classification
+        logits = self.classifier(combined)
+        
+        return logits, ca_history
+
+
+class CalciumEEGDataset(Dataset):
+    """Dataset that provides EEG data split by time windows"""
+    def __init__(self, coco_path, annotations_path, split='train', 
+                 max_samples=None, trials_to_average=1):
+        self.coco_path = Path(coco_path)
+        
+        # Load dataset
+        print(f"Loading Alljoined ({split})...")
+        self.dataset = load_dataset("Alljoined/05_125", split=split, streaming=False)
+        
+        if max_samples:
+            self.dataset = self.dataset.select(range(min(int(max_samples), len(self.dataset))))
+        
+        # Load COCO annotations
+        print(f"Loading COCO annotations...")
+        with open(annotations_path, 'r') as f:
+            coco_data = json.load(f)
+        
+        self.image_categories = defaultdict(set)
+        for ann in coco_data['annotations']:
+            img_id = ann['image_id']
+            if ann['category_id'] in CATEGORY_NAMES:
+                self.image_categories[img_id].add(ann['category_id'])
+        
+        # Pre-cache samples
+        print("Pre-caching EEG data...")
+        self.samples = []
+        for idx, sample in enumerate(self.dataset):
+            coco_id = sample['coco_id']
+            if coco_id in self.image_categories and len(self.image_categories[coco_id]) > 0:
+                label = torch.zeros(len(TARGET_CATEGORIES))
+                for cat_id in self.image_categories[coco_id]:
+                    if cat_id in CATEGORY_NAMES:
+                        cat_idx = list(TARGET_CATEGORIES.values()).index(cat_id)
+                        label[cat_idx] = 1.0
+                
+                if label.sum() > 0:
+                    self.samples.append((idx, label))
+        
+        print(f"Cached {len(self.samples)} samples")
+    
+    def __len__(self):
+        return len(self.samples)
+    
+    def __getitem__(self, idx):
+        sample_idx, label = self.samples[idx]
+        sample = self.dataset[sample_idx]
+        
+        eeg_data = np.array(sample['EEG'], dtype=np.float32)
+        
+        # Extract time windows
+        eeg_windows = []
+        for start_ms, end_ms, _ in TIME_WINDOWS:
+            start_idx = int((start_ms / 1000.0) * EEG_SAMPLE_RATE)
+            end_idx = int((end_ms / 1000.0) * EEG_SAMPLE_RATE)
+            
+            if eeg_data.shape[1] >= end_idx:
+                window = eeg_data[:, start_idx:end_idx]
+            else:
+                window = eeg_data[:, start_idx:]
+                # Pad if needed
+                if window.shape[1] < (end_idx - start_idx):
+                    pad_width = (end_idx - start_idx) - window.shape[1]
+                    window = np.pad(window, ((0,0), (0, pad_width)), mode='edge')
+            
+            # Normalize
+            window = (window - window.mean(axis=1, keepdims=True)) / \
+                     (window.std(axis=1, keepdims=True) + 1e-8)
+            
+            eeg_windows.append(torch.from_numpy(window).float())
+        
+        return eeg_windows, label
+
+
+class CalciumEEGTrainerGUI(tk.Tk):
+    def __init__(self):
+        super().__init__()
+        self.title("Calcium-Bridged Temporal EEG Decoder V2")
+        self.geometry("1200x850")
+        
+        self.coco_path = ""
+        self.annotations_path = ""
+        self.train_thread = None
+        self.stop_flag = threading.Event()
+        self.log_queue = queue.Queue()
+        
+        self.setup_gui()
+        self.process_logs()
+    
+    def setup_gui(self):
+        # Title
+        title = tk.Label(self, text="Calcium-Bridged Temporal EEG Decoder (V2 - Extended Window)", 
+                        font=("Arial", 14, "bold"))
+        title.pack(pady=10)
+        
+        info = tk.Label(self, 
+                       text="Sequential constraint satisfaction across 4 ERP time windows up to 550ms\n"
+                            "Calcium/W state from early windows primes later windows",
+                       fg="blue", font=("Arial", 9))
+        info.pack(pady=5)
+        
+        # Paths
+        path_frame = ttk.LabelFrame(self, text="Dataset")
+        path_frame.pack(pady=5, padx=10, fill=tk.X)
+        
+        tk.Label(path_frame, text="COCO:").grid(row=0, column=0, sticky=tk.W, padx=5, pady=3)
+        self.coco_var = tk.StringVar()
+        ttk.Entry(path_frame, textvariable=self.coco_var, width=50).grid(row=0, column=1, padx=5)
+        ttk.Button(path_frame, text="Browse", command=self.browse_coco).grid(row=0, column=2)
+        
+        tk.Label(path_frame, text="Annotations:").grid(row=1, column=0, sticky=tk.W, padx=5, pady=3)
+        self.ann_var = tk.StringVar()
+        ttk.Entry(path_frame, textvariable=self.ann_var, width=50).grid(row=1, column=1, padx=5)
+        ttk.Button(path_frame, text="Browse", command=self.browse_ann).grid(row=1, column=2)
+        
+        # Settings
+        settings_frame = ttk.LabelFrame(self, text="Training Settings")
+        settings_frame.pack(pady=5, padx=10, fill=tk.X)
+        
+        tk.Label(settings_frame, text="Max Samples:").grid(row=0, column=0, padx=5)
+        self.max_var = tk.IntVar(value=3000)
+        tk.Spinbox(settings_frame, from_=1000, to=10000, increment=1000,
+                  textvariable=self.max_var, width=10).grid(row=0, column=1)
+        
+        tk.Label(settings_frame, text="Epochs:").grid(row=0, column=2, padx=5)
+        self.epochs_var = tk.IntVar(value=100)
+        tk.Spinbox(settings_frame, from_=50, to=500, increment=50,
+                  textvariable=self.epochs_var, width=10).grid(row=0, column=3)
+        
+        # --- CRITICAL CHANGE V2 ---
+        # Updated GUI to reflect the new 4-stage process.
+        windows_frame = ttk.LabelFrame(self, text="Constraint Satisfaction Stages")
+        windows_frame.pack(pady=5, padx=10, fill=tk.X)
+        
+        for start, end, label in TIME_WINDOWS:
+            desc = {
+                "EarlyVisual": "Low-level visual features (edges, textures)",
+                "MidFeature": "Mid-level binding (parts, shapes)",
+                "LateSemantic": "High-level semantics (concepts, context)",
+                "CognitiveEvaluation": "Memory, context check, final decision"
+            }
+            tk.Label(windows_frame, 
+                    text=f"{label} ({start}-{end}ms): {desc[label]}", 
+                    font=("Courier", 9)).pack(anchor=tk.W, padx=10, pady=2)
+        
+        # Buttons
+        btn_frame = tk.Frame(self)
+        btn_frame.pack(pady=10)
+        
+        self.train_btn = tk.Button(btn_frame, text="Train Extended Model (V2)", 
+                                   command=self.start_train,
+                                   bg="#4CAF50", fg="white", font=("Arial", 10, "bold"))
+        self.train_btn.pack(side=tk.LEFT, padx=5)
+        
+        self.stop_btn = tk.Button(btn_frame, text="Stop", 
+                                  command=self.stop_train,
+                                  bg="#f44336", fg="white",
+                                  state=tk.DISABLED)
+        self.stop_btn.pack(side=tk.LEFT, padx=5)
+        
+        # Progress
+        self.progress = ttk.Progressbar(self, mode='determinate')
+        self.progress.pack(fill=tk.X, padx=10, pady=5)
+        
+        # Log
+        log_frame = ttk.LabelFrame(self, text="Training Log")
+        log_frame.pack(pady=5, padx=10, fill=tk.BOTH, expand=True)
+        
+        self.log_text = tk.Text(log_frame, height=20, bg='black', fg='lightgreen',
+                               font=('Courier', 8))
+        self.log_text.pack(fill=tk.BOTH, expand=True)
+    
+    def browse_coco(self):
+        path = filedialog.askdirectory()
+        if path:
+            self.coco_var.set(path)
+            self.coco_path = path
+    
+    def browse_ann(self):
+        path = filedialog.askopenfilename(filetypes=[("JSON", "*.json")])
+        if path:
+            self.ann_var.set(path)
+            self.annotations_path = path
+    
+    def log(self, msg):
+        self.log_queue.put(msg)
+    
+    def process_logs(self):
+        try:
+            while not self.log_queue.empty():
+                msg = self.log_queue.get_nowait()
+                self.log_text.insert(tk.END, msg + "\n")
+                self.log_text.see(tk.END)
+        except queue.Empty:
+            pass
+        self.after(100, self.process_logs)
+    
+    def start_train(self):
+        if not self.coco_path or not self.annotations_path:
+            messagebox.showerror("Error", "Select paths first")
+            return
+        
+        self.stop_flag.clear()
+        self.train_btn.config(state=tk.DISABLED)
+        self.stop_btn.config(state=tk.NORMAL)
+        
+        self.train_thread = threading.Thread(target=self._train_model, daemon=True)
+        self.train_thread.start()
+    
+    def stop_train(self):
+        self.stop_flag.set()
+    
+    def _train_model(self):
+        try:
+            self.log("="*70)
+            self.log("CALCIUM-BRIDGED TEMPORAL EEG DECODER (V2 - Extended Window)")
+            self.log("="*70)
+            self.log("\nConcept: Sequential constraint satisfaction across FOUR time windows")
+            self.log("Now capturing late-stage cognitive evaluation signals up to 550ms\n")
+            
+            # Create dataset
+            dataset = CalciumEEGDataset(
+                self.coco_path,
+                self.annotations_path,
+                'train',
+                self.max_var.get()
+            )
+            
+            total = len(dataset)
+            train_size = int(0.8 * total)
+            val_size = total - train_size
+            
+            train_set, val_set = torch.utils.data.random_split(
+                dataset, [train_size, val_size],
+                generator=torch.Generator().manual_seed(42)
+            )
+            
+            self.log(f"Train: {train_size}, Val: {val_size}")
+            
+            train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, 
+                                    shuffle=True, num_workers=0, pin_memory=True)
+            val_loader = DataLoader(val_set, batch_size=BATCH_SIZE, 
+                                   shuffle=False, num_workers=0, pin_memory=True)
+            
+            # Create model (will be automatically sized for 4 windows)
+            model = TemporalConstraintEEGModel().to(DEVICE)
+            self.log(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")
+            
+            optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)
+            criterion = nn.BCEWithLogitsLoss()
+            scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=20, T_mult=2)
+            
+            best_val_loss = float('inf')
+            
+            for epoch in range(self.epochs_var.get()):
+                if self.stop_flag.is_set():
+                    break
+                
+                # Train
+                model.train()
+                train_loss = 0
+                for eeg_windows, labels in train_loader:
+                    if self.stop_flag.is_set():
+                        break
+                    
+                    eeg_windows = [w.to(DEVICE) for w in eeg_windows]
+                    labels = labels.to(DEVICE)
+                    
+                    optimizer.zero_grad()
+                    logits, _ = model(eeg_windows)
+                    loss = criterion(logits, labels)
+                    loss.backward()
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                    optimizer.step()
+                    
+                    train_loss += loss.item()
+                
+                # Validate
+                model.eval()
+                val_loss = 0
+                with torch.no_grad():
+                    for eeg_windows, labels in val_loader:
+                        eeg_windows = [w.to(DEVICE) for w in eeg_windows]
+                        labels = labels.to(DEVICE)
+                        logits, _ = model(eeg_windows)
+                        loss = criterion(logits, labels)
+                        val_loss += loss.item()
+                
+                train_loss /= len(train_loader)
+                val_loss /= len(val_loader)
+                
+                scheduler.step()
+                
+                self.progress['value'] = ((epoch + 1) / self.epochs_var.get()) * 100
+                
+                if epoch % 5 == 0:
+                    self.log(f"Epoch {epoch+1}/{self.epochs_var.get()}: "
+                            f"TrLoss={train_loss:.4f} ValLoss={val_loss:.4f}")
+                
+                if val_loss < best_val_loss:
+                    best_val_loss = val_loss
+                    torch.save({
+                        'model_state_dict': model.state_dict(),
+                        'val_loss': val_loss,
+                        'epoch': epoch
+                    }, "calcium_bridge_eeg_model_v2.pth") # Save as V2
+                    if epoch % 5 == 0:
+                        self.log(f"  -> Saved V2 model (val_loss={val_loss:.4f})")
+            
+            self.log("\n" + "="*70)
+            self.log("TRAINING COMPLETE")
+            self.log(f"Best Val Loss: {best_val_loss:.4f}")
+            self.log("="*70)
+            
+        except Exception as e:
+            self.log(f"ERROR: {e}")
+            import traceback
+            self.log(traceback.format_exc())
+        finally:
+            self.train_btn.config(state=tk.NORMAL)
+            self.stop_btn.config(state=tk.DISABLED)
+
+
+if __name__ == "__main__":
+    app = CalciumEEGTrainerGUI()
+    app.mainloop()

pkas_cal_viewer_gemini2.py

@@ -0,0 +1,349 @@
+"""
+Calcium-Bridge EEG Constraint Viewer (V2.1 - Fixed)
+Visualizes how constraint satisfaction unfolds across four temporal windows up to 550ms.
+
+Shows:
+1. Original COCO image
+2. EEG heatmaps for each of the 4 time windows
+3. Calcium "attention" evolution (what the model focuses on at each stage)
+4. Top predictions crystallizing across the 4 windows
+
+V2.1 Fixes:
+- Corrected 'figsize' argument placement during figure creation.
+- Corrected colorbar creation to use the figure object directly, resolving warnings.
+"""
+
+import os
+import tkinter as tk
+from tkinter import filedialog, messagebox, ttk
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image, ImageTk
+import matplotlib.pyplot as plt
+from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
+import json
+from pathlib import Path
+from collections import defaultdict
+import random
+
+try:
+    from datasets import load_dataset
+except ImportError:
+    print("Missing datasets library.")
+    exit()
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+EEG_SAMPLE_RATE = 512
+
+TIME_WINDOWS = [
+    (50, 150, "EarlyVisual"),
+    (150, 250, "MidFeature"),
+    (250, 350, "LateSemantic"),
+    (350, 550, "CognitiveEvaluation")
+]
+
+TARGET_CATEGORIES = {
+    'elephant': 22, 'giraffe': 25, 'bear': 23, 'zebra': 24,
+    'cow': 21, 'sheep': 20, 'horse': 19, 'dog': 18, 'cat': 17, 'bird': 16,
+    'airplane': 5, 'train': 7, 'boat': 9, 'bus': 6, 'truck': 8,
+    'motorcycle': 4, 'bicycle': 2, 'car': 3,
+    'traffic light': 10, 'fire hydrant': 11, 'stop sign': 13,
+    'parking meter': 14, 'bench': 15,
+}
+
+CATEGORY_NAMES = {v: k for k, v in TARGET_CATEGORIES.items()}
+TARGET_IDS = set(TARGET_CATEGORIES.values())
+ALL_COCO_IDS = list(range(1, 91))
+EXCLUDED_IDS = set(ALL_COCO_IDS) - TARGET_IDS
+
+
+# === MODEL ARCHITECTURE (Must match V2 training code) ===
+
+class CalciumAttentionModule(nn.Module):
+    def __init__(self, n_features, d_model=256):
+        super().__init__()
+        self.n_features = n_features
+        self.d_model = d_model
+        self.phase_proj = nn.Linear(n_features, d_model)
+        self.ca_gate = nn.Sequential(
+            nn.Linear(d_model, d_model // 2),
+            nn.Sigmoid()
+        )
+        self.W = nn.Parameter(torch.randn(d_model, d_model) * 0.01)
+        self.norm = nn.LayerNorm(d_model)
+        
+    def forward(self, x, prev_ca=None, prev_W=None):
+        batch_size = x.size(0)
+        phi = self.phase_proj(x)
+        
+        if prev_ca is None:
+            ca = torch.zeros(batch_size, self.d_model, device=x.device)
+        else:
+            ca = prev_ca.clone()
+        
+        W = self.W if prev_W is None else prev_W
+        
+        coherence = torch.abs(torch.cos(phi[:, :, None] - phi[:, None, :]))
+        ca_update = torch.mean(coherence, dim=2)
+        ca = ca * 0.95 + ca_update * 0.05
+        
+        ca_gate = self.ca_gate(ca)
+        coupled = torch.matmul(phi, W)
+        ca_gate_full = torch.cat([ca_gate, ca_gate], dim=1)
+        features = coupled * ca_gate_full
+        features = self.norm(features + phi)
+        
+        return features, ca, W
+
+
+class TemporalConstraintEEGModel(nn.Module):
+    def __init__(self, n_channels=64, num_classes=len(TARGET_CATEGORIES)):
+        super().__init__()
+        self.n_channels = n_channels
+        
+        self.window_encoders = nn.ModuleList([
+            self._build_cnn_encoder() for _ in TIME_WINDOWS
+        ])
+        
+        self.ca_modules = nn.ModuleList([
+            CalciumAttentionModule(256, d_model=256) for _ in TIME_WINDOWS
+        ])
+        
+        self.classifier = nn.Sequential(
+            nn.Linear(256 * len(TIME_WINDOWS), 512),
+            nn.BatchNorm1d(512),
+            nn.GELU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, num_classes)
+        )
+        
+    def _build_cnn_encoder(self):
+        return nn.Sequential(
+            nn.Conv1d(self.n_channels, 128, kernel_size=15, padding=7),
+            nn.BatchNorm1d(128),
+            nn.ELU(),
+            nn.MaxPool1d(2),
+            nn.Conv1d(128, 256, kernel_size=7, padding=3),
+            nn.BatchNorm1d(256),
+            nn.ELU(),
+            nn.AdaptiveAvgPool1d(1)
+        )
+    
+    def forward(self, eeg_windows, return_intermediates=False):
+        batch_size = eeg_windows[0].size(0)
+        
+        window_features, ca_history, W_history, window_logits_list = [], [], [], []
+        ca_state, W_state = None, None
+        
+        for i, (encoder, ca_module, eeg_window) in enumerate(
+            zip(self.window_encoders, self.ca_modules, eeg_windows)
+        ):
+            cnn_features = encoder(eeg_window).squeeze(-1)
+            features, ca_state, W_state = ca_module(cnn_features, ca_state, W_state)
+            
+            window_features.append(features)
+            if return_intermediates:
+                ca_history.append(ca_state.detach())
+                W_history.append(W_state.detach())
+                
+                padded_features = window_features + [
+                    torch.zeros_like(features) for _ in range(len(TIME_WINDOWS) - len(window_features))
+                ]
+                intermediate_logits = self.classifier(torch.cat(padded_features, dim=1))
+                window_logits_list.append(intermediate_logits.detach())
+        
+        combined = torch.cat(window_features, dim=1)
+        logits = self.classifier(combined)
+        
+        if return_intermediates:
+            return logits, ca_history, W_history, window_logits_list
+        return logits, ca_history
+
+
+# === DATA LOADER ===
+class FilteredTestDataset:
+    def __init__(self, annotations_path, max_samples=1000):
+        print("Loading and filtering test dataset...")
+        self.eeg_dataset = load_dataset("Alljoined/05_125", split='test', streaming=False).select(range(max_samples))
+        with open(annotations_path, 'r') as f:
+            coco_data = json.load(f)
+        
+        image_annotations = defaultdict(set)
+        for ann in coco_data['annotations']:
+            image_annotations[ann['image_id']].add(ann['category_id'])
+
+        self.filtered_samples = []
+        for idx, sample in enumerate(self.eeg_dataset):
+            ann_ids = image_annotations.get(sample['coco_id'], set())
+            if not any(cat_id in EXCLUDED_IDS for cat_id in ann_ids) and any(cat_id in TARGET_IDS for cat_id in ann_ids):
+                self.filtered_samples.append({
+                    'coco_id': sample['coco_id'],
+                    'eeg_data': np.array(sample['EEG'], dtype=np.float32)
+                })
+        print(f"Loaded {len(self.filtered_samples)} filtered test samples.")
+        if not self.filtered_samples: raise RuntimeError("No suitable test samples found.")
+
+    def get_eeg_windows(self, sample_info):
+        eeg_data = sample_info['eeg_data']
+        eeg_windows = []
+        for start_ms, end_ms, _ in TIME_WINDOWS:
+            start_idx, end_idx = int(start_ms / 1000 * EEG_SAMPLE_RATE), int(end_ms / 1000 * EEG_SAMPLE_RATE)
+            n_timepoints = end_idx - start_idx
+            window = eeg_data[:, start_idx:end_idx] if eeg_data.shape[1] >= end_idx else eeg_data[:, start_idx:]
+            
+            if window.shape[1] != n_timepoints:
+                pad_width = n_timepoints - window.shape[1]
+                window = np.pad(window, ((0,0), (0, pad_width)), 'edge') if pad_width > 0 else window[:, :n_timepoints]
+            
+            window = (window - window.mean(axis=1, keepdims=True)) / (window.std(axis=1, keepdims=True) + 1e-8)
+            eeg_windows.append(window)
+        return eeg_windows
+
+    def get_random_sample_info(self):
+        return random.choice(self.filtered_samples)
+
+# === VIEWER APPLICATION ===
+class CalciumBridgeViewer(tk.Tk):
+    def __init__(self):
+        super().__init__()
+        self.title("Calcium-Bridge EEG Constraint Viewer V2 (Extended Window)")
+        self.geometry("2000x1000")
+        self.model, self.test_data = None, None
+        self.setup_gui()
+
+    def setup_gui(self):
+        control_frame = ttk.Frame(self); control_frame.pack(pady=10, padx=10, fill=tk.X)
+        ttk.Label(control_frame, text="COCO Path:").pack(side=tk.LEFT, padx=5)
+        self.coco_var = tk.StringVar(); ttk.Entry(control_frame, textvariable=self.coco_var, width=20).pack(side=tk.LEFT, padx=2)
+        ttk.Button(control_frame, text="Browse", command=self.browse_coco).pack(side=tk.LEFT, padx=5)
+        ttk.Label(control_frame, text="Annotations:").pack(side=tk.LEFT, padx=5)
+        self.ann_var = tk.StringVar(); ttk.Entry(control_frame, textvariable=self.ann_var, width=20).pack(side=tk.LEFT, padx=2)
+        ttk.Button(control_frame, text="Browse", command=self.browse_ann).pack(side=tk.LEFT, padx=5)
+        ttk.Button(control_frame, text="Load V2 Model", command=self.load_model).pack(side=tk.LEFT, padx=20)
+        self.test_btn = ttk.Button(control_frame, text="Test Random Sample", command=self.test_sample, state=tk.DISABLED); self.test_btn.pack(side=tk.LEFT, padx=5)
+        self.status_label = tk.Label(control_frame, text="Model: Not loaded", fg="gray"); self.status_label.pack(side=tk.LEFT, padx=20)
+
+        main_paned = ttk.PanedWindow(self, orient=tk.HORIZONTAL); main_paned.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
+        image_frame = ttk.Frame(main_paned, width=400); main_paned.add(image_frame, weight=0)
+        ttk.Label(image_frame, text="COCO Image", font=("Arial", 12, "bold")).pack(pady=5)
+        self.image_canvas = tk.Canvas(image_frame, width=400, height=400, bg='lightgray'); self.image_canvas.pack()
+        self.coco_id_label = ttk.Label(image_frame, text="COCO ID: N/A"); self.coco_id_label.pack(pady=5)
+        
+        self.notebook = ttk.Notebook(main_paned); main_paned.add(self.notebook, weight=1)
+        self.create_tabs()
+
+    def create_tabs(self):
+        self.constraint_fig, self.constraint_canvas = self.create_tab("Constraint Satisfaction", "How predictions crystallize as constraints are satisfied")
+        self.calcium_fig, self.calcium_canvas = self.create_tab("Calcium Attention", "Calcium state evolution: What the model 'focuses on' at each stage")
+        self.eeg_fig, self.eeg_canvas = self.create_tab("EEG Heatmaps", "Raw EEG signals for each time window")
+
+    def create_tab(self, title, description):
+        tab = ttk.Frame(self.notebook); self.notebook.add(tab, text=title)
+        ttk.Label(tab, text=description, font=("Arial", 11)).pack(pady=5)
+        fig = plt.Figure()
+        canvas = FigureCanvasTkAgg(fig, tab); canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)
+        return fig, canvas
+
+    def browse_coco(self):
+        path = filedialog.askdirectory(); self.coco_var.set(path); self.coco_path = path
+    
+    def browse_ann(self):
+        path = filedialog.askopenfilename(filetypes=[("JSON", "*.json")]); self.ann_var.set(path); self.annotations_path = path
+
+    def load_model(self):
+        model_path = filedialog.askopenfilename(filetypes=[("PyTorch Model", "*.pth")], title="Select calcium_bridge_eeg_model_v2.pth")
+        if not model_path or not self.annotations_path: return
+        try:
+            checkpoint = torch.load(model_path, map_location=DEVICE)
+            self.model = TemporalConstraintEEGModel().to(DEVICE)
+            self.model.load_state_dict(checkpoint['model_state_dict'])
+            self.model.eval()
+            self.test_data = FilteredTestDataset(self.annotations_path)
+            self.status_label.config(text="Model: V2 Loaded ✓", fg="green")
+            self.test_btn.config(state=tk.NORMAL)
+        except Exception as e: messagebox.showerror("Error", f"Failed to load model:\n{e}"); print(traceback.format_exc())
+
+    def _fetch_image(self, coco_id):
+        formatted_id = f"{coco_id:012d}.jpg"
+        for s in ["train2017", "val2017", "test2017"]:
+            path = os.path.join(self.coco_path, s, formatted_id)
+            if os.path.exists(path): return Image.open(path).convert("RGB")
+        return None
+
+    def test_sample(self):
+        if not self.model: return
+        try:
+            sample_info = self.test_data.get_random_sample_info()
+            image = self._fetch_image(sample_info['coco_id'])
+            if image: self.display_image(image, sample_info['coco_id'])
+            
+            eeg_windows_np = self.test_data.get_eeg_windows(sample_info)
+            eeg_windows = [torch.from_numpy(w).unsqueeze(0).to(DEVICE) for w in eeg_windows_np]
+            
+            with torch.no_grad():
+                logits, ca_history, _, window_logits = self.model(eeg_windows, return_intermediates=True)
+            
+            self.visualize_constraint_satisfaction(window_logits, logits)
+            self.visualize_calcium_evolution(ca_history)
+            self.visualize_eeg_heatmaps(eeg_windows_np)
+        except Exception as e: messagebox.showerror("Error", f"Failed to process sample:\n{e}"); print(traceback.format_exc())
+
+    def display_image(self, image, coco_id):
+        ratio = min(400/image.width, 400/image.height)
+        resized = image.resize((int(image.width * ratio), int(image.height * ratio)), Image.LANCZOS)
+        self.pil_image_tk = ImageTk.PhotoImage(resized)
+        self.image_canvas.create_image(200, 200, image=self.pil_image_tk)
+        self.coco_id_label.config(text=f"COCO ID: {coco_id}")
+
+    def visualize_constraint_satisfaction(self, window_logits, final_logits):
+        self.constraint_fig.clear()
+        cat_list = list(TARGET_CATEGORIES.keys())
+        n_windows = len(window_logits)
+        final_probs = torch.sigmoid(final_logits).squeeze(0).cpu().numpy()
+        top_indices = np.argsort(final_probs)[::-1][:10]
+        axes = self.constraint_fig.subplots(1, n_windows + 1)
+        
+        for i, (ax, wl) in enumerate(zip(axes[:-1], window_logits)):
+            probs = torch.sigmoid(wl).squeeze(0).cpu().numpy()[top_indices]
+            ax.barh([cat_list[idx] for idx in top_indices], probs, color='steelblue')
+            ax.set_title(f"{TIME_WINDOWS[i][2]}\n({TIME_WINDOWS[i][0]}-{TIME_WINDOWS[i][1]}ms)", fontsize=10)
+            ax.set_xlim(0, 1); ax.invert_yaxis(); ax.tick_params(axis='y', labelsize=8)
+        
+        axes[-1].barh([cat_list[idx] for idx in top_indices], final_probs[top_indices], color='darkgreen')
+        axes[-1].set_title("Final\n(Combined)", fontsize=10); axes[-1].set_xlim(0, 1); axes[-1].invert_yaxis(); axes[-1].tick_params(axis='y', labelsize=8)
+        self.constraint_fig.suptitle("Constraint Satisfaction: Predictions Crystallizing Over Time", fontsize=14, fontweight='bold')
+        self.constraint_fig.tight_layout(); self.constraint_canvas.draw()
+
+    def visualize_calcium_evolution(self, ca_history):
+        self.calcium_fig.clear()
+        n_windows = len(ca_history)
+        axes = self.calcium_fig.subplots(2, n_windows)
+        
+        for i, ca_state in enumerate(ca_history):
+            ca_np = ca_state.squeeze(0).cpu().numpy()
+            top_20_idx = np.argsort(ca_np)[::-1][:20]
+            axes[0, i].plot(ca_np, 'r'); axes[0, i].fill_between(range(len(ca_np)), ca_np, color='r', alpha=0.3)
+            axes[0, i].set_title(f"{TIME_WINDOWS[i][2]}\n({TIME_WINDOWS[i][0]}-{TIME_WINDOWS[i][1]}ms)", fontsize=10)
+            axes[1, i].barh([f"F{idx}" for idx in top_20_idx], ca_np[top_20_idx], color='darkred')
+            axes[1, i].invert_yaxis(); axes[1, i].tick_params(axis='y', labelsize=7)
+        self.calcium_fig.suptitle("Calcium Attention: What the Model Focuses On", fontsize=14, fontweight='bold')
+        self.calcium_fig.tight_layout(); self.calcium_canvas.draw()
+    
+    def visualize_eeg_heatmaps(self, eeg_windows_np):
+        self.eeg_fig.clear()
+        n_windows = len(eeg_windows_np)
+        axes = self.eeg_fig.subplots(1, n_windows)
+        
+        for i, (ax, eeg_data) in enumerate(zip(axes, eeg_windows_np)):
+            im = ax.imshow(eeg_data, aspect='auto', cmap='RdBu_r', vmin=-3, vmax=3)
+            ax.set_title(f"{TIME_WINDOWS[i][2]}\n({TIME_WINDOWS[i][0]}-{TIME_WINDOWS[i][1]}ms)", fontsize=10)
+            if i == 0: ax.set_ylabel("Channel")
+            self.eeg_fig.colorbar(im, ax=ax) # CORRECTED
+        self.eeg_fig.suptitle("Raw EEG Signals by Time Window", fontsize=14, fontweight='bold')
+        self.eeg_fig.tight_layout(); self.eeg_canvas.draw()
+
+if __name__ == "__main__":
+    app = CalciumBridgeViewer()
+    app.mainloop()