app.py

"""
Multimodal Brain Encoder - Gradio Application
=============================================
Full end-to-end system:
  Input → CLIP Features → Brain Prediction → ROI Analysis → LLM Q&A → Visualization

Uses real trained weights from NSD dataset.
LLM is an INTERPRETER only - grounded in model predictions, not independent.
"""

import os
import sys
import json
import time
import logging
import pickle
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from pathlib import Path
from datetime import datetime
from collections import OrderedDict

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ============================================================
# Configuration (must match training)
# ============================================================
MODEL_REPO = os.environ.get("MODEL_REPO", "ryu34/multimodal-brain-encoder")

ROI_NAMES = {
    1: "V1v", 2: "V1d", 3: "V2v", 4: "V2d", 5: "V3v", 6: "V3d", 7: "hV4",
    8: "EBA", 9: "FBA-1", 10: "FBA-2", 11: "mTL-bodies",
    12: "OFA", 13: "FFA-1", 14: "FFA-2", 15: "mTL-faces", 16: "aTL-faces",
    17: "OPA", 18: "PPA", 19: "RSC",
    20: "OWFA", 21: "VWFA-1", 22: "VWFA-2", 23: "mfs-words", 24: "mTL-words",
}

FUNCTIONAL_NETWORKS = {
    "early_visual": [1, 2, 3, 4, 5, 6, 7],
    "body_selective": [8, 9, 10, 11],
    "face_selective": [12, 13, 14, 15, 16],
    "place_selective": [17, 18, 19],
    "word_selective": [20, 21, 22, 23, 24],
}

# Known neuroscience associations for grounded Q&A
ROI_FUNCTIONS = {
    "V1v": "Primary visual cortex (ventral); processes basic visual features: edges, orientations, spatial frequencies",
    "V1d": "Primary visual cortex (dorsal); processes basic visual features with dorsal visual stream emphasis",
    "V2v": "Secondary visual cortex (ventral); processes texture, figure-ground segregation",
    "V2d": "Secondary visual cortex (dorsal); processes contour and border ownership",
    "V3v": "Third visual area (ventral); contributes to form perception and shape processing",
    "V3d": "Third visual area (dorsal); processes dynamic form and motion boundaries",
    "hV4": "Human V4; processes color, pattern, moderate object features, texture discrimination",
    "EBA": "Extrastriate Body Area; selectively responds to bodies and body parts",
    "FBA-1": "Fusiform Body Area 1; body processing in ventral temporal cortex",
    "FBA-2": "Fusiform Body Area 2; complementary body processing region",
    "mTL-bodies": "Medial temporal lobe body area; body recognition with memory component",
    "OFA": "Occipital Face Area; early face-selective processing, face parts detection",
    "FFA-1": "Fusiform Face Area 1; core face recognition and identity processing",
    "FFA-2": "Fusiform Face Area 2; complementary face processing, holistic face representation",
    "mTL-faces": "Medial temporal lobe face area; face recognition with episodic memory",
    "aTL-faces": "Anterior temporal lobe face area; person identity and semantic knowledge",
    "OPA": "Occipital Place Area; processes local scene elements and spatial boundaries",
    "PPA": "Parahippocampal Place Area; processes scenes, buildings, spatial layouts",
    "RSC": "Retrosplenial Cortex; spatial navigation, scene-to-map coordinate transformation",
    "OWFA": "Occipital Word Form Area; early visual word processing",
    "VWFA-1": "Visual Word Form Area 1; processes written words and letter strings",
    "VWFA-2": "Visual Word Form Area 2; higher-level word form processing",
    "mfs-words": "Mid-fusiform sulcus word area; intermediate word processing",
    "mTL-words": "Medial temporal lobe word area; word recognition with memory",
}

NETWORK_FUNCTIONS = {
    "early_visual": "Early visual processing: edges, orientations, spatial frequencies, textures, colors. Active for all visual stimuli.",
    "body_selective": "Body-selective cortex: responds to human bodies, body parts, biological motion. Key for person perception.",
    "face_selective": "Face-selective cortex: responds to faces, facial features, identity. Critical for social perception.",
    "place_selective": "Place/scene-selective cortex: responds to spatial layouts, buildings, scenes, navigation cues.",
    "word_selective": "Word/reading-selective cortex: responds to written text, letter strings, word forms.",
}


# ============================================================
# Helper: enable only Dropout for MC sampling (keep BatchNorm in eval)
# ============================================================
def enable_dropout_only(model):
    """Enable Dropout layers while keeping BatchNorm in eval mode.
    
    This is needed for MC Dropout uncertainty estimation with batch_size=1,
    because BatchNorm1d requires batch_size > 1 in training mode.
    """
    for module in model.modules():
        if isinstance(module, nn.Dropout):
            module.train()


# ============================================================
# BrainEncoder model (must match training architecture exactly)
# ============================================================
class BrainEncoder(nn.Module):
    def __init__(self, input_dim=4096, n_voxels=15724, hidden_dims=None, dropout=0.3, n_rois=24):
        super().__init__()
        if hidden_dims is None:
            hidden_dims = [2048, 2048, 1024]
        
        self.input_dim = input_dim
        self.n_voxels = n_voxels
        self.n_rois = n_rois
        
        layers = []
        prev_dim = input_dim
        for h_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, h_dim),
                nn.BatchNorm1d(h_dim),
                nn.GELU(),
                nn.Dropout(dropout),
            ])
            prev_dim = h_dim
        self.backbone = nn.Sequential(*layers)
        
        self.general_head = nn.Linear(hidden_dims[-1], n_voxels)
        
        self.roi_attention = nn.ModuleDict()
        self.roi_heads = nn.ModuleDict()
        self.network_names = ["early_visual", "body_selective", "face_selective",
                              "place_selective", "word_selective"]
        
        for net_name in self.network_names:
            self.roi_attention[net_name] = nn.Sequential(
                nn.Linear(hidden_dims[-1], 256),
                nn.GELU(),
                nn.Linear(256, hidden_dims[-1]),
                nn.Sigmoid(),
            )
            self.roi_heads[net_name] = nn.Linear(hidden_dims[-1], n_voxels)
        
        self.register_buffer('roi_mask', torch.zeros(n_voxels, dtype=torch.long))
    
    def set_roi_assignments(self, annot):
        for net_idx, (net_name, roi_ids) in enumerate(FUNCTIONAL_NETWORKS.items()):
            for roi_id in roi_ids:
                mask = (annot == roi_id)
                if len(mask) <= self.n_voxels:
                    self.roi_mask[:len(mask)][mask[:self.n_voxels]] = net_idx + 1
    
    def forward(self, x, return_intermediates=False):
        intermediates = {}
        backbone_out = self.backbone(x)
        intermediates['backbone'] = backbone_out.detach()
        
        pred = self.general_head(backbone_out)
        intermediates['general_pred'] = pred.detach()
        
        for net_idx, net_name in enumerate(self.network_names):
            if net_name in self.roi_attention:
                attn = self.roi_attention[net_name](backbone_out)
                weighted = backbone_out * attn
                roi_pred = self.roi_heads[net_name](weighted)
                mask = (self.roi_mask == net_idx + 1)
                if mask.any():
                    pred[:, mask] = roi_pred[:, mask]
                intermediates[f'roi_{net_name}'] = roi_pred.detach()
        
        if return_intermediates:
            return pred, intermediates
        return pred


# ============================================================
# Model Manager - loads and caches models
# ============================================================
class ModelManager:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.brain_encoder = None
        self.ridge_model = None
        self.clip_model = None
        self.clip_processor = None
        self.roi_annotations = None
        self.config = None
        self._loaded = False
    
    def load(self):
        if self._loaded:
            return
        
        from huggingface_hub import hf_hub_download
        
        logger.info(f"Loading models from {MODEL_REPO}...")
        
        # Load config
        try:
            config_path = hf_hub_download(repo_id=MODEL_REPO, filename="config.json")
            with open(config_path) as f:
                self.config = json.load(f)
            logger.info(f"Config loaded: {self.config.get('architecture', 'unknown')}")
        except Exception as e:
            logger.warning(f"Config load failed: {e}")
            self.config = {}
        
        # Load ROI annotations
        try:
            annot_path = hf_hub_download(repo_id=MODEL_REPO, filename="roi_annotations.npy")
            self.roi_annotations = np.load(annot_path).flatten()
            logger.info(f"ROI annotations: {self.roi_annotations.shape}")
        except Exception as e:
            logger.warning(f"ROI annotations load failed: {e}")
        
        # Load brain encoder (optional - ridge is primary)
        try:
            model_path = hf_hub_download(repo_id=MODEL_REPO, filename="best_model.pt")
            checkpoint = torch.load(model_path, map_location=self.device, weights_only=False)
            
            model_config = checkpoint.get('config', {})
            self.brain_encoder = BrainEncoder(
                input_dim=model_config.get('input_dim', 4096),
                n_voxels=model_config.get('n_voxels', 15724),
                hidden_dims=model_config.get('hidden_dims', [2048, 2048, 1024]),
                dropout=model_config.get('dropout', 0.3),
            )
            self.brain_encoder.load_state_dict(checkpoint['model_state_dict'])
            self.brain_encoder.to(self.device).eval()
            
            if self.roi_annotations is not None:
                self.brain_encoder.set_roi_assignments(self.roi_annotations)
            
            # Free checkpoint memory
            del checkpoint
            
            logger.info("Brain encoder loaded successfully")
        except Exception as e:
            logger.warning(f"Brain encoder load failed (will use ridge only): {e}")
            self.brain_encoder = None
        
        # Load ridge model
        try:
            ridge_path = hf_hub_download(repo_id=MODEL_REPO, filename="ridge_model.pkl")
            with open(ridge_path, 'rb') as f:
                self.ridge_model = pickle.load(f)
            logger.info("Ridge model loaded successfully")
        except Exception as e:
            logger.warning(f"Ridge model load failed: {e}")
        
        # Load CLIP
        try:
            from transformers import CLIPModel, CLIPProcessor
            self.clip_model = CLIPModel.from_pretrained(
                "openai/clip-vit-large-patch14",
                torch_dtype=torch.float32,
            ).to(self.device).eval()
            self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
            logger.info("CLIP model loaded")
        except Exception as e:
            logger.error(f"CLIP load failed: {e}")
            raise
        
        self._loaded = True
        logger.info("All models loaded successfully!")
    
    def extract_features(self, image=None, text=None, audio=None):
        """Extract multimodal CLIP features."""
        features_dict = {}
        
        if image is not None:
            if isinstance(image, np.ndarray):
                image = Image.fromarray(image)
            
            inputs = self.clip_processor(images=image, return_tensors="pt")
            inputs = {k: v.to(self.device) for k, v in inputs.items() if isinstance(v, torch.Tensor)}
            
            with torch.no_grad():
                outputs = self.clip_model.vision_model(**inputs, output_hidden_states=True)
                
                cls_features = outputs.last_hidden_state[:, 0, :]
                projected = self.clip_model.visual_projection(cls_features)
                
                hidden_concat = []
                for layer_idx in [6, 12, 18, 23]:
                    h = outputs.hidden_states[layer_idx][:, 0, :]
                    hidden_concat.append(h)
                multi_layer = torch.cat(hidden_concat, dim=-1)
                
                features_dict['image_projected'] = projected.cpu().float()
                features_dict['image_multi_layer'] = multi_layer.cpu().float()
                features_dict['modality'] = 'image'
        
        if text is not None and text.strip():
            inputs = self.clip_processor(text=[text], return_tensors="pt", padding=True, truncation=True)
            inputs = {k: v.to(self.device) for k, v in inputs.items() if isinstance(v, torch.Tensor)}
            
            with torch.no_grad():
                text_outputs = self.clip_model.text_model(**inputs)
                pooled = text_outputs.pooler_output
                projected = self.clip_model.text_projection(pooled)
                
                # For text, repeat the projected features to match multi-layer dim
                # Text goes through the same brain encoder by tiling to 4096
                text_multi = projected.repeat(1, 4096 // projected.shape[1] + 1)[:, :4096]
                
                features_dict['text_projected'] = projected.cpu().float()
                features_dict['text_multi_layer'] = text_multi.cpu().float()
                if 'modality' not in features_dict:
                    features_dict['modality'] = 'text'
                else:
                    features_dict['modality'] = 'image+text'
        
        if audio is not None:
            # Convert audio to spectrogram image for CLIP processing
            sr, audio_data = audio if isinstance(audio, tuple) else (16000, audio)
            
            if len(audio_data.shape) > 1:
                audio_data = audio_data.mean(axis=1)
            audio_data = audio_data.astype(np.float32)
            
            # Create mel spectrogram visualization
            import matplotlib
            matplotlib.use('Agg')
            import matplotlib.pyplot as plt
            
            fig, ax = plt.subplots(1, 1, figsize=(4, 4))
            
            # Simple spectrogram using STFT
            n_fft = min(1024, len(audio_data))
            hop_length = n_fft // 4
            
            if len(audio_data) > n_fft:
                # Manual STFT
                n_frames = (len(audio_data) - n_fft) // hop_length + 1
                spec = np.zeros((n_fft // 2 + 1, n_frames))
                window = np.hanning(n_fft)
                
                for i in range(n_frames):
                    start = i * hop_length
                    frame = audio_data[start:start + n_fft] * window
                    fft = np.fft.rfft(frame)
                    spec[:, i] = np.abs(fft)
                
                spec_db = 20 * np.log10(spec + 1e-10)
                ax.imshow(spec_db, aspect='auto', origin='lower', cmap='viridis')
            else:
                ax.plot(audio_data[:1000])
            
            ax.set_title("Audio Spectrogram")
            ax.axis('off')
            
            fig.canvas.draw()
            
            # Convert to image
            buf = fig.canvas.buffer_rgba()
            spec_img = Image.frombytes('RGBA', fig.canvas.get_width_height(), buf).convert('RGB')
            plt.close(fig)
            
            # Process through CLIP as image
            inputs = self.clip_processor(images=spec_img, return_tensors="pt")
            inputs = {k: v.to(self.device) for k, v in inputs.items() if isinstance(v, torch.Tensor)}
            
            with torch.no_grad():
                outputs = self.clip_model.vision_model(**inputs, output_hidden_states=True)
                cls_features = outputs.last_hidden_state[:, 0, :]
                projected = self.clip_model.visual_projection(cls_features)
                
                hidden_concat = []
                for layer_idx in [6, 12, 18, 23]:
                    h = outputs.hidden_states[layer_idx][:, 0, :]
                    hidden_concat.append(h)
                multi_layer = torch.cat(hidden_concat, dim=-1)
                
                features_dict['audio_projected'] = projected.cpu().float()
                features_dict['audio_multi_layer'] = multi_layer.cpu().float()
                
                if features_dict.get('modality') is None:
                    features_dict['modality'] = 'audio'
                else:
                    features_dict['modality'] = features_dict['modality'] + '+audio'
        
        return features_dict
    
    def predict_brain_activity(self, features_dict):
        """Run brain encoder forward pass using BOTH ridge and deep models."""
        # Determine which features to use
        if 'image_multi_layer' in features_dict:
            input_features = features_dict['image_multi_layer']
        elif 'text_multi_layer' in features_dict:
            input_features = features_dict['text_multi_layer']
        elif 'audio_multi_layer' in features_dict:
            input_features = features_dict['audio_multi_layer']
        else:
            raise ValueError("No features available for prediction")
        
        # If multimodal, average features
        all_modality_features = []
        for key in ['image_multi_layer', 'text_multi_layer', 'audio_multi_layer']:
            if key in features_dict:
                all_modality_features.append(features_dict[key])
        
        if len(all_modality_features) > 1:
            input_features = torch.mean(torch.stack(all_modality_features), dim=0)
        
        input_features_np = input_features.cpu().numpy()
        input_features = input_features.to(self.device)
        
        # ── Primary: Ridge Model (proven baseline from Algonauts 2023) ──
        if self.ridge_model is not None:
            ridge = self.ridge_model
            X_norm = (input_features_np - ridge['feat_mean']) / ridge['feat_std']
            pred_z = ridge['model'].predict(X_norm)
            pred_np = (pred_z * ridge['fmri_std'] + ridge['fmri_mean']).flatten()
            
            # Clip extreme values for better visualization (keep 99.5th percentile)
            clip_val = np.percentile(np.abs(pred_np), 99.5)
            pred_np = np.clip(pred_np, -clip_val, clip_val)
        else:
            # Fallback to deep encoder
            with torch.no_grad():
                predictions, _ = self.brain_encoder(input_features, return_intermediates=True)
            pred_np = predictions.cpu().numpy().flatten()
        
        # ── Deep encoder for intermediates and uncertainty ──
        intermediates = {}
        if self.brain_encoder is not None:
            with torch.no_grad():
                deep_pred, intermediates = self.brain_encoder(input_features, return_intermediates=True)
            
            # Compute uncertainty via MC Dropout
            # IMPORTANT: Only enable Dropout layers, keep BatchNorm in eval mode.
            # BatchNorm1d requires batch_size > 1 in training mode, but we have batch_size=1.
            self.brain_encoder.eval()  # Ensure everything is in eval mode first
            enable_dropout_only(self.brain_encoder)  # Selectively enable only Dropout
            mc_predictions = []
            for _ in range(10):
                with torch.no_grad():
                    mc_pred = self.brain_encoder(input_features)
                    mc_predictions.append(mc_pred.cpu().numpy().flatten())
            self.brain_encoder.eval()  # Restore full eval mode
            mc_predictions = np.array(mc_predictions)
            uncertainty = np.std(mc_predictions, axis=0)
        else:
            # Estimate uncertainty from ridge prediction variance across feature perturbation
            ridge = self.ridge_model
            mc_predictions = []
            for _ in range(10):
                noise = np.random.normal(0, 0.01, size=input_features_np.shape)
                X_noisy = (input_features_np + noise - ridge['feat_mean']) / ridge['feat_std']
                mp = ridge['model'].predict(X_noisy).flatten()
                mc_predictions.append(mp)
            mc_predictions = np.array(mc_predictions)
            uncertainty = np.std(mc_predictions, axis=0)
        
        # Compute modality contributions using ridge
        modality_contributions = {}
        if self.ridge_model is not None:
            ridge = self.ridge_model
            for key in ['image_multi_layer', 'text_multi_layer', 'audio_multi_layer']:
                if key in features_dict:
                    modality_name = key.split('_')[0]
                    feat_np = features_dict[key].cpu().numpy()
                    X_n = (feat_np - ridge['feat_mean']) / ridge['feat_std']
                    mp = (ridge['model'].predict(X_n) * ridge['fmri_std'] + ridge['fmri_mean']).flatten()
                    clip_val_mod = np.percentile(np.abs(mp), 99.5)
                    mp = np.clip(mp, -clip_val_mod, clip_val_mod)
                    modality_contributions[modality_name] = mp
        
        # Compute ROI summaries using z-scored per-voxel predictions
        # This shows which regions are MORE or LESS activated compared to baseline
        if self.ridge_model is not None:
            baseline_mean = self.ridge_model['fmri_mean']
            baseline_std = self.ridge_model['fmri_std']
            # Z-score predictions relative to training distribution
            n_v = min(len(pred_np), len(baseline_mean))
            pred_z = (pred_np[:n_v] - baseline_mean[:n_v]) / (baseline_std[:n_v] + 1e-8)
        else:
            pred_z = pred_np
        
        roi_summary = self._compute_roi_summary(pred_z, uncertainty)
        
        # Validation checks
        warnings = self._validate_predictions(pred_np)
        
        result = {
            'predictions': pred_np,
            'uncertainty': uncertainty,
            'roi_summary': roi_summary,
            'modality_contributions': modality_contributions,
            'modality': features_dict.get('modality', 'unknown'),
            'intermediates': {k: v.cpu().numpy() if torch.is_tensor(v) else v 
                           for k, v in intermediates.items()},
            'warnings': warnings,
            'timestamp': datetime.now().isoformat(),
        }
        
        return result
    
    def _compute_roi_summary(self, predictions, uncertainty):
        """Compute per-ROI activation summaries."""
        if self.roi_annotations is None:
            return {}
        
        annot = self.roi_annotations
        n_voxels = len(predictions)
        
        roi_summary = {}
        
        for roi_id, roi_name in ROI_NAMES.items():
            mask = (annot[:n_voxels] == roi_id) if len(annot) >= n_voxels else np.zeros(n_voxels, dtype=bool)
            
            if mask.sum() == 0:
                continue
            
            roi_activations = predictions[mask]
            roi_uncertainty = uncertainty[mask]
            
            roi_summary[roi_name] = {
                'mean_activation': float(np.mean(roi_activations)),
                'max_activation': float(np.max(roi_activations)),
                'min_activation': float(np.min(roi_activations)),
                'std_activation': float(np.std(roi_activations)),
                'mean_uncertainty': float(np.mean(roi_uncertainty)),
                'n_voxels': int(mask.sum()),
                'abs_mean': float(np.mean(np.abs(roi_activations))),
                'known_function': ROI_FUNCTIONS.get(roi_name, "Unknown"),
            }
        
        return roi_summary
    
    def _validate_predictions(self, predictions):
        """Validation safeguards."""
        warnings = []
        
        if np.std(predictions) < 1e-6:
            warnings.append("⚠️ CONSTANT OUTPUT DETECTED: All voxels have near-identical values")
        
        if np.any(np.isnan(predictions)):
            warnings.append("⚠️ NaN VALUES DETECTED in predictions")
        
        if np.any(np.isinf(predictions)):
            warnings.append("⚠️ Infinite VALUES DETECTED in predictions")
        
        if np.max(np.abs(predictions)) > 50:
            warnings.append(f"⚠️ Unusually large activations detected (max |activation| = {np.max(np.abs(predictions)):.2f})")
        
        return warnings


# ============================================================
# Grounded Q&A System
# ============================================================
class GroundedQA:
    """
    RAG-grounded Q&A system.
    The LLM is an INTERPRETER - it only explains model predictions.
    It does NOT generate independent neuroscience claims.
    """
    
    def __init__(self):
        self.inference_client = None
        self._init_client()
    
    def _init_client(self):
        try:
            from huggingface_hub import InferenceClient
            self.inference_client = InferenceClient(
                provider="hf-inference",
                api_key=os.environ.get("HF_TOKEN", ""),
            )
            logger.info("HF Inference Client initialized")
        except Exception as e:
            logger.warning(f"Inference client init failed: {e}")
    
    def build_context(self, brain_result):
        """Build structured context from model predictions for LLM grounding."""
        
        roi_summary = brain_result.get('roi_summary', {})
        modality = brain_result.get('modality', 'unknown')
        warnings = brain_result.get('warnings', [])
        modality_contributions = brain_result.get('modality_contributions', {})
        
        # Sort ROIs by absolute mean activation
        sorted_rois = sorted(
            roi_summary.items(),
            key=lambda x: abs(x[1]['abs_mean']),
            reverse=True
        )
        
        # Top activated regions
        top_regions = []
        for roi_name, data in sorted_rois[:10]:
            top_regions.append(
                f"- {roi_name}: mean_activation={data['mean_activation']:.4f}, "
                f"abs_mean={data['abs_mean']:.4f}, uncertainty={data['mean_uncertainty']:.4f}, "
                f"n_voxels={data['n_voxels']}"
            )
        
        # Network-level summaries
        network_summaries = {}
        for net_name, roi_ids in FUNCTIONAL_NETWORKS.items():
            roi_names_in_net = [ROI_NAMES[rid] for rid in roi_ids if rid in ROI_NAMES]
            activations = []
            for rn in roi_names_in_net:
                if rn in roi_summary:
                    activations.append(roi_summary[rn]['abs_mean'])
            
            if activations:
                network_summaries[net_name] = {
                    'mean_abs_activation': np.mean(activations),
                    'max_abs_activation': np.max(activations),
                    'function': NETWORK_FUNCTIONS.get(net_name, ""),
                }
        
        sorted_networks = sorted(
            network_summaries.items(),
            key=lambda x: x[1]['mean_abs_activation'],
            reverse=True
        )
        
        # Modality contributions
        modality_info = ""
        if modality_contributions:
            modality_info = "\n## Modality Contributions\n"
            for mod_name, mod_pred in modality_contributions.items():
                modality_info += f"- {mod_name}: mean_abs_activation={np.mean(np.abs(mod_pred)):.4f}, std={np.std(mod_pred):.4f}\n"
        
        # Global prediction stats
        predictions = brain_result['predictions']
        global_stats = (
            f"- Total voxels predicted: {len(predictions)}\n"
            f"- Global mean activation: {np.mean(predictions):.4f}\n"
            f"- Global std: {np.std(predictions):.4f}\n"
            f"- Global range: [{np.min(predictions):.4f}, {np.max(predictions):.4f}]\n"
            f"- Mean uncertainty: {np.mean(brain_result['uncertainty']):.4f}\n"
        )
        
        context = f"""## Brain Activity Prediction Summary
Input modality: {modality}

## Global Statistics
{global_stats}

## Top 10 Activated Brain Regions (by absolute activation strength)
{chr(10).join(top_regions)}

## Functional Network Activations (ranked by strength)
"""
        for net_name, net_data in sorted_networks:
            context += (
                f"- {net_name}: mean_abs={net_data['mean_abs_activation']:.4f}, "
                f"max_abs={net_data['max_abs_activation']:.4f}\n"
                f"  Known function: {net_data['function']}\n"
            )
        
        context += modality_info
        
        if warnings:
            context += "\n## Warnings\n"
            for w in warnings:
                context += f"- {w}\n"
        
        # ROI functional labels
        context += "\n## ROI Functional Reference\n"
        for roi_name in [r[0] for r in sorted_rois[:10]]:
            if roi_name in ROI_FUNCTIONS:
                context += f"- {roi_name}: {ROI_FUNCTIONS[roi_name]}\n"
        
        return context
    
    def answer(self, question, brain_result):
        """Answer a question grounded in model predictions."""
        
        context = self.build_context(brain_result)
        
        system_prompt = """You are a neuroscience interpreter for a brain encoding model. 
Your role is STRICTLY to interpret and explain the model's predicted brain activity patterns.

CRITICAL RULES:
1. ONLY reference data provided in the context below. Never invent neuroscience claims.
2. Always distinguish between:
   - "Predicted activation" (what the model outputs)
   - "Known neuroscience association" (established findings about brain regions)
   - "Possible interpretation" (your inference connecting the two)
3. Include uncertainty statements. Use phrases like "the model predicts", "this is consistent with", "one possible interpretation is"
4. NEVER make definitive claims about emotions, consciousness, or behavior from brain activity alone.
5. Always cite specific regions, activation values, and confidence levels from the context.
6. If the question cannot be answered from the provided data, say so explicitly.
7. Keep answers concise but informative (2-4 paragraphs max).

You are an INTERPRETER of model outputs, not an independent neuroscience oracle."""

        user_prompt = f"""## Model Prediction Context
{context}

## User Question
{question}

Please answer based ONLY on the model prediction data above. Cite specific regions and values."""

        if self.inference_client is None:
            return self._fallback_answer(question, brain_result, context)
        
        try:
            response = self.inference_client.chat.completions.create(
                model="Qwen/Qwen2.5-72B-Instruct",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                max_tokens=800,
                temperature=0.3,
            )
            answer = response.choices[0].message.content
            
            # Add grounding footer
            answer += "\n\n---\n*This interpretation is based on model predictions with "
            mean_unc = np.mean(brain_result['uncertainty'])
            answer += f"mean uncertainty={mean_unc:.4f}. "
            answer += "Predictions are from a brain encoder trained on NSD (Natural Scenes Dataset) fMRI data.*"
            
            return answer
            
        except Exception as e:
            logger.warning(f"LLM inference failed: {e}")
            return self._fallback_answer(question, brain_result, context)
    
    def _fallback_answer(self, question, brain_result, context):
        """Structured fallback when LLM is unavailable."""
        roi_summary = brain_result.get('roi_summary', {})
        sorted_rois = sorted(
            roi_summary.items(),
            key=lambda x: abs(x[1]['abs_mean']),
            reverse=True
        )
        
        answer = "## Brain Activity Interpretation\n\n"
        answer += f"**Input modality:** {brain_result.get('modality', 'unknown')}\n\n"
        
        answer += "### Top Activated Regions\n"
        for roi_name, data in sorted_rois[:5]:
            answer += (
                f"- **{roi_name}** (activation={data['mean_activation']:.4f}, "
                f"uncertainty={data['mean_uncertainty']:.4f}): "
                f"{ROI_FUNCTIONS.get(roi_name, 'Unknown function')}\n"
            )
        
        answer += "\n### Network-Level Summary\n"
        for net_name, roi_ids in FUNCTIONAL_NETWORKS.items():
            roi_names_in_net = [ROI_NAMES[rid] for rid in roi_ids if rid in ROI_NAMES]
            activations = [roi_summary[rn]['abs_mean'] for rn in roi_names_in_net if rn in roi_summary]
            if activations:
                mean_act = np.mean(activations)
                answer += f"- **{net_name}**: mean_abs_activation={mean_act:.4f} — {NETWORK_FUNCTIONS.get(net_name, '')}\n"
        
        answer += f"\n*Note: LLM interpretation unavailable. Showing structured prediction summary. "
        answer += f"Mean uncertainty: {np.mean(brain_result['uncertainty']):.4f}*"
        
        return answer


# ============================================================
# Transparency Logger
# ============================================================
class TransparencyLogger:
    """Logs all inputs, intermediates, and outputs for traceability."""
    
    def __init__(self):
        self.logs = []
    
    def log_inference(self, inputs, features_dict, brain_result, qa_answer=None):
        entry = {
            'timestamp': datetime.now().isoformat(),
            'inputs': {
                'has_image': inputs.get('image') is not None,
                'has_text': inputs.get('text') is not None and inputs.get('text', '').strip() != '',
                'has_audio': inputs.get('audio') is not None,
                'text_content': inputs.get('text', '')[:200],
            },
            'features': {
                'modality': features_dict.get('modality', 'unknown'),
                'feature_norms': {},
            },
            'predictions': {
                'n_voxels': len(brain_result['predictions']),
                'pred_mean': float(np.mean(brain_result['predictions'])),
                'pred_std': float(np.std(brain_result['predictions'])),
                'pred_range': [float(np.min(brain_result['predictions'])), 
                              float(np.max(brain_result['predictions']))],
                'uncertainty_mean': float(np.mean(brain_result['uncertainty'])),
            },
            'roi_summary_sent_to_llm': list(brain_result.get('roi_summary', {}).keys()),
            'warnings': brain_result.get('warnings', []),
            'qa_answer_length': len(qa_answer) if qa_answer else 0,
        }
        
        # Feature norms
        for key in ['image_multi_layer', 'text_multi_layer', 'audio_multi_layer']:
            if key in features_dict:
                entry['features']['feature_norms'][key] = float(features_dict[key].norm().item())
        
        self.logs.append(entry)
        return entry
    
    def get_log_text(self):
        return json.dumps(self.logs[-5:], indent=2, default=str)


# ============================================================
# Visualization helpers
# ============================================================
def create_brain_activation_plot(brain_result, roi_annotations):
    """Create brain activation visualization."""
    import plotly.graph_objects as go
    from plotly.subplots import make_subplots
    
    roi_summary = brain_result.get('roi_summary', {})
    
    if not roi_summary:
        fig = go.Figure()
        fig.add_annotation(text="No ROI data available", x=0.5, y=0.5)
        return fig
    
    # Create multi-panel figure
    fig = make_subplots(
        rows=2, cols=2,
        subplot_titles=(
            "ROI Activation Strengths",
            "Functional Network Summary",
            "Activation Uncertainty",
            "Activation Distribution",
        ),
        specs=[
            [{"type": "bar"}, {"type": "bar"}],
            [{"type": "bar"}, {"type": "histogram"}],
        ]
    )
    
    # Panel 1: ROI activations
    sorted_rois = sorted(roi_summary.items(), key=lambda x: abs(x[1]['abs_mean']), reverse=True)[:15]
    roi_names = [r[0] for r in sorted_rois]
    roi_activations = [r[1]['mean_activation'] for r in sorted_rois]
    roi_colors = []
    for r in sorted_rois:
        name = r[0]
        for net_name, roi_ids in FUNCTIONAL_NETWORKS.items():
            roi_names_in_net = [ROI_NAMES[rid] for rid in roi_ids if rid in ROI_NAMES]
            if name in roi_names_in_net:
                color_map = {
                    "early_visual": "#4CAF50",
                    "body_selective": "#FF9800",
                    "face_selective": "#E91E63",
                    "place_selective": "#2196F3",
                    "word_selective": "#9C27B0",
                }
                roi_colors.append(color_map.get(net_name, "#666"))
                break
        else:
            roi_colors.append("#666")
    
    fig.add_trace(
        go.Bar(x=roi_names, y=roi_activations, marker_color=roi_colors, name="Activation"),
        row=1, col=1
    )
    
    # Panel 2: Network summary
    net_names = []
    net_activations = []
    net_colors_list = []
    color_map = {
        "early_visual": "#4CAF50",
        "body_selective": "#FF9800",
        "face_selective": "#E91E63",
        "place_selective": "#2196F3",
        "word_selective": "#9C27B0",
    }
    
    for net_name, roi_ids in FUNCTIONAL_NETWORKS.items():
        roi_names_in_net = [ROI_NAMES[rid] for rid in roi_ids if rid in ROI_NAMES]
        activations = [roi_summary[rn]['abs_mean'] for rn in roi_names_in_net if rn in roi_summary]
        if activations:
            net_names.append(net_name.replace("_", " ").title())
            net_activations.append(np.mean(activations))
            net_colors_list.append(color_map.get(net_name, "#666"))
    
    fig.add_trace(
        go.Bar(x=net_names, y=net_activations, marker_color=net_colors_list, name="Network"),
        row=1, col=2
    )
    
    # Panel 3: Uncertainty
    roi_uncertainty = [r[1]['mean_uncertainty'] for r in sorted_rois]
    fig.add_trace(
        go.Bar(x=roi_names, y=roi_uncertainty, marker_color='rgba(255,0,0,0.5)', name="Uncertainty"),
        row=2, col=1
    )
    
    # Panel 4: Distribution
    predictions = brain_result['predictions']
    fig.add_trace(
        go.Histogram(x=predictions[::10], nbinsx=50, name="Activations", marker_color='#4CAF50'),
        row=2, col=2
    )
    
    fig.update_layout(
        height=700,
        showlegend=False,
        title_text="Brain Activity Predictions",
        template="plotly_white",
    )
    
    return fig


def create_modality_contribution_plot(brain_result):
    """Create modality contribution visualization."""
    import plotly.graph_objects as go
    
    contributions = brain_result.get('modality_contributions', {})
    
    if len(contributions) <= 1:
        fig = go.Figure()
        fig.add_annotation(text="Single modality input - no comparison available", x=0.5, y=0.5)
        return fig
    
    fig = go.Figure()
    
    for mod_name, mod_pred in contributions.items():
        # Show distribution of activations per modality
        fig.add_trace(go.Histogram(
            x=mod_pred[::10],
            name=mod_name.capitalize(),
            opacity=0.6,
            nbinsx=50,
        ))
    
    fig.update_layout(
        title="Modality Contributions to Brain Activity",
        xaxis_title="Predicted Activation",
        yaxis_title="Count",
        barmode='overlay',
        template="plotly_white",
        height=400,
    )
    
    return fig


# ============================================================
# Gradio Application
# ============================================================
def build_gradio_app():
    import gradio as gr
    
    # Global state
    manager = ModelManager()
    qa_system = GroundedQA()
    transparency_log = TransparencyLogger()
    
    current_result = {"value": None}
    
    def initialize():
        try:
            manager.load()
            return "✅ Models loaded successfully!"
        except Exception as e:
            return f"❌ Error loading models: {e}"
    
    def process_input(image, text, audio):
        """Main inference pipeline."""
        if not manager._loaded:
            manager.load()
        
        if image is None and (text is None or text.strip() == '') and audio is None:
            return "Please provide at least one input (image, text, or audio).", None, None, ""
        
        try:
            # Step 1: Extract features
            features = manager.extract_features(image=image, text=text, audio=audio)
            
            # Step 2: Predict brain activity
            result = manager.predict_brain_activity(features)
            current_result["value"] = result
            
            # Step 3: Create visualizations
            brain_plot = create_brain_activation_plot(result, manager.roi_annotations)
            modality_plot = create_modality_contribution_plot(result)
            
            # Step 4: Log for transparency
            log_entry = transparency_log.log_inference(
                {'image': image, 'text': text, 'audio': audio},
                features, result
            )
            
            # Summary text
            roi_summary = result.get('roi_summary', {})
            sorted_rois = sorted(roi_summary.items(), key=lambda x: abs(x[1]['abs_mean']), reverse=True)
            
            summary = f"**Modality:** {result['modality']}\n"
            summary += f"**Voxels predicted:** {len(result['predictions'])}\n"
            summary += f"**Mean uncertainty:** {np.mean(result['uncertainty']):.4f}\n\n"
            summary += "**Top 5 Activated Regions:**\n"
            for roi_name, data in sorted_rois[:5]:
                summary += f"- {roi_name}: {data['mean_activation']:.4f} (±{data['mean_uncertainty']:.4f})\n"
            
            if result['warnings']:
                summary += "\n**Warnings:**\n"
                for w in result['warnings']:
                    summary += f"- {w}\n"
            
            return summary, brain_plot, modality_plot, json.dumps(log_entry, indent=2, default=str)
            
        except Exception as e:
            import traceback
            return f"Error: {e}\n{traceback.format_exc()}", None, None, ""
    
    def ask_question(question, history):
        """Q&A with grounded interpretation."""
        if current_result["value"] is None:
            history = history or []
            history.append({"role": "user", "content": question})
            history.append({"role": "assistant", "content": "Please run an inference first (provide an input in the Stimulus tab) before asking questions."})
            return history, ""
        
        history = history or []
        history.append({"role": "user", "content": question})
        
        answer = qa_system.answer(question, current_result["value"])
        history.append({"role": "assistant", "content": answer})
        
        # Log Q&A
        transparency_log.log_inference(
            {'text': question},
            {'modality': 'qa'},
            current_result["value"],
            qa_answer=answer,
        )
        
        return history, ""
    
    def get_transparency_log():
        return transparency_log.get_log_text()
    
    # Build UI
    with gr.Blocks(title="Multimodal Brain Encoder") as demo:
        gr.Markdown("""
        # 🧠 Multimodal Brain Encoder
        
        **A real brain encoding model trained on the Natural Scenes Dataset (NSD)**
        
        This system predicts brain activity (fMRI voxel responses) from multimodal inputs using:
        - **CLIP ViT-L/14** for feature extraction (multi-layer: layers 6, 12, 18, 24)
        - **Deep Brain Encoder** with ROI-specific attention heads (trained on NSD subj01)
        - **Ridge Regression** baseline (Algonauts 2023 recipe)
        - **Grounded LLM Q&A** that only interprets model predictions
        
        All predictions are from real model forward passes with learned weights.
        """)
        
        status = gr.Textbox(label="Status", value="Click 'Load Models' to initialize")
        load_btn = gr.Button("🔄 Load Models", variant="primary")
        load_btn.click(fn=initialize, outputs=status)
        
        with gr.Tabs():
            # Tab 1: Input & Prediction
            with gr.Tab("🎯 Stimulus Input & Brain Prediction"):
                with gr.Row():
                    with gr.Column(scale=1):
                        image_input = gr.Image(type="pil", label="Visual Stimulus (Image)")
                        text_input = gr.Textbox(
                            label="Text Input",
                            placeholder="Enter a description or sentence...",
                            lines=3,
                        )
                        audio_input = gr.Audio(type="numpy", label="Audio Input")
                        predict_btn = gr.Button("🧠 Predict Brain Activity", variant="primary", size="lg")
                    
                    with gr.Column(scale=2):
                        summary_output = gr.Markdown(label="Prediction Summary")
                        brain_plot = gr.Plot(label="Brain Activity Visualization")
                        modality_plot = gr.Plot(label="Modality Contributions")
                
                predict_btn.click(
                    fn=process_input,
                    inputs=[image_input, text_input, audio_input],
                    outputs=[summary_output, brain_plot, modality_plot, gr.Textbox(visible=False)],
                )
            
            # Tab 2: Q&A
            with gr.Tab("💬 Grounded Q&A"):
                gr.Markdown("""
                ### Ask questions about the predicted brain activity
                
                The LLM interpreter will answer based ONLY on:
                - Predicted activation maps and ROI summaries
                - Known functional labels from brain atlases
                - Modality attribution outputs
                - Uncertainty estimates
                
                It will NOT make independent neuroscience claims.
                """)
                
                chatbot = gr.Chatbot(
                    type="messages",
                    label="Brain Activity Q&A",
                    height=400,
                )
                
                with gr.Row():
                    question_input = gr.Textbox(
                        label="Your Question",
                        placeholder="e.g., Which brain regions are most activated? What does the face-selective network response mean?",
                        scale=4,
                    )
                    ask_btn = gr.Button("Ask", variant="primary", scale=1)
                
                ask_btn.click(
                    fn=ask_question,
                    inputs=[question_input, chatbot],
                    outputs=[chatbot, question_input],
                )
                question_input.submit(
                    fn=ask_question,
                    inputs=[question_input, chatbot],
                    outputs=[chatbot, question_input],
                )
                
                gr.Markdown("""
                **Example questions:**
                - "What are the most activated brain regions for this input?"
                - "Is the face-selective network responding? What might that mean?"
                - "How confident is the model in these predictions?"
                - "How does the visual input differ from the text input in brain response?"
                - "What does high PPA activation suggest about this image?"
                """)
            
            # Tab 3: Transparency Log
            with gr.Tab("📋 Transparency Log"):
                gr.Markdown("### Full inference traceability log")
                gr.Markdown("Every inference is logged with inputs, features, predictions, and Q&A answers.")
                
                log_output = gr.Code(language="json", label="Recent Logs")
                refresh_log_btn = gr.Button("🔄 Refresh Log")
                refresh_log_btn.click(fn=get_transparency_log, outputs=log_output)
            
            # Tab 4: Model Info
            with gr.Tab("ℹ️ Model Information"):
                gr.Markdown(f"""
                ### Architecture Details
                
                | Component | Details |
                |-----------|---------|
                | Feature Extractor | CLIP ViT-L/14 (openai/clip-vit-large-patch14) |
                | Feature Layers | Layers 6, 12, 18, 24 (CLS tokens concatenated = 4096-dim) |
                | Brain Encoder | 4096 → 2048 → 2048 → 1024 → N_voxels |
                | Activations | GELU + BatchNorm + Dropout(0.3) |
                | ROI Heads | 5 functional network heads with learned attention |
                | Ridge Baseline | sklearn RidgeCV with 17 alphas (1e-2 to 1e6) |
                | Training Data | NSD subj01 (~8,859 train, ~300 val images) |
                | fMRI Resolution | 7T, ~15,724 voxels (NSD general cortical mask) |
                | Uncertainty | MC Dropout (10 forward passes) |
                
                ### Brain Regions (24 ROIs from NSD)
                
                | Network | Regions | Function |
                |---------|---------|----------|
                | Early Visual | V1v, V1d, V2v, V2d, V3v, V3d, hV4 | Basic visual processing |
                | Body Selective | EBA, FBA-1, FBA-2, mTL-bodies | Body/person perception |
                | Face Selective | OFA, FFA-1, FFA-2, mTL-faces, aTL-faces | Face recognition |
                | Place Selective | OPA, PPA, RSC | Scene/navigation |
                | Word Selective | OWFA, VWFA-1, VWFA-2, mfs-words, mTL-words | Reading/text |
                
                ### References
                - Natural Scenes Dataset: Allen et al. 2022, Nature Neuroscience
                - Algonauts 2023: Gifford et al. 2023
                - CLIP: Radford et al. 2021
                - Model repo: [{MODEL_REPO}](https://huggingface.co/{MODEL_REPO})
                """)
    
    return demo


if __name__ == "__main__":
    demo = build_gradio_app()
    demo.launch(server_name="0.0.0.0", server_port=7860)