"""
Training script for EEG-to-fMRI translation model.

Uses publicly available paired EEG-fMRI datasets to train the MLP
to map EEG features → fMRI BOLD signals across 96 Harvard-Oxford regions.

Public datasets used:
1. Simultaneous EEG-fMRI dataset from OpenNeuro (ds003505)
   - Resting-state EEG-fMRI from multiple subjects
   - EEG: 64 channels, 5000 Hz
   - fMRI: 3T, TR=2s, Harvard-Oxford parcellation

Usage:
    python train.py --data_dir ./data --epochs 100 --output model_weights.pt

If no real data is available, generates structured synthetic paired data
that preserves known EEG-fMRI correlations (alpha-occipital, beta-motor, etc.)
"""

import argparse
import os
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from model import EEGtoFMRIModel
from eeg_processor import (
    extract_all_features,
    get_feature_dimension,
    BANDS,
    BAND_NAMES,
)


class EEGfMRIDataset(Dataset):
    """Dataset of paired EEG features and fMRI BOLD values."""

    def __init__(self, features: np.ndarray, bold: np.ndarray):
        self.features = torch.FloatTensor(features)
        self.bold = torch.FloatTensor(bold)

    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        return self.features[idx], self.bold[idx]


def generate_synthetic_paired_data(
    n_samples: int = 5000,
    n_channels: int = 8,
    sfreq: float = 256.0,
    n_regions: int = 96,
    epoch_duration: float = 2.0,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Generate synthetic paired EEG-fMRI data that preserves known
    neurophysiological relationships between EEG bands and brain regions.

    Known correlations used:
    - Alpha power (8-13 Hz) negatively correlates with visual cortex BOLD
    - Beta power (13-30 Hz) correlates with motor cortex BOLD
    - Theta power (4-8 Hz) correlates with hippocampal/default mode BOLD
    - Gamma power (30-45 Hz) correlates with attention/executive BOLD
    - Delta power (1-4 Hz) correlates with deep structures during rest
    - Cross-hemisphere coherence correlates with bilateral activation

    This is NOT a substitute for real data training — it's a structured
    initialization that produces physiologically plausible outputs while
    we collect/process real paired datasets.
    """
    print(f"Generating {n_samples} synthetic paired EEG-fMRI samples...")

    n_samples_per_epoch = int(sfreq * epoch_duration)
    all_features = []
    all_bold = []

    # Define which regions respond to which EEG features
    # Based on known EEG-fMRI correlations from literature
    # Region indices follow Harvard-Oxford atlas ordering

    # Visual regions (occipital)
    visual_regions = [21, 22, 23, 31, 35, 38, 39, 46, 47, 69, 70, 71, 79, 83, 86, 87, 94, 95]
    # Motor regions
    motor_regions = [6, 16, 17, 25, 54, 64, 65, 73]
    # Language regions (temporal)
    language_regions = [4, 5, 8, 9, 10, 11, 18, 19, 20, 44, 45, 52, 53, 56, 57, 58, 59, 66, 67, 68, 92, 93]
    # Default mode network
    dmn_regions = [0, 24, 27, 28, 29, 30, 33, 34, 48, 72, 75, 76, 77, 78, 81, 82]
    # Executive/attention
    exec_regions = [1, 2, 3, 26, 32, 40, 41, 42, 49, 50, 51, 74, 80, 88, 89, 90]

    for sample_idx in range(n_samples):
        # Generate different "brain states" for diversity
        state = np.random.choice(["rest", "visual", "motor", "cognitive", "mixed"])

        # Generate EEG with state-specific spectral properties
        eeg = np.zeros((n_channels, n_samples_per_epoch))
        t = np.arange(n_samples_per_epoch) / sfreq

        # Base noise
        eeg += np.random.randn(n_channels, n_samples_per_epoch) * 0.5

        # State-dependent oscillations
        alpha_power = 0.3
        beta_power = 0.3
        theta_power = 0.3
        gamma_power = 0.2
        delta_power = 0.2

        if state == "rest":
            alpha_power = 0.8 + np.random.rand() * 0.4
            theta_power = 0.5 + np.random.rand() * 0.3
            delta_power = 0.4 + np.random.rand() * 0.2
        elif state == "visual":
            alpha_power = 0.1 + np.random.rand() * 0.2  # Alpha suppression
            gamma_power = 0.6 + np.random.rand() * 0.4
        elif state == "motor":
            beta_power = 0.7 + np.random.rand() * 0.5
            alpha_power = 0.2 + np.random.rand() * 0.2
        elif state == "cognitive":
            theta_power = 0.7 + np.random.rand() * 0.4
            gamma_power = 0.5 + np.random.rand() * 0.3
            alpha_power = 0.2 + np.random.rand() * 0.2
        else:  # mixed
            alpha_power = np.random.rand() * 0.8
            beta_power = np.random.rand() * 0.8
            theta_power = np.random.rand() * 0.8
            gamma_power = np.random.rand() * 0.5

        # Add oscillations to EEG channels with spatial distribution
        for ch in range(n_channels):
            # Each channel gets slightly different weights
            ch_weight = 0.7 + 0.6 * np.random.rand()

            # Posterior channels (6,7) get more alpha
            alpha_w = alpha_power * (1.5 if ch >= 6 else 0.7) * ch_weight
            # Central channels (2,3,4,5) get more beta
            beta_w = beta_power * (1.5 if 2 <= ch <= 5 else 0.7) * ch_weight
            # Frontal channels (0,1) get more theta
            theta_w = theta_power * (1.5 if ch <= 1 else 0.7) * ch_weight

            freq_alpha = 10 + np.random.randn() * 0.5
            freq_beta = 20 + np.random.randn() * 2
            freq_theta = 6 + np.random.randn() * 0.5
            freq_gamma = 40 + np.random.randn() * 3
            freq_delta = 2 + np.random.randn() * 0.3

            eeg[ch] += alpha_w * np.sin(2 * np.pi * freq_alpha * t + np.random.rand() * 2 * np.pi)
            eeg[ch] += beta_w * np.sin(2 * np.pi * freq_beta * t + np.random.rand() * 2 * np.pi)
            eeg[ch] += theta_w * np.sin(2 * np.pi * freq_theta * t + np.random.rand() * 2 * np.pi)
            eeg[ch] += gamma_power * ch_weight * np.sin(2 * np.pi * freq_gamma * t + np.random.rand() * 2 * np.pi)
            eeg[ch] += delta_power * ch_weight * np.sin(2 * np.pi * freq_delta * t + np.random.rand() * 2 * np.pi)

        # Extract real features from this synthetic EEG
        features = extract_all_features(eeg, sfreq, target_channels=n_channels)

        # Generate corresponding BOLD values based on known correlations
        bold = np.ones(n_regions) * 0.3 + np.random.randn(n_regions) * 0.05

        # Visual cortex: negatively correlated with alpha, positively with gamma
        for r in visual_regions:
            if r < n_regions:
                bold[r] = 0.3 + (1.0 - alpha_power) * 0.4 + gamma_power * 0.2 + np.random.randn() * 0.05

        # Motor cortex: correlated with beta
        for r in motor_regions:
            if r < n_regions:
                bold[r] = 0.3 + beta_power * 0.5 + np.random.randn() * 0.05

        # Language regions: mix of theta and gamma
        for r in language_regions:
            if r < n_regions:
                bold[r] = 0.3 + theta_power * 0.2 + gamma_power * 0.3 + np.random.randn() * 0.05

        # Default mode: correlated with alpha and theta, anti-correlated with task
        for r in dmn_regions:
            if r < n_regions:
                bold[r] = 0.3 + alpha_power * 0.3 + theta_power * 0.2 - gamma_power * 0.15 + np.random.randn() * 0.05

        # Executive: correlated with gamma and theta, anti-correlated with alpha
        for r in exec_regions:
            if r < n_regions:
                bold[r] = 0.3 + gamma_power * 0.35 + theta_power * 0.15 - alpha_power * 0.1 + np.random.randn() * 0.05

        # Clip to [0, 1]
        bold = np.clip(bold, 0.05, 0.95)

        all_features.append(features)
        all_bold.append(bold)

    features_array = np.array(all_features)
    bold_array = np.array(all_bold)

    print(f"Generated features shape: {features_array.shape}")
    print(f"Generated BOLD shape: {bold_array.shape}")

    return features_array, bold_array


def load_real_data(data_dir: str):
    """
    Load real paired EEG-fMRI data from preprocessed files.

    Expected format:
        data_dir/features.npy — (n_samples, 236) EEG feature vectors
        data_dir/bold.npy     — (n_samples, 96) fMRI BOLD values

    To create these files from raw data:
    1. Download dataset from OpenNeuro (e.g., ds003505)
    2. Run preprocessing with MNE (EEG) and nilearn (fMRI)
    3. Extract EEG features using eeg_processor.extract_all_features()
    4. Extract fMRI ROI values using Harvard-Oxford atlas
    5. Save as .npy files
    """
    features_path = os.path.join(data_dir, "features.npy")
    bold_path = os.path.join(data_dir, "bold.npy")

    if os.path.exists(features_path) and os.path.exists(bold_path):
        features = np.load(features_path)
        bold = np.load(bold_path)
        print(f"Loaded real data: {features.shape[0]} samples")
        print(f"  Features: {features.shape}")
        print(f"  BOLD: {bold.shape}")
        return features, bold

    return None


def train(
    features: np.ndarray,
    bold: np.ndarray,
    epochs: int = 100,
    batch_size: int = 64,
    lr: float = 1e-3,
    val_split: float = 0.15,
    output_path: str = "model_weights.pt",
):
    """Train the EEG-to-fMRI translation model."""

    # Normalize features
    feature_mean = features.mean(axis=0)
    feature_std = features.std(axis=0)
    feature_std[feature_std < 1e-8] = 1.0
    features_norm = (features - feature_mean) / feature_std

    # Train/val split
    n = len(features_norm)
    n_val = int(n * val_split)
    indices = np.random.permutation(n)

    train_idx = indices[n_val:]
    val_idx = indices[:n_val]

    train_dataset = EEGfMRIDataset(features_norm[train_idx], bold[train_idx])
    val_dataset = EEGfMRIDataset(features_norm[val_idx], bold[val_idx])

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size)

    print(f"Training: {len(train_dataset)} samples")
    print(f"Validation: {len(val_dataset)} samples")
    print(f"Feature dim: {features.shape[1]}")
    print(f"Output dim: {bold.shape[1]}")

    # Model
    model = EEGtoFMRIModel(input_dim=features.shape[1], output_dim=bold.shape[1])
    device = torch.device("cpu")
    model.to(device)

    # Loss and optimizer
    criterion = nn.MSELoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)

    best_val_loss = float("inf")
    patience = 15
    patience_counter = 0

    for epoch in range(epochs):
        # Training
        model.train()
        train_loss = 0
        for batch_features, batch_bold in train_loader:
            batch_features = batch_features.to(device)
            batch_bold = batch_bold.to(device)

            pred = model(batch_features)
            loss = criterion(pred, batch_bold)

            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

            train_loss += loss.item()

        train_loss /= len(train_loader)

        # Validation
        model.eval()
        val_loss = 0
        with torch.no_grad():
            for batch_features, batch_bold in val_loader:
                batch_features = batch_features.to(device)
                batch_bold = batch_bold.to(device)

                pred = model(batch_features)
                loss = criterion(pred, batch_bold)
                val_loss += loss.item()

        val_loss /= len(val_loader)
        scheduler.step()

        if (epoch + 1) % 10 == 0 or epoch == 0:
            print(
                f"Epoch {epoch+1:3d}/{epochs} | "
                f"Train Loss: {train_loss:.6f} | "
                f"Val Loss: {val_loss:.6f} | "
                f"LR: {scheduler.get_last_lr()[0]:.2e}"
            )

        # Early stopping
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            patience_counter = 0
            # Save best model
            torch.save(
                {
                    "model_state_dict": model.state_dict(),
                    "feature_mean": feature_mean,
                    "feature_std": feature_std,
                    "val_loss": best_val_loss,
                    "epoch": epoch,
                    "input_dim": features.shape[1],
                    "output_dim": bold.shape[1],
                },
                output_path,
            )
        else:
            patience_counter += 1
            if patience_counter >= patience:
                print(f"Early stopping at epoch {epoch+1}")
                break

    print(f"\nBest validation loss: {best_val_loss:.6f}")
    print(f"Model saved to {output_path}")

    # Compute per-region correlation on validation set
    model.load_state_dict(
        torch.load(output_path, map_location=device, weights_only=False)["model_state_dict"]
    )
    model.eval()

    all_preds = []
    all_true = []
    with torch.no_grad():
        for batch_features, batch_bold in val_loader:
            pred = model(batch_features.to(device)).cpu().numpy()
            all_preds.append(pred)
            all_true.append(batch_bold.numpy())

    all_preds = np.vstack(all_preds)
    all_true = np.vstack(all_true)

    # Per-region Pearson correlation
    from scipy.stats import pearsonr

    correlations = []
    for r in range(bold.shape[1]):
        if np.std(all_preds[:, r]) > 0 and np.std(all_true[:, r]) > 0:
            corr, _ = pearsonr(all_preds[:, r], all_true[:, r])
            correlations.append(corr)
        else:
            correlations.append(0.0)

    mean_corr = np.mean(correlations)
    print(f"Mean per-region correlation: {mean_corr:.4f}")
    print(f"Top 5 regions by correlation:")
    top_indices = np.argsort(correlations)[-5:][::-1]
    for idx in top_indices:
        from model import HARVARD_OXFORD_REGIONS
        print(f"  {HARVARD_OXFORD_REGIONS[idx]}: r={correlations[idx]:.4f}")

    return model


def main():
    parser = argparse.ArgumentParser(description="Train EEG-to-fMRI translation model")
    parser.add_argument("--data_dir", type=str, default="./data", help="Directory with paired data")
    parser.add_argument("--epochs", type=int, default=100, help="Training epochs")
    parser.add_argument("--batch_size", type=int, default=64, help="Batch size")
    parser.add_argument("--lr", type=float, default=1e-3, help="Learning rate")
    parser.add_argument("--n_synthetic", type=int, default=5000, help="Number of synthetic samples if no real data")
    parser.add_argument("--output", type=str, default="model_weights.pt", help="Output model path")
    args = parser.parse_args()

    # Try to load real data first
    real_data = load_real_data(args.data_dir)

    if real_data is not None:
        features, bold = real_data
        print("Training on REAL paired EEG-fMRI data")
    else:
        print("No real data found — generating structured synthetic data")
        print("(To train on real data, place features.npy and bold.npy in data_dir)")
        features, bold = generate_synthetic_paired_data(n_samples=args.n_synthetic)

    train(
        features,
        bold,
        epochs=args.epochs,
        batch_size=args.batch_size,
        lr=args.lr,
        output_path=args.output,
    )


if __name__ == "__main__":
    main()