"""
Demo script to visualize FC matrices from real fMRI data using nilearn's built-in datasets.
This script uses the consolidated visualization utilities.
"""

import numpy as np
# Configure matplotlib for headless environment
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend
import matplotlib.pyplot as plt
from nilearn import datasets
from nilearn import input_data, connectome
from fc_visualization import FCVisualizer
from visualization import plot_fc_matrices

def visualize_from_nilearn_dataset():
    """Download and visualize FC matrices from nilearn's ADHD dataset."""
    print("Downloading a sample fMRI dataset (ADHD)...")
    adhd_dataset = datasets.fetch_adhd(n_subjects=1)
    
    # Get the fMRI file path
    func_file = adhd_dataset.func[0]
    confound_file = adhd_dataset.confounds[0]
    
    print(f"Downloaded fMRI file: {func_file}")
    
    # Get Power atlas coordinates
    power = datasets.fetch_coords_power_2011()
    coords = np.vstack((power.rois['x'], power.rois['y'], power.rois['z'])).T
    
    print(f"Using Power atlas with {len(coords)} ROIs")
    
    # Create a masker to extract time series from the ROIs
    masker = input_data.NiftiSpheresMasker(
        coords, 
        radius=8,  # 8mm radius
        standardize=True,
        memory='nilearn_cache', 
        memory_level=1,
        verbose=1,
        detrend=True,
        low_pass=0.08,
        high_pass=0.01,
        t_r=2.0  # ADHD dataset TR
    )
    
    # Extract time series, including confounds
    print("Extracting time series from ROIs...")
    time_series = masker.fit_transform(func_file, confounds=confound_file)
    print(f"Time series shape: {time_series.shape}")
    
    # Compute correlation matrix (FC matrix)
    correlation_measure = connectome.ConnectivityMeasure(
        kind='correlation',
        vectorize=False,
        discard_diagonal=False
    )
    
    fc_matrix = correlation_measure.fit_transform([time_series])[0]
    print(f"FC matrix shape: {fc_matrix.shape}")
    
    # Save the FC matrix for future use
    np.save('adhd_fc_matrix.npy', fc_matrix)
    print("Saved FC matrix to adhd_fc_matrix.npy")
    
    # Visualize using both methods for comparison
    
    # Method 1: Using FCVisualizer (OOP approach)
    visualizer = FCVisualizer(cmap='RdBu_r', vmin=-1, vmax=1)
    fig1, _ = visualizer.plot_single_matrix(fc_matrix, title="ADHD FC Matrix (FCVisualizer)")
    fig1.savefig('adhd_fc_matrix_method1.png', dpi=300, bbox_inches='tight')
    
    # Method 2: Using plot_fc_matrices (functional approach)
    # Create a zero matrix for the "reconstructed" and "generated" placeholders
    zero_matrix = np.zeros_like(fc_matrix)
    fig2 = plot_fc_matrices(fc_matrix, zero_matrix, zero_matrix)
    fig2.savefig('adhd_fc_matrix_method2.png', dpi=300, bbox_inches='tight')
    
    print("Saved visualizations using both methods")

if __name__ == "__main__":
    visualize_from_nilearn_dataset()