#!/usr/bin/env python
"""
Direct FC Matrix Visualization Script.

This script creates and visualizes FC matrices directly, without relying on fMRI data.
It uses the consolidated visualization utilities.
"""

import os
import numpy as np
# Configure matplotlib for headless environment
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend
import matplotlib.pyplot as plt
from visualization import vector_to_matrix, plot_fc_matrices

def create_synthetic_fc_matrices(n_subjects=10, n_rois=264, seed=42):
    """
    Create synthetic FC matrices for visualization.
    
    Args:
        n_subjects: Number of synthetic subjects
        n_rois: Number of regions of interest
        seed: Random seed for reproducibility
    
    Returns:
        dict: Dictionary with original FC matrices, latent features, and reconstructions
    """
    np.random.seed(seed)
    
    # Calculate the size of upper triangular part
    n_triu = n_rois * (n_rois - 1) // 2
    
    # Create synthetic FC matrices
    print(f"Creating {n_subjects} synthetic FC matrices with {n_rois} ROIs each")
    
    # Create original FC matrices (upper triangular vectors)
    original_fc_vectors = []
    for i in range(n_subjects):
        # Create random correlation values
        np.random.seed(seed + i)  # For reproducibility
        # Generate values between -0.8 and 0.8 (typical FC range)
        fc_triu = np.random.rand(n_triu) * 1.6 - 0.8
        original_fc_vectors.append(fc_triu)
    
    # Simulate latent features (much lower dimensional)
    latent_dim = 16
    latent_features = np.random.randn(n_subjects, latent_dim)
    
    # Simulate reconstructions with some error
    reconstructed_fc_vectors = []
    for i in range(n_subjects):
        # Add some noise to original to simulate reconstruction error
        recon = original_fc_vectors[i] + np.random.randn(n_triu) * 0.1
        # Clip to realistic correlation range
        recon = np.clip(recon, -0.99, 0.99)
        reconstructed_fc_vectors.append(recon)
    
    # Simulate a newly generated FC matrix
    generated_fc_vector = np.random.rand(n_triu) * 1.6 - 0.8
    
    return {
        'original_vectors': original_fc_vectors,
        'reconstructed_vectors': reconstructed_fc_vectors,
        'generated_vector': generated_fc_vector,
        'latent_features': latent_features
    }

def calculate_metrics(original, reconstructed):
    """
    Calculate reconstruction metrics
    
    Args:
        original: Original FC matrix
        reconstructed: Reconstructed FC matrix
    
    Returns:
        dict: Dictionary of metrics
    """
    from sklearn.metrics import mean_squared_error, r2_score
    
    # Flatten matrices
    orig_flat = original.flatten()
    recon_flat = reconstructed.flatten()
    
    # Calculate metrics
    mse = mean_squared_error(orig_flat, recon_flat)
    rmse = np.sqrt(mse)
    r2 = r2_score(orig_flat, recon_flat)
    corr = np.corrcoef(orig_flat, recon_flat)[0, 1]
    
    return {
        'MSE': mse,
        'RMSE': rmse,
        'R²': r2,
        'Correlation': corr
    }

def main():
    """Run the visualization script"""
    print("Creating direct FC matrix visualization without fMRI data")
    
    # Create synthetic FC data
    fc_data = create_synthetic_fc_matrices(n_subjects=10)
    
    # Get the vectors for visualization
    original_vector = fc_data['original_vectors'][0]
    reconstructed_vector = fc_data['reconstructed_vectors'][0]
    generated_vector = fc_data['generated_vector']
    
    # Create visualization using the consolidated function
    fig = plot_fc_matrices(original_vector, reconstructed_vector, generated_vector)
    
    # Save the figure
    output_file = "fc_visualization.png"
    fig.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"Saved visualization to {output_file}")
    
    # Save matrices for inspection
    original_matrix = vector_to_matrix(original_vector)
    reconstructed_matrix = vector_to_matrix(reconstructed_vector)
    generated_matrix = vector_to_matrix(generated_vector)
    
    np.save('original_fc.npy', original_matrix)
    np.save('reconstructed_fc.npy', reconstructed_matrix)
    np.save('generated_fc.npy', generated_matrix)
    print("Saved matrices to NPY files")
    
    # Calculate and display metrics
    metrics = calculate_metrics(original_matrix, reconstructed_matrix)
    print("\nFC Reconstruction Metrics:")
    for name, value in metrics.items():
        print(f"  {name}: {value:.6f}")
    
    print("\nVisualization complete!")

if __name__ == "__main__":
    main()