model.py

"""
EEG-to-fMRI Translation Model.

MLP that maps EEG features (236-dim) to fMRI BOLD signals (96 regions).
Architecture: 236 → 512 → 256 → 128 → 96

Harvard-Oxford Cortical Atlas: 96 regions covering the full cortex.
"""

import torch
import torch.nn as nn
import numpy as np
import os
from typing import Optional


# Harvard-Oxford atlas region names (96 regions)
HARVARD_OXFORD_REGIONS = [
    "Frontal Pole",
    "Insular Cortex",
    "Superior Frontal Gyrus",
    "Middle Frontal Gyrus",
    "Inferior Frontal Gyrus pars triangularis",
    "Inferior Frontal Gyrus pars opercularis",
    "Precentral Gyrus",
    "Temporal Pole",
    "Superior Temporal Gyrus anterior",
    "Superior Temporal Gyrus posterior",
    "Middle Temporal Gyrus anterior",
    "Middle Temporal Gyrus posterior",
    "Middle Temporal Gyrus temporooccipital",
    "Inferior Temporal Gyrus anterior",
    "Inferior Temporal Gyrus posterior",
    "Inferior Temporal Gyrus temporooccipital",
    "Postcentral Gyrus",
    "Superior Parietal Lobule",
    "Supramarginal Gyrus anterior",
    "Supramarginal Gyrus posterior",
    "Angular Gyrus",
    "Lateral Occipital Cortex superior",
    "Lateral Occipital Cortex inferior",
    "Intracalcarine Cortex",
    "Frontal Medial Cortex",
    "Juxtapositional Lobule Cortex",
    "Subcallosal Cortex",
    "Paracingulate Gyrus",
    "Cingulate Gyrus anterior",
    "Cingulate Gyrus posterior",
    "Precuneous Cortex",
    "Cuneal Cortex",
    "Frontal Orbital Cortex",
    "Parahippocampal Gyrus anterior",
    "Parahippocampal Gyrus posterior",
    "Lingual Gyrus",
    "Temporal Fusiform Cortex anterior",
    "Temporal Fusiform Cortex posterior",
    "Temporal Occipital Fusiform Cortex",
    "Occipital Fusiform Gyrus",
    "Frontal Operculum Cortex",
    "Central Opercular Cortex",
    "Parietal Operculum Cortex",
    "Planum Polare",
    "Heschl Gyrus",
    "Planum Temporale",
    "Supracalcarine Cortex",
    "Occipital Pole",
    # Right hemisphere (mirrored)
    "Frontal Pole R",
    "Insular Cortex R",
    "Superior Frontal Gyrus R",
    "Middle Frontal Gyrus R",
    "Inferior Frontal Gyrus pars triangularis R",
    "Inferior Frontal Gyrus pars opercularis R",
    "Precentral Gyrus R",
    "Temporal Pole R",
    "Superior Temporal Gyrus anterior R",
    "Superior Temporal Gyrus posterior R",
    "Middle Temporal Gyrus anterior R",
    "Middle Temporal Gyrus posterior R",
    "Middle Temporal Gyrus temporooccipital R",
    "Inferior Temporal Gyrus anterior R",
    "Inferior Temporal Gyrus posterior R",
    "Inferior Temporal Gyrus temporooccipital R",
    "Postcentral Gyrus R",
    "Superior Parietal Lobule R",
    "Supramarginal Gyrus anterior R",
    "Supramarginal Gyrus posterior R",
    "Angular Gyrus R",
    "Lateral Occipital Cortex superior R",
    "Lateral Occipital Cortex inferior R",
    "Intracalcarine Cortex R",
    "Frontal Medial Cortex R",
    "Juxtapositional Lobule Cortex R",
    "Subcallosal Cortex R",
    "Paracingulate Gyrus R",
    "Cingulate Gyrus anterior R",
    "Cingulate Gyrus posterior R",
    "Precuneous Cortex R",
    "Cuneal Cortex R",
    "Frontal Orbital Cortex R",
    "Parahippocampal Gyrus anterior R",
    "Parahippocampal Gyrus posterior R",
    "Lingual Gyrus R",
    "Temporal Fusiform Cortex anterior R",
    "Temporal Fusiform Cortex posterior R",
    "Temporal Occipital Fusiform Cortex R",
    "Occipital Fusiform Gyrus R",
    "Frontal Operculum Cortex R",
    "Central Opercular Cortex R",
    "Parietal Operculum Cortex R",
    "Planum Polare R",
    "Heschl Gyrus R",
    "Planum Temporale R",
    "Supracalcarine Cortex R",
    "Occipital Pole R",
]

# Network assignments for each region
NETWORK_ASSIGNMENTS = {
    "Visual": [21, 22, 23, 31, 35, 38, 39, 46, 47, 69, 70, 71, 79, 83, 86, 87, 94, 95],
    "Motor": [6, 16, 17, 25, 54, 64, 65, 73],
    "Language": [4, 5, 8, 9, 10, 11, 18, 19, 20, 44, 45, 52, 53, 56, 57, 58, 59, 66, 67, 68, 92, 93],
    "Default Mode": [0, 24, 27, 28, 29, 30, 33, 34, 48, 72, 75, 76, 77, 78, 81, 82],
    "Executive": [1, 2, 3, 26, 32, 40, 41, 42, 49, 50, 51, 74, 80, 88, 89, 90],
}


class EEGtoFMRIModel(nn.Module):
    """
    MLP for translating EEG features to fMRI BOLD signals.

    Architecture:
        Input (236) → Linear(512) → BN → GELU → Dropout
        → Linear(256) → BN → GELU → Dropout
        → Linear(128) → BN → GELU → Dropout
        → Linear(96) → Sigmoid

    Output is in [0, 1] range representing normalized BOLD signal.
    """

    def __init__(
        self,
        input_dim: int = 236,
        output_dim: int = 96,
        hidden_dims: list[int] = [512, 256, 128],
        dropout: float = 0.2,
    ):
        super().__init__()

        self.input_dim = input_dim
        self.output_dim = output_dim

        layers = []
        prev_dim = input_dim

        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.BatchNorm1d(hidden_dim),
                nn.GELU(),
                nn.Dropout(dropout),
            ])
            prev_dim = hidden_dim

        layers.append(nn.Linear(prev_dim, output_dim))
        layers.append(nn.Sigmoid())

        self.network = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.network(x)


class TranslationModel:
    """
    Wrapper for loading and running the EEG-to-fMRI translation model.
    """

    def __init__(self, model_path: Optional[str] = None):
        self.device = torch.device("cpu")  # CPU inference — fast enough for MLP
        self.model = EEGtoFMRIModel()

        if model_path and os.path.exists(model_path):
            checkpoint = torch.load(model_path, map_location=self.device, weights_only=False)
            self.model.load_state_dict(checkpoint["model_state_dict"])
            self.feature_mean = checkpoint.get("feature_mean", None)
            self.feature_std = checkpoint.get("feature_std", None)
            print(f"Loaded trained model from {model_path}")
        else:
            self.feature_mean = None
            self.feature_std = None
            print("No trained model found — using untrained model")

        self.model.to(self.device)
        self.model.eval()

    def predict(self, features: np.ndarray) -> dict:
        """
        Translate EEG features to fMRI BOLD estimates.

        Args:
            features: (236,) feature vector from eeg_processor

        Returns:
            dict with region activations and network summaries
        """
        # Normalize features if we have training stats
        if self.feature_mean is not None and self.feature_std is not None:
            features = (features - self.feature_mean) / (self.feature_std + 1e-8)

        # Run model
        with torch.no_grad():
            x = torch.FloatTensor(features).unsqueeze(0).to(self.device)
            bold = self.model(x).squeeze(0).cpu().numpy()

        # Build response
        regions = []
        for i, (name, value) in enumerate(zip(HARVARD_OXFORD_REGIONS, bold)):
            # Find which network this region belongs to
            network = "Other"
            for net_name, indices in NETWORK_ASSIGNMENTS.items():
                if i in indices:
                    network = net_name
                    break

            regions.append({
                "index": i,
                "name": name,
                "bold_value": round(float(value), 4),
                "network": network,
            })

        # Network summaries
        networks = {}
        for net_name, indices in NETWORK_ASSIGNMENTS.items():
            values = [float(bold[i]) for i in indices if i < len(bold)]
            networks[net_name] = {
                "mean_activation": round(np.mean(values), 4) if values else 0,
                "max_activation": round(np.max(values), 4) if values else 0,
                "active_regions": sum(1 for v in values if v > 0.5),
                "total_regions": len(values),
            }

        # Top active regions
        sorted_regions = sorted(regions, key=lambda r: r["bold_value"], reverse=True)

        return {
            "bold_values": [round(float(v), 4) for v in bold],
            "regions": regions,
            "networks": networks,
            "top_regions": sorted_regions[:10],
            "global_mean_activation": round(float(np.mean(bold)), 4),
            "global_max_activation": round(float(np.max(bold)), 4),
        }