vis.py

#!/usr/bin/env python3
"""
Comprehensive Visualization System for Diffusion Models
=======================================================

This module provides a unified visualization system for DDPM, SM, and NCSM models.
Includes training visualizations, sampling animations, and score comparisons.
"""

import os
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
from typing import Dict, List, Any, Optional, Tuple
import warnings
warnings.filterwarnings('ignore')

# Set matplotlib backend
plt.switch_backend('Agg')

def _get_noise_levels_from_config(config: Dict[str, Any]) -> List[float]:
    """
    Helper function to generate noise levels from config parameters.
    Uses the same logic as training to ensure consistency.
    """
    from models.ncsm import create_noise_schedule
    return create_noise_schedule(
        sigma_min=config['sigma_min'],
        sigma_max=config['sigma_max'],
        num_scales=config['num_noise_scales']
    )

# Global plot parameters for consistent styling
PLOT_BOUNDS = (-0.5, 0.5, -0.5, 0.5)  # Match actual dataset bounds
POINT_SIZE = 20  # Bigger, more visible points
FIGURE_SIZE = (6, 6)  # Smaller figures make text and points appear bigger
DPI = 150

# Unified color scheme
COLORS = {
    'training': 'blue',        # Training data always blue
    'ddpm': 'red',            # DDPM model red
    'sm': 'green',            # Score Matching model green  
    'ncsm': 'orange',         # NCSM model orange
    'loss_raw': 'blue',       # Raw loss curve
    'loss_smooth': 'red',     # Smoothed loss curve
    'contour': 'lightgray',   # PDF contour lines
}

# Unified font sizes
FONT_SIZES = {
    'title': 16,              # Plot titles
    'axis_label': 14,         # X/Y axis labels
    'legend': 12,             # Legend text
    'annotation': 12,         # Animation annotations
    'tick_label': 12,         # Axis tick labels
}





def plot_losses(losses: List[float], model_type: str, config: Dict[str, Any], show: bool = True):
    """
    Plot training losses over time.
    
    Args:
        losses: List of loss values
        model_type: Type of model ('ddpm', 'sm', 'ncsm')
        config: Configuration dictionary
        show: Whether to show the plot
    """
    plt.figure(figsize=(FIGURE_SIZE[0]*2, FIGURE_SIZE[1]))
    
    # Convert to numpy array for easier manipulation
    losses_array = np.array(losses)
    
    # Plot raw losses
    plt.plot(losses_array, alpha=0.6, color=COLORS['loss_raw'], linewidth=1)
    
    # Add smoothed version if we have enough data points
    if len(losses) > 50:
        window_size = min(50, len(losses) // 10)
        smoothed = np.convolve(losses_array, np.ones(window_size)/window_size, mode='valid')
        plt.plot(range(window_size-1, len(losses)), smoothed, 
                color=COLORS['loss_smooth'], linewidth=2, label='Smoothed')
        plt.legend(fontsize=FONT_SIZES['legend'])
    
    plt.title(f'{model_type.upper()} Training Loss', fontsize=FONT_SIZES['title'])
    plt.xlabel('Epoch' if model_type == 'ddpm' else 'Step', fontsize=FONT_SIZES['axis_label'])
    plt.ylabel('Loss', fontsize=FONT_SIZES['axis_label'])
    plt.grid(True, alpha=0.3)
    plt.tick_params(axis='both', which='major', labelsize=FONT_SIZES['tick_label'])
    
    # Save to figures directory
    os.makedirs('./figures', exist_ok=True)
    save_path = os.path.join('./figures', f'{model_type}_losses.png')
    plt.savefig(save_path, dpi=DPI, bbox_inches='tight')
    print(f"Saved loss plot: {save_path}")
    
    if show:
        plt.show()
    else:
        plt.close()


def plot_samples(samples: torch.Tensor,
                title: str = "Generated Samples",
                save_path: Optional[str] = None,
                show_plot: bool = True,
                figsize: Tuple[int, int] = (5, 5)) -> None:
    """
    Plot generated samples.
    
    Args:
        samples: Generated samples to plot
        title: Title for the plot
        save_path: Path to save the plot
        show_plot: Whether to show the plot
        figsize: Figure size
    """
    fig, ax = plt.subplots(figsize=figsize)
    
    # Convert to numpy for plotting
    if isinstance(samples, torch.Tensor):
        samples_np = samples.detach().cpu().numpy()
    else:
        samples_np = samples
    
    # Plot samples
    ax.scatter(samples_np[:, 0], samples_np[:, 1], 
               s=POINT_SIZE, alpha=0.7, color='blue', label='Generated')
    
    # Set bounds and styling
    ax.set_xlim(PLOT_BOUNDS[0], PLOT_BOUNDS[1])
    ax.set_ylim(PLOT_BOUNDS[2], PLOT_BOUNDS[3])
    ax.set_xlabel('X', fontsize=FONT_SIZES['axis_label'])
    ax.set_ylabel('Y', fontsize=FONT_SIZES['axis_label'])
    ax.set_title(title, fontsize=FONT_SIZES['title'])
    ax.grid(True, alpha=0.3)
    ax.legend(fontsize=FONT_SIZES['legend'], loc='upper right')
    
    # Make tick labels larger
    ax.tick_params(axis='both', which='major', labelsize=FONT_SIZES['tick_label'])
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
    if show_plot:
        plt.show()
    else:
        plt.close()


def plot_data_comparison(training_data: torch.Tensor, 
                        generated_data: torch.Tensor,
                        model_type: str,
                        save_path: Optional[str] = None,
                        show_plot: bool = True,
                        dataset=None,
                        figsize: Tuple[int, int] = (10, 5)) -> None:
    """
    Plot comparison between training data and generated samples with true PDF contours.
    
    Args:
        training_data: Training data samples (with fixed seed)
        generated_data: Generated samples from model
        model_type: Type of model ('ddpm', 'sm', 'ncsm')
        save_path: Path to save the plot
        show_plot: Whether to show the plot
        dataset: Dataset instance (for PDF contours)
        figsize: Figure size (smaller for larger fonts)
    """
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)
    
    # Convert to numpy for plotting
    if isinstance(training_data, torch.Tensor):
        training_np = training_data.detach().cpu().numpy()
    else:
        training_np = training_data
        
    if isinstance(generated_data, torch.Tensor):
        generated_np = generated_data.detach().cpu().numpy()
    else:
        generated_np = generated_data
    
    # Create PDF contour background if dataset is provided
    if dataset is not None:
        x_min, x_max, y_min, y_max = PLOT_BOUNDS
        x_grid = np.linspace(x_min, x_max, 100)
        y_grid = np.linspace(y_min, y_max, 100)
        X, Y = np.meshgrid(x_grid, y_grid)
        points = np.stack([X.ravel(), Y.ravel()], axis=1)
        
        # Compute PDF values
        pdf_values = dataset.pdf(X.ravel(), Y.ravel())
        pdf_grid = pdf_values.reshape(X.shape)
        
        # Add contour lines to both plots
        levels = np.linspace(0.1, pdf_values.max(), 8)
        ax1.contour(X, Y, pdf_grid, levels=levels, colors='lightgray', alpha=0.6, linewidths=0.8)
        ax2.contour(X, Y, pdf_grid, levels=levels, colors='lightgray', alpha=0.6, linewidths=0.8)
    
    # Plot training data
    ax1.scatter(training_np[:, 0], training_np[:, 1], 
               s=POINT_SIZE, alpha=0.7, color=COLORS['training'], label='Training')
    ax1.set_xlim(PLOT_BOUNDS[0], PLOT_BOUNDS[1])
    ax1.set_ylim(PLOT_BOUNDS[2], PLOT_BOUNDS[3])
    ax1.set_xlabel('X', fontsize=FONT_SIZES['axis_label'])
    ax1.set_ylabel('Y', fontsize=FONT_SIZES['axis_label'])
    ax1.set_title('Training Data', fontsize=FONT_SIZES['title'])
    ax1.grid(True, alpha=0.3)
    ax1.legend(fontsize=FONT_SIZES['legend'], loc='upper right')
    ax1.tick_params(axis='both', which='major', labelsize=FONT_SIZES['tick_label'])
    
    # Plot generated samples with model-specific color
    model_color = COLORS.get(model_type, 'blue')  # Default to blue if model not found
    ax2.scatter(generated_np[:, 0], generated_np[:, 1], 
               s=POINT_SIZE, alpha=0.7, color=model_color, label='Generated')
    ax2.set_xlim(PLOT_BOUNDS[0], PLOT_BOUNDS[1])
    ax2.set_ylim(PLOT_BOUNDS[2], PLOT_BOUNDS[3])
    ax2.set_xlabel('X', fontsize=FONT_SIZES['axis_label'])
    ax2.set_ylabel('Y', fontsize=FONT_SIZES['axis_label'])
    ax2.set_title(f'{model_type.upper()} Generated', fontsize=FONT_SIZES['title'])
    ax2.grid(True, alpha=0.3)
    ax2.legend(fontsize=FONT_SIZES['legend'], loc='upper right')
    ax2.tick_params(axis='both', which='major', labelsize=FONT_SIZES['tick_label'])
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
    if show_plot:
        plt.show()
    else:
        plt.close()


def save_samples_grid(samples: torch.Tensor,
                     model_type: str,
                     config: Dict[str, Any],
                     save_name: Optional[str] = None,
                     figsize: Tuple[int, int] = (8, 8)) -> str:
    """
    Save a grid of generated samples.
    
    Args:
        samples: Generated samples
        model_type: Type of model ('ddpm', 'sm', 'ncsm')
        config: Configuration dictionary
        save_name: Optional custom save name
        figsize: Figure size
        
    Returns:
        Path to saved figure
    """
    plt.figure(figsize=figsize)
    
    # Convert to numpy
    if isinstance(samples, torch.Tensor):
        samples_np = samples.detach().cpu().numpy()
    else:
        samples_np = samples
    
    # Create scatter plot with uniform styling
    model_color = COLORS.get(model_type, 'blue')  # Default to blue if model not found
    plt.scatter(samples_np[:, 0], samples_np[:, 1], 
               s=POINT_SIZE, alpha=0.6, c=model_color, edgecolors='none')
    
    plt.xlim(PLOT_BOUNDS[0], PLOT_BOUNDS[1])
    plt.ylim(PLOT_BOUNDS[2], PLOT_BOUNDS[3])
    plt.gca().set_aspect('equal')
    plt.grid(True, alpha=0.3)
    plt.title(f'{model_type.upper()} Generated Samples', fontsize=FONT_SIZES['title'])
    plt.xlabel('X', fontsize=FONT_SIZES['axis_label'])
    plt.ylabel('Y', fontsize=FONT_SIZES['axis_label'])
    plt.tick_params(axis='both', which='major', labelsize=FONT_SIZES['tick_label'])
    
    # Save to figures directory
    os.makedirs('./figures', exist_ok=True)
    if save_name is None:
        save_name = f'{model_type}_samples.png'
    save_path = os.path.join('./figures', save_name)
    plt.savefig(save_path, dpi=DPI, bbox_inches='tight')
    print(f"Saved samples: {save_path}")
    
    plt.close()
    return save_path


def plot_dataset_overview(dataset, n_samples: int = 2000, n_grid: int = 200, 
                         save_path: Optional[str] = None, show: bool = True,
                         figsize: Tuple[int, int] = (12, 5)) -> str:
    """
    Create a double plot showing training data samples and PDF contours.
    
    Args:
        dataset: SmilingFaceDataset instance
        n_samples: Number of samples to generate and plot
        n_grid: Grid resolution for PDF contour plot
        save_path: Optional path to save the plot
        show: Whether to display the plot
        figsize: Figure size tuple
        
    Returns:
        Path where the plot was saved
    """
    # Generate samples
    samples = dataset.sample(n_samples, return_tensor=False)
    
    # Create evaluation grid for PDF
    x_min, x_max, y_min, y_max = dataset.get_bounds()
    xs = np.linspace(x_min, x_max, n_grid)
    ys = np.linspace(y_min, y_max, n_grid)
    X, Y = np.meshgrid(xs, ys)
    
    # Evaluate PDF
    PDF_vals = dataset.pdf(X, Y)
    
    # Create figure with subplots
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)
    fig.suptitle('Smiling Face Dataset Overview', fontsize=16, fontweight='bold')
    
    # Left plot: Training data samples
    ax1.scatter(samples[:, 0], samples[:, 1], alpha=0.6, s=POINT_SIZE, c='blue')
    ax1.set_xlim(x_min, x_max)
    ax1.set_ylim(y_min, y_max)
    ax1.set_xlabel('X coordinate')
    ax1.set_ylabel('Y coordinate')
    ax1.set_title(f'Training Data Samples (n={n_samples})')
    ax1.grid(True, alpha=0.3)
    ax1.set_aspect('equal', adjustable='box')
    
    # Right plot: PDF contours
    contour = ax2.contourf(X, Y, PDF_vals, levels=50, cmap='viridis')
    ax2.set_xlim(x_min, x_max)
    ax2.set_ylim(y_min, y_max)
    ax2.set_xlabel('X coordinate')
    ax2.set_ylabel('Y coordinate')
    ax2.set_title('Probability Density Function')
    ax2.grid(True, alpha=0.3)
    ax2.set_aspect('equal', adjustable='box')
    
    # Add colorbar for PDF
    cbar = plt.colorbar(contour, ax=ax2)
    cbar.set_label('PDF Value')
    
    plt.tight_layout()
    
    # Save the plot
    if save_path is None:
        os.makedirs("figures", exist_ok=True)
        save_path = "figures/dataset_overview.png"
    
    plt.savefig(save_path, dpi=DPI, bbox_inches='tight', facecolor='white')
    print(f"📊 Dataset overview saved to: {save_path}")
    
    if show:
        plt.show()
    else:
        plt.close()
    
    return save_path


def plot_score_comparison(model,
                         model_type: str,
                         training_data: torch.Tensor,
                         generated_data: torch.Tensor,
                         save_path: Optional[str] = None,
                         show_plot: bool = True,
                         noise_levels: Optional[List[float]] = None,
                         n_grid: int = 15,
                         dataset=None) -> None:
    """
    Plot comparison between predicted and true score fields side by side.
    
    Args:
        model: Trained model
        model_type: Type of model ('sm', 'ncsm')
        training_data: Training data samples
        generated_data: Generated data samples
        save_path: Path to save the plot
        show_plot: Whether to show the plot
        noise_levels: Noise levels for NCSM (optional)
        n_grid: Grid size for score visualization
        dataset: Dataset object (for true score function)
    """
    # Create side-by-side plots
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(FIGURE_SIZE[0]*2, FIGURE_SIZE[1]))
    
    # Create grid for score evaluation
    x = np.linspace(PLOT_BOUNDS[0], PLOT_BOUNDS[1], n_grid)
    y = np.linspace(PLOT_BOUNDS[2], PLOT_BOUNDS[3], n_grid)
    X, Y = np.meshgrid(x, y)
    
    # Prepare grid points
    grid_points = np.column_stack([X.ravel(), Y.ravel()])
    grid_tensor = torch.FloatTensor(grid_points).to(next(model.parameters()).device)
    
    # Add PDF contour lines if dataset is available
    if dataset is not None:
        # Create finer grid for PDF contours
        x_pdf = np.linspace(PLOT_BOUNDS[0], PLOT_BOUNDS[1], 100)
        y_pdf = np.linspace(PLOT_BOUNDS[2], PLOT_BOUNDS[3], 100)
        X_pdf, Y_pdf = np.meshgrid(x_pdf, y_pdf)
        
        # Compute PDF values
        pdf_values = dataset.pdf(X_pdf.ravel(), Y_pdf.ravel())
        pdf_grid = pdf_values.reshape(X_pdf.shape)
        
        # Add contour lines to both plots
        levels = np.linspace(0.1, pdf_values.max(), 8)
        ax1.contour(X_pdf, Y_pdf, pdf_grid, levels=levels, colors=COLORS['contour'], 
                   alpha=0.6, linewidths=0.8)
        ax2.contour(X_pdf, Y_pdf, pdf_grid, levels=levels, colors=COLORS['contour'], 
                   alpha=0.6, linewidths=0.8)
    
    # Get true scores (if dataset available)
    if dataset is not None:
        true_scores = dataset.true_score(grid_tensor[:, 0], grid_tensor[:, 1])
        if isinstance(true_scores, torch.Tensor):
            true_scores = true_scores.detach().cpu().numpy()
        
        # Plot true score field
        U_true = true_scores[:, 0].reshape(n_grid, n_grid)
        V_true = true_scores[:, 1].reshape(n_grid, n_grid)
        
        # Calculate magnitude and scale separately for visibility
        M_true = np.sqrt(U_true**2 + V_true**2)
        M_true_max = np.max(M_true)
        
        if M_true_max > 0:
            scale_factor = 1.0 / M_true_max
            U_true_scaled = U_true * scale_factor
            V_true_scaled = V_true * scale_factor
        else:
            U_true_scaled = U_true
            V_true_scaled = V_true
        
        ax1.quiver(X, Y, U_true_scaled, V_true_scaled, 
                  scale=10.0, scale_units='xy', angles='xy', alpha=0.7, color=COLORS['training'])
        ax1.set_title('True Score Field', fontsize=FONT_SIZES['title'])
    else:
        ax1.text(0.5, 0.5, 'Dataset not available\nfor true score', 
                transform=ax1.transAxes, ha='center', va='center', fontsize=FONT_SIZES['annotation'])
        ax1.set_title('True Score Field (N/A)', fontsize=FONT_SIZES['title'])
    
    # Get predicted scores
    model.eval()
    with torch.no_grad():
        if model_type == 'sm':
            predicted_scores = model(grid_tensor).detach().cpu().numpy()
        elif model_type == 'ncsm':
            # Use the first noise level for NCSM
            noise_level = noise_levels[0] if noise_levels else 0.1
            noise_tensor = torch.full((grid_tensor.shape[0], 1), noise_level, 
                                    device=grid_tensor.device, dtype=grid_tensor.dtype)
            predicted_scores = model(grid_tensor, noise_tensor).detach().cpu().numpy()
        else:
            predicted_scores = np.zeros_like(grid_points)
    
    # Plot predicted score field
    U_pred = predicted_scores[:, 0].reshape(n_grid, n_grid)
    V_pred = predicted_scores[:, 1].reshape(n_grid, n_grid)
    
    # Calculate magnitude and scale separately for visibility
    M_pred = np.sqrt(U_pred**2 + V_pred**2)
    M_pred_max = np.max(M_pred)
    
    if M_pred_max > 0:
        scale_factor = 1.0 / M_pred_max
        U_pred_scaled = U_pred * scale_factor
        V_pred_scaled = V_pred * scale_factor
    else:
        U_pred_scaled = U_pred
        V_pred_scaled = V_pred
    
    # Use model-specific color for predicted scores
    model_color = COLORS.get(model_type, 'blue')
    ax2.quiver(X, Y, U_pred_scaled, V_pred_scaled, 
              scale=10.0, scale_units='xy', angles='xy', alpha=0.7, color=model_color)
    
    title = f'Predicted Score Field ({model_type.upper()})'
    if model_type == 'ncsm' and noise_levels:
        title += f' (σ={noise_levels[0]:.3f})'
    ax2.set_title(title, fontsize=FONT_SIZES['title'])
    
    # Set bounds and styling for both plots
    for ax in [ax1, ax2]:
        ax.set_xlim(PLOT_BOUNDS[0], PLOT_BOUNDS[1])
        ax.set_ylim(PLOT_BOUNDS[2], PLOT_BOUNDS[3])
        ax.set_xlabel('X', fontsize=FONT_SIZES['axis_label'])
        ax.set_ylabel('Y', fontsize=FONT_SIZES['axis_label'])
        ax.grid(True, alpha=0.3)
        ax.tick_params(axis='both', which='major', labelsize=FONT_SIZES['tick_label'])
        ax.set_aspect('equal')
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Saved score comparison: {save_path}")
    
    if show_plot:
        plt.show()
    else:
        plt.close()


def plot_score_distance(model_type: str,
                       model,
                       dataset,
                       config: Dict[str, Any],
                       device: torch.device,
                       n_grid: int = 50,
                       noise_level: float = 0.1,
                       save_name: Optional[str] = None,
                       show: bool = True) -> str:
    """
    Plot distance between predicted and true scores.
    
    Args:
        model_type: Type of model ('sm', 'ncsm')
        model: Trained model
        dataset: Dataset
        config: Configuration dictionary
        device: Device to run on
        n_grid: Grid resolution
        noise_level: Noise level (for NCSM)
        save_name: Optional custom save name
        show: Whether to show the plot
        
    Returns:
        Path to saved figure
    """
    # Create grid
    x = np.linspace(PLOT_BOUNDS[0], PLOT_BOUNDS[1], n_grid)
    y = np.linspace(PLOT_BOUNDS[2], PLOT_BOUNDS[3], n_grid)
    X, Y = np.meshgrid(x, y)
    
    # Prepare grid points
    grid_points = np.column_stack([X.ravel(), Y.ravel()])
    grid_tensor = torch.FloatTensor(grid_points).to(device)
    
    # Get true scores
    true_scores = dataset.true_score(grid_tensor[:, 0], grid_tensor[:, 1])
    if isinstance(true_scores, torch.Tensor):
        true_scores = true_scores.detach().cpu().numpy()
    
    # Get predicted scores
    model.eval()
    with torch.no_grad():
        if model_type == 'sm':
            predicted_scores = model(grid_tensor).detach().cpu().numpy()
        elif model_type == 'ncsm':
            # Add noise level as separate input
            noise_tensor = torch.full((grid_tensor.shape[0], 1), noise_level, 
                                    device=device, dtype=grid_tensor.dtype)
            predicted_scores = model(grid_tensor, noise_tensor).detach().cpu().numpy()
    
    # Compute distances
    distances = np.linalg.norm(predicted_scores - true_scores, axis=1)
    distances = distances.reshape(n_grid, n_grid)
    
    # Create heatmap
    plt.figure(figsize=FIGURE_SIZE)
    plt.imshow(distances, extent=[PLOT_BOUNDS[0], PLOT_BOUNDS[1], PLOT_BOUNDS[2], PLOT_BOUNDS[3]], 
              origin='lower', cmap='viridis', aspect='equal')
    plt.colorbar(label='L2 Distance')
    
    title = f'{model_type.upper()} Score Distance'
    if model_type == 'ncsm':
        title += f' (σ={noise_level:.3f})'
    plt.title(title)
    plt.xlabel('x')
    plt.ylabel('y')
    
    # Save to figures directory
    os.makedirs('./figures', exist_ok=True)
    if save_name is None:
        if model_type == 'ncsm':
            save_name = f'ncsm_score_distance_sigma_{noise_level:.3f}.png'
        else:
            save_name = f'{model_type}_score_distance.png'
    save_path = os.path.join('./figures', save_name)
    plt.savefig(save_path, dpi=DPI, bbox_inches='tight')
    print(f"Saved score distance: {save_path}")
    
    if show:
        plt.show()
    else:
        plt.close()
    
    return save_path


def plot_noise_schedules(config: Dict[str, Any],
                       save_path: Optional[str] = None,
                       show: bool = True,
                       figsize: Tuple[int, int] = (8, 6)) -> str:
    """
    Plot linear and cosine noise schedules for DDPM.
    
    Args:
        config: Configuration dictionary containing DDPM parameters
        save_path: Optional path to save the plot
        show: Whether to display the plot
        figsize: Figure size tuple
        
    Returns:
        Path where the plot was saved
    """
    from models.ddpm import NoiseScheduler
    
    # Get parameters from config or use defaults
    num_timesteps = config.get('num_timesteps', 1000)
    beta_start = config.get('beta_start', 0.0001)
    beta_end = config.get('beta_end', 0.02)
    
    # Create both schedulers
    linear_scheduler = NoiseScheduler(
        num_timesteps=num_timesteps,
        beta_start=beta_start,
        beta_end=beta_end,
        schedule='linear'
    )
    
    cosine_scheduler = NoiseScheduler(
        num_timesteps=num_timesteps,
        beta_start=beta_start,
        beta_end=beta_end,
        schedule='cosine'
    )
    
    # Create timesteps array
    timesteps = torch.arange(num_timesteps)
    
    # Create figure with single plot
    fig, ax = plt.subplots(1, 1, figsize=figsize)
    fig.suptitle('DDPM Noise Schedules Comparison', fontsize=16, fontweight='bold')
    
    # Plot cumulative alpha values
    ax.plot(timesteps.numpy(), linear_scheduler.alphas_cumprod.numpy(), 
            label='Linear', color='blue', linewidth=2)
    ax.plot(timesteps.numpy(), cosine_scheduler.alphas_cumprod.numpy(), 
            label='Cosine', color='red', linewidth=2)
    ax.set_xlabel('Timestep')
    ax.set_ylabel('ᾱ(t) = ∏α(s)')
    ax.set_title('Cumulative Alpha Schedule')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    
    # Save the plot
    if save_path is None:
        os.makedirs("figures", exist_ok=True)
        save_path = f"figures/noise_schedules_comparison.png"
    
    plt.savefig(save_path, dpi=DPI, bbox_inches='tight', facecolor='white')
    print(f"📊 Noise schedules comparison saved to: {save_path}")
    
    if show:
        plt.show()
    else:
        plt.close()
    
    return save_path


def create_ddpm_sampling_animation(model, 
                                 num_samples: int,
                                 config: Dict[str, Any],
                                 device: torch.device,
                                 save_name: Optional[str] = None,
                                 interval: int = 100,
                                 subsample_steps: int = 10,
                                 final_pause_frames: int = 20) -> str:
    """
    Create GIF animation of DDPM sampling process.
    
    Args:
        model: Trained DDPM model
        num_samples: Number of samples to generate
        config: Configuration dictionary
        device: Device to run on
        save_name: Optional custom save name
        interval: Animation frame interval in ms
        subsample_steps: How many steps to skip between frames
        final_pause_frames: Number of extra frames to show final result
        
    Returns:
        Path to saved animation
    """
    from sampling import DDPMSampler
    
    # Create sampler
    sampler = DDPMSampler(model, device)
    
    # Generate samples with trajectory tracking
    samples, info = sampler.sample(num_samples, store_trajectory=True, 
                                  store_every=subsample_steps, verbose=False)
    
    trajectory = info['trajectory']
    timesteps = info['timesteps']
    
    # Create animation
    fig, ax = plt.subplots(figsize=FIGURE_SIZE)
    ax.set_xlim(PLOT_BOUNDS[0], PLOT_BOUNDS[1])
    ax.set_ylim(PLOT_BOUNDS[2], PLOT_BOUNDS[3])
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.3)
    ax.set_title('DDPM Sampling Process', fontsize=FONT_SIZES['title'])
    
    # Initialize scatter plot with bigger points
    scat = ax.scatter([], [], s=POINT_SIZE, alpha=0.6, c=COLORS['ddpm'], edgecolors='none')
    
    # Add timestep text
    time_text = ax.text(0.02, 0.98, '', transform=ax.transAxes, 
                       fontsize=FONT_SIZES['annotation'], verticalalignment='top',
                       bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
    
    def animate(frame):
        # Show normal trajectory frames
        if frame < len(trajectory):
            data = trajectory[frame].detach().cpu().numpy()
            scat.set_offsets(data)
            
            # Update timestep text
            if frame < len(timesteps):
                t = timesteps[frame]
                time_text.set_text(f'Timestep: {t}')
        else:
            # Show final frame for pause frames
            data = trajectory[-1].detach().cpu().numpy()
            scat.set_offsets(data)
            time_text.set_text('Final Result')
            
        return [scat, time_text]
    
    # Create animation with extended frames (original + pause frames)
    total_frames = len(trajectory) + final_pause_frames
    anim = FuncAnimation(fig, animate, frames=total_frames, 
                        interval=interval, blit=True, repeat=True)
    
    # Save animation to figures directory
    os.makedirs('./figures', exist_ok=True)
    if save_name is None:
        save_name = f'ddpm_sampling_animation.gif'
    save_path = os.path.join('./figures', save_name)
    
    # Save with PillowWriter
    writer = PillowWriter(fps=10)
    anim.save(save_path, writer=writer)
    plt.close()
    
    print(f"Saved DDPM sampling animation: {save_path}")
    return save_path


def create_langevin_sampling_animation(score_network,
                                     dataset, 
                                     config: Dict[str, Any],
                                     device: torch.device,
                                     num_samples: int = 100,
                                     num_steps: int = None,
                                     step_size: float = None,
                                     save_name: Optional[str] = None,
                                     interval: int = 50,
                                     subsample_steps: int = 10,
                                     final_pause_frames: int = 20) -> str:
    """
    Create GIF animation of Langevin sampling process for Score Matching.
    
    Args:
        score_network: Trained score network
        dataset: Dataset
        config: Configuration dictionary
        device: Device to run on
        num_samples: Number of samples to generate
        num_steps: Number of Langevin steps
        step_size: Langevin step size
        save_name: Optional custom save name
        interval: Animation frame interval in ms
        subsample_steps: How many steps to skip between frames
        final_pause_frames: Number of extra frames to show final result
        
    Returns:
        Path to saved animation
    """
    from sampling import create_learned_sampler
    
    # Use config values if not provided
    if num_steps is None:
        num_steps = config['langevin_num_steps']
    if step_size is None:
        step_size = config['langevin_step_size']
    
    # Create sampler
    bounds = dataset.get_bounds()
    sampler = create_learned_sampler(score_network, bounds, device, dataset)
    
    # Generate samples with trajectory tracking
    samples, info = sampler.sample(num_samples, num_steps, step_size, 
                                  store_trajectory=True, store_every=subsample_steps, 
                                  verbose=False)
    
    trajectory = info['trajectory']
    
    # Create animation
    fig, ax = plt.subplots(figsize=FIGURE_SIZE)
    ax.set_xlim(PLOT_BOUNDS[0], PLOT_BOUNDS[1])
    ax.set_ylim(PLOT_BOUNDS[2], PLOT_BOUNDS[3])
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.3)
    ax.set_title('Score Matching Langevin Sampling', fontsize=FONT_SIZES['title'])
    
    # Initialize scatter plot with bigger points
    scat = ax.scatter([], [], s=POINT_SIZE, alpha=0.6, c=COLORS['sm'], edgecolors='none')
    
    # Add step text
    step_text = ax.text(0.02, 0.98, '', transform=ax.transAxes, 
                       fontsize=FONT_SIZES['annotation'], verticalalignment='top',
                       bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
    
    def animate(frame):
        # Show normal trajectory frames
        if frame < len(trajectory):
            data = trajectory[frame].detach().cpu().numpy()
            scat.set_offsets(data)
            
            # Update step text
            step = frame * subsample_steps
            step_text.set_text(f'Step: {step}')
        else:
            # Show final frame for pause frames
            data = trajectory[-1].detach().cpu().numpy()
            scat.set_offsets(data)
            step_text.set_text('Final Result')
            
        return [scat, step_text]
    
    # Create animation with extended frames (original + pause frames)
    total_frames = len(trajectory) + final_pause_frames
    anim = FuncAnimation(fig, animate, frames=total_frames, 
                        interval=interval, blit=True, repeat=True)
    
    # Save animation to figures directory
    os.makedirs('./figures', exist_ok=True)
    if save_name is None:
        save_name = f'sm_sampling_animation.gif'
    save_path = os.path.join('./figures', save_name)
    
    # Save with PillowWriter
    writer = PillowWriter(fps=20)
    anim.save(save_path, writer=writer)
    plt.close()
    
    print(f"Saved SM sampling animation: {save_path}")
    return save_path


def create_annealed_langevin_animation(score_network,
                                     noise_levels: List[float],
                                     dataset,
                                     config: Dict[str, Any],
                                     device: torch.device,
                                     num_samples: int = 100,
                                     steps_per_noise: int = None,
                                     base_step_size: float = None,
                                     save_name: Optional[str] = None,
                                     interval: int = 50,
                                     subsample_steps: int = 5,
                                     final_pause_frames: int = 20) -> str:
    """
    Create GIF animation of annealed Langevin sampling process for NCSM.
    
    Args:
        score_network: Trained score network
        noise_levels: List of noise levels (high to low)
        dataset: Dataset
        config: Configuration dictionary
        device: Device to run on
        num_samples: Number of samples to generate
        steps_per_noise: Number of steps per noise level
        save_name: Optional custom save name
        interval: Animation frame interval in ms
        subsample_steps: How many steps to skip between frames
        final_pause_frames: Number of extra frames to show final result
        
    Returns:
        Path to saved animation
    """
    from sampling import create_annealed_sampler_from_ncsn
    
    # Use config values if not provided
    if steps_per_noise is None:
        steps_per_noise = config['langevin_steps_per_noise']
    if base_step_size is None:
        base_step_size = config['langevin_base_step_size']
    
    # Create sampler
    sampler = create_annealed_sampler_from_ncsn(score_network, noise_levels, 
                                              dataset, device)
    
    # Generate samples with trajectory tracking
    samples, info = sampler.sample(num_samples, steps_per_noise, 
                                  base_step_size=base_step_size,
                                  store_trajectory=True, store_every=subsample_steps, 
                                  verbose=False)
    
    trajectory = info['trajectory']
    noise_levels_used = info.get('noise_levels', noise_levels)
    
    # Create animation
    fig, ax = plt.subplots(figsize=FIGURE_SIZE)
    ax.set_xlim(PLOT_BOUNDS[0], PLOT_BOUNDS[1])
    ax.set_ylim(PLOT_BOUNDS[2], PLOT_BOUNDS[3])
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.3)
    ax.set_title('NCSM Annealed Langevin Sampling', fontsize=FONT_SIZES['title'])
    
    # Initialize scatter plot with bigger points
    scat = ax.scatter([], [], s=POINT_SIZE, alpha=0.6, c=COLORS['ncsm'], edgecolors='none')
    
    # Add noise level text
    noise_text = ax.text(0.02, 0.98, '', transform=ax.transAxes, 
                        fontsize=FONT_SIZES['annotation'], verticalalignment='top',
                        bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
    
    def animate(frame):
        # Show normal trajectory frames
        if frame < len(trajectory):
            data = trajectory[frame].detach().cpu().numpy()
            scat.set_offsets(data)
            
            # Update noise level text - USE ACTUAL PARAMETERS
            frames_per_noise = steps_per_noise // subsample_steps
            noise_idx = min(frame // frames_per_noise, 
                          len(noise_levels_used) - 1)
            if noise_idx < len(noise_levels_used):
                noise_level = noise_levels_used[noise_idx]
                noise_text.set_text(f'Noise Level: {noise_level:.3f}')
        else:
            # Show final frame for pause frames
            data = trajectory[-1].detach().cpu().numpy()
            scat.set_offsets(data)
            noise_text.set_text('Final Result')
            
        return [scat, noise_text]
    
    # Create animation with extended frames (original + pause frames)
    total_frames = len(trajectory) + final_pause_frames
    anim = FuncAnimation(fig, animate, frames=total_frames, 
                        interval=interval, blit=True, repeat=True)
    
    # Save animation to figures directory
    os.makedirs('./figures', exist_ok=True)
    if save_name is None:
        save_name = f'ncsm_sampling_animation.gif'
    save_path = os.path.join('./figures', save_name)
    
    # Save with PillowWriter
    writer = PillowWriter(fps=20)
    anim.save(save_path, writer=writer)
    plt.close()
    
    print(f"Saved NCSM sampling animation: {save_path}")
    return save_path


def create_sampling_animation(model_type: str,
                            model,
                            dataset,
                            config: Dict[str, Any],
                            device: torch.device,
                            num_samples: int = 100,
                            final_pause_frames: int = 20,
                            **kwargs) -> str:
    """
    Create sampling animation based on model type.
    
    Args:
        model_type: Type of model ('ddpm', 'sm', 'ncsm')
        model: Trained model
        dataset: Dataset
        config: Configuration dictionary
        device: Device to run on
        num_samples: Number of samples to generate
        final_pause_frames: Number of extra frames to show final result
        **kwargs: Additional arguments for specific animation types
        
    Returns:
        Path to saved animation
    """
    if model_type == 'ddpm':
        return create_ddpm_sampling_animation(model, num_samples, config, device, 
                                            final_pause_frames=final_pause_frames, **kwargs)
    elif model_type == 'sm':
        return create_langevin_sampling_animation(model, dataset, config, device, 
                                                num_samples=num_samples, 
                                                final_pause_frames=final_pause_frames, **kwargs)
    elif model_type == 'ncsm':
        noise_levels = kwargs.get('noise_levels', _get_noise_levels_from_config(config))
        return create_annealed_langevin_animation(model, noise_levels, dataset, 
                                                config, device, num_samples=num_samples, 
                                                final_pause_frames=final_pause_frames, **kwargs)
    else:
        raise ValueError(f"Unknown model type: {model_type}")


def create_score_visualizations(model_type: str,
                              model,
                              dataset,
                              config: Dict[str, Any],
                              device: torch.device,
                              noise_levels: Optional[List[float]] = None,
                              training_samples=None,
                              generated_samples=None,
                              **kwargs):
    """
    Create comprehensive score visualizations for SM and NCSM models.
    
    Args:
        model_type: Type of model ('sm', 'ncsm')
        model: Trained model
        dataset: Dataset
        config: Configuration dictionary
        device: Device to run on
        noise_levels: Noise levels for NCSM (optional)
        **kwargs: Additional arguments
    """
    if model_type not in ['sm', 'ncsm']:
        print(f"Score visualizations only available for SM and NCSM models, skipping for {model_type}")
        return
    
    print(f"Creating score visualizations for {model_type.upper()}...")
    
    if model_type == 'sm':
        # For SM, create single comparison and distance plot
        if training_samples is not None and generated_samples is not None:
            save_path_comp = os.path.join('./figures', 'sm_score_comparison.png')
            plot_score_comparison(model, model_type, training_samples, generated_samples, 
                                save_path=save_path_comp, show_plot=False, dataset=dataset)
        plot_score_distance(model_type, model, dataset, config, device, show=False)
    
    elif model_type == 'ncsm':
        # For NCSM, create both comparison and distance plot
        if noise_levels is None:
            noise_levels = _get_noise_levels_from_config(config)
        
        # Use only the first noise level
        noise_level = noise_levels[0]
        
        # Create score comparison
        if training_samples is not None and generated_samples is not None:
            save_path_comp = os.path.join('./figures', 'ncsm_score_comparison.png')
            plot_score_comparison(model, model_type, training_samples, generated_samples, 
                                save_path=save_path_comp, show_plot=False, dataset=dataset,
                                noise_levels=[noise_level])
        
        # Create score distance with simple filename
        plot_score_distance(model_type, model, dataset, config, device,
                           noise_level=noise_level, save_name='ncsm_score_distance.png', show=False)
    
    print(f"Score visualizations for {model_type.upper()} completed!")


def plot_true_pdf_scores(dataset,
                        n_grid: int = 15,
                        save_path: Optional[str] = None,
                        show_plot: bool = True,
                        figsize: Tuple[int, int] = FIGURE_SIZE) -> str:
    """
    Plot true PDF scores (normalized) with PDF contour lines.
    
    Args:
        dataset: Dataset instance with true_score and pdf methods
        n_grid: Grid size for score visualization
        save_path: Path to save the plot
        show_plot: Whether to show the plot
        figsize: Figure size
        
    Returns:
        Path to saved figure
    """
    # Create grid for score evaluation
    x = np.linspace(PLOT_BOUNDS[0], PLOT_BOUNDS[1], n_grid)
    y = np.linspace(PLOT_BOUNDS[2], PLOT_BOUNDS[3], n_grid)
    X, Y = np.meshgrid(x, y)
    
    # Prepare grid points
    grid_points = np.column_stack([X.ravel(), Y.ravel()])
    grid_tensor = torch.FloatTensor(grid_points)
    
    # Create figure
    fig, ax = plt.subplots(figsize=figsize)
    
    # Create finer grid for PDF contours
    x_pdf = np.linspace(PLOT_BOUNDS[0], PLOT_BOUNDS[1], 100)
    y_pdf = np.linspace(PLOT_BOUNDS[2], PLOT_BOUNDS[3], 100)
    X_pdf, Y_pdf = np.meshgrid(x_pdf, y_pdf)
    
    # Compute PDF values for contours
    pdf_values = dataset.pdf(X_pdf.ravel(), Y_pdf.ravel())
    pdf_grid = pdf_values.reshape(X_pdf.shape)
    
    # Add contour lines
    levels = np.linspace(0.1, pdf_values.max(), 8)
    ax.contour(X_pdf, Y_pdf, pdf_grid, levels=levels, colors=COLORS['contour'], 
               alpha=0.6, linewidths=0.8)
    
    # Get true scores
    true_scores = dataset.true_score(grid_tensor[:, 0], grid_tensor[:, 1])
    if isinstance(true_scores, torch.Tensor):
        true_scores = true_scores.detach().cpu().numpy()
    
    # Plot true score field
    U_true = true_scores[:, 0].reshape(n_grid, n_grid)
    V_true = true_scores[:, 1].reshape(n_grid, n_grid)
    
    # Calculate magnitude and normalize for visibility
    M_true = np.sqrt(U_true**2 + V_true**2)
    M_true_max = np.max(M_true)
    
    if M_true_max > 0:
        scale_factor = 1.0 / M_true_max
        U_true_scaled = U_true * scale_factor
        V_true_scaled = V_true * scale_factor
    else:
        U_true_scaled = U_true
        V_true_scaled = V_true
    
    # Plot score field as arrows
    ax.quiver(X, Y, U_true_scaled, V_true_scaled, 
              scale=10.0, scale_units='xy', angles='xy', alpha=0.7, 
              color=COLORS['training'], width=0.003)
    
    # Set bounds and styling
    ax.set_xlim(PLOT_BOUNDS[0], PLOT_BOUNDS[1])
    ax.set_ylim(PLOT_BOUNDS[2], PLOT_BOUNDS[3])
    ax.set_xlabel('X', fontsize=FONT_SIZES['axis_label'])
    ax.set_ylabel('Y', fontsize=FONT_SIZES['axis_label'])
    ax.set_title('True PDF Score Field with PDF Contours', fontsize=FONT_SIZES['title'])
    ax.grid(True, alpha=0.3)
    ax.tick_params(axis='both', which='major', labelsize=FONT_SIZES['tick_label'])
    ax.set_aspect('equal')
    
    plt.tight_layout()
    
    # Save to figures directory
    if save_path is None:
        os.makedirs('./figures', exist_ok=True)
        save_path = os.path.join('./figures', 'true_pdf_scores.png')
    
    plt.savefig(save_path, dpi=DPI, bbox_inches='tight')
    print(f"Saved true PDF scores plot: {save_path}")
    
    if show_plot:
        plt.show()
    else:
        plt.close()
    
    return save_path





def create_noise_effect_animation(dataset,
                                 noise_levels: Optional[List[float]] = None,
                                 n_grid: int = 15,
                                 n_pdf_grid: int = 100,
                                 save_name: Optional[str] = None,
                                 interval: int = 500,
                                 final_pause_frames: int = 20,
                                 figsize: Tuple[int, int] = FIGURE_SIZE) -> str:
    """
    Create animation showing how noise addition affects both the PDF and score field in NCSM.
    
    Args:
        dataset: Dataset instance
        noise_levels: List of noise levels (σ values) to animate
        n_grid: Grid size for score visualization
        n_pdf_grid: Grid size for PDF contour visualization
        save_name: Optional custom save name
        interval: Animation frame interval in ms
        final_pause_frames: Number of extra frames to show final result
        figsize: Figure size
        
    Returns:
        Path to saved animation
    """
    # Default noise levels if not provided
    if noise_levels is None:
        noise_levels = np.linspace(0.1, 0.001, 20).tolist()  # From 0.1 to 0.001 (linear, high to low noise)
    
    # Create grids for score and PDF evaluation
    x_score = np.linspace(PLOT_BOUNDS[0], PLOT_BOUNDS[1], n_grid)
    y_score = np.linspace(PLOT_BOUNDS[2], PLOT_BOUNDS[3], n_grid)
    X_score, Y_score = np.meshgrid(x_score, y_score)
    
    x_pdf = np.linspace(PLOT_BOUNDS[0], PLOT_BOUNDS[1], n_pdf_grid)
    y_pdf = np.linspace(PLOT_BOUNDS[2], PLOT_BOUNDS[3], n_pdf_grid)
    X_pdf, Y_pdf = np.meshgrid(x_pdf, y_pdf)
    
    # Prepare grid points
    score_grid_points = np.column_stack([X_score.ravel(), Y_score.ravel()])
    score_grid_tensor = torch.FloatTensor(score_grid_points)
    
    # Pre-compute score fields and PDF fields for all noise levels
    print("Computing score fields and PDFs for different noise levels...")
    score_fields = []
    pdf_fields = []
    max_magnitude = 0
    
    for i, sigma in enumerate(noise_levels):
        # Compute noisy score field
        noisy_scores = dataset.true_score_with_noise(score_grid_tensor[:, 0], score_grid_tensor[:, 1], sigma)
        if isinstance(noisy_scores, torch.Tensor):
            noisy_scores = noisy_scores.detach().cpu().numpy()
        
        # Store the score field
        U = noisy_scores[:, 0].reshape(n_grid, n_grid)
        V = noisy_scores[:, 1].reshape(n_grid, n_grid)
        score_fields.append((U, V))
        
        # Track maximum magnitude for consistent scaling
        M = np.sqrt(U**2 + V**2)
        max_magnitude = max(max_magnitude, np.max(M))
        
        # Compute noisy PDF field
        noisy_pdf = dataset.pdf_with_noise(X_pdf.ravel(), Y_pdf.ravel(), sigma)
        pdf_grid = noisy_pdf.reshape(n_pdf_grid, n_pdf_grid)
        pdf_fields.append(pdf_grid)
        
        if (i + 1) % 5 == 0:
            print(f"  Computed {i + 1}/{len(noise_levels)} noise levels")
    
    # Create animation
    fig, ax = plt.subplots(figsize=figsize)
    ax.set_xlim(PLOT_BOUNDS[0], PLOT_BOUNDS[1])
    ax.set_ylim(PLOT_BOUNDS[2], PLOT_BOUNDS[3])
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.3)
    
    # Initialize PDF contour plot
    pdf_init = pdf_fields[0]
    levels = np.linspace(0.1, np.max(pdf_init), 8)
    contour = ax.contour(X_pdf, Y_pdf, pdf_init, levels=levels, colors=COLORS['contour'], 
                        alpha=0.6, linewidths=0.8)
    
    # Initialize quiver plot
    U_init, V_init = score_fields[0]
    if max_magnitude > 0:
        scale_factor = 1.0 / max_magnitude
        U_scaled = U_init * scale_factor
        V_scaled = V_init * scale_factor
    else:
        U_scaled = U_init
        V_scaled = V_init
    
    quiver = ax.quiver(X_score, Y_score, U_scaled, V_scaled, 
                      scale=10.0, scale_units='xy', angles='xy', alpha=0.7, 
                      color=COLORS['ncsm'], width=0.003)
    
    # Add noise level text
    noise_text = ax.text(0.02, 0.98, '', transform=ax.transAxes, 
                        fontsize=FONT_SIZES['annotation'], verticalalignment='top',
                        bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
    
    ax.set_xlabel('X', fontsize=FONT_SIZES['axis_label'])
    ax.set_ylabel('Y', fontsize=FONT_SIZES['axis_label'])
    ax.set_title('Noise Effect on PDF and Score Field', fontsize=FONT_SIZES['title'])
    ax.tick_params(axis='both', which='major', labelsize=FONT_SIZES['tick_label'])
    
    def animate(frame):
        # Clear previous contour
        for coll in ax.collections:
            if coll != quiver:  # Don't remove quiver plot
                coll.remove()
        
        # Show normal animation frames
        if frame < len(noise_levels):
            sigma = noise_levels[frame]
            U, V = score_fields[frame]
            pdf_grid = pdf_fields[frame]
            
            # Update PDF contours
            levels = np.linspace(0.1, np.max(pdf_grid), 8)
            ax.contour(X_pdf, Y_pdf, pdf_grid, levels=levels, colors=COLORS['contour'], 
                      alpha=0.6, linewidths=0.8)
            
            # Apply consistent scaling to score field
            if max_magnitude > 0:
                scale_factor = 1.0 / max_magnitude
                U_scaled = U * scale_factor
                V_scaled = V * scale_factor
            else:
                U_scaled = U
                V_scaled = V
            
            # Update quiver plot
            quiver.set_UVC(U_scaled, V_scaled)
            
            # Update noise level text
            noise_text.set_text(f'Noise Level σ = {sigma:.4f} (denoising)')
        else:
            # Show final frame for pause frames
            sigma = noise_levels[-1]
            pdf_grid = pdf_fields[-1]
            levels = np.linspace(0.1, np.max(pdf_grid), 8)
            ax.contour(X_pdf, Y_pdf, pdf_grid, levels=levels, colors=COLORS['contour'], 
                      alpha=0.6, linewidths=0.8)
            noise_text.set_text(f'True Distribution (σ = {sigma:.4f})')
            
        return [quiver, noise_text]
    
    # Create animation with extended frames (original + pause frames)
    total_frames = len(noise_levels) + final_pause_frames
    anim = FuncAnimation(fig, animate, frames=total_frames, 
                        interval=interval, blit=False, repeat=True)  # blit=False for contour updates
    
    # Save animation to figures directory
    os.makedirs('./figures', exist_ok=True)
    if save_name is None:
        save_name = 'noise_effect_animation.gif'
    save_path = os.path.join('./figures', save_name)
    
    # Save with PillowWriter
    writer = PillowWriter(fps=2)  # Slow FPS to see the effect clearly
    anim.save(save_path, writer=writer)
    plt.close()
    
    print(f"Saved noise effect animation: {save_path}")
    return save_path


def create_true_langevin_animation(dataset,
                                   num_samples: int = 100,
                                   num_steps: int = 1000,
                                   step_size: float = 0.01,
                                   save_name: Optional[str] = None,
                                   interval: int = 50,
                                   subsample_steps: int = 10,
                                   final_pause_frames: int = 20,
                                   figsize: Tuple[int, int] = FIGURE_SIZE) -> str:
    """
    Create animation of Langevin sampling using true PDF scores.
    
    Args:
        dataset: Dataset instance with true_score method
        num_samples: Number of samples to simulate
        num_steps: Number of Langevin steps
        step_size: Langevin step size (epsilon)
        save_name: Optional custom save name
        interval: Animation frame interval in ms
        subsample_steps: How many steps to skip between frames
        final_pause_frames: Number of extra frames to show final result
        figsize: Figure size
        
    Returns:
        Path to saved animation
    """
    print(f"Running Langevin sampling with true scores ({num_samples} samples, {num_steps} steps)...")
    
    # Initialize samples randomly within bounds
    x_min, x_max, y_min, y_max = dataset.get_bounds()
    current_samples = torch.rand(num_samples, 2)
    current_samples[:, 0] = current_samples[:, 0] * (x_max - x_min) + x_min
    current_samples[:, 1] = current_samples[:, 1] * (y_max - y_min) + y_min
    
    # Store trajectory
    trajectory = []
    step_indices = []
    
    # Langevin dynamics: x_{i+1} = x_i + ε * ∇_x log p(x_i) + √(2ε) * z_i
    for step in range(num_steps):
        # Store current state for animation
        if step % subsample_steps == 0:
            trajectory.append(current_samples.clone())
            step_indices.append(step)
        
        # Get true scores at current positions
        true_scores = dataset.true_score(current_samples[:, 0], current_samples[:, 1])
        if isinstance(true_scores, torch.Tensor):
            true_scores = true_scores.detach()
        else:
            true_scores = torch.tensor(true_scores, dtype=torch.float32)
        
        # Add noise term
        noise = torch.randn_like(current_samples) * np.sqrt(2 * step_size)
        
        # Langevin update
        current_samples = current_samples + step_size * true_scores + noise
        
        # Optional: clip to bounds to prevent samples from going too far
        current_samples[:, 0] = torch.clamp(current_samples[:, 0], x_min, x_max)
        current_samples[:, 1] = torch.clamp(current_samples[:, 1], y_min, y_max)
        
        if (step + 1) % 100 == 0:
            print(f"  Completed step {step + 1}/{num_steps}")
    
    # Add final state
    trajectory.append(current_samples.clone())
    step_indices.append(num_steps)
    
    print(f"Langevin sampling completed. Creating animation with {len(trajectory)} frames...")
    
    # Create animation
    fig, ax = plt.subplots(figsize=figsize)
    
    # Add PDF contour background
    x_pdf = np.linspace(x_min, x_max, 100)
    y_pdf = np.linspace(y_min, y_max, 100)
    X_pdf, Y_pdf = np.meshgrid(x_pdf, y_pdf)
    
    # Compute PDF values
    pdf_values = dataset.pdf(X_pdf.ravel(), Y_pdf.ravel())
    pdf_grid = pdf_values.reshape(X_pdf.shape)
    
    # Add contour lines
    levels = np.linspace(0.1, pdf_values.max(), 8)
    ax.contour(X_pdf, Y_pdf, pdf_grid, levels=levels, colors=COLORS['contour'], 
               alpha=0.6, linewidths=0.8)
    
    ax.set_xlim(x_min, x_max)
    ax.set_ylim(y_min, y_max)
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.3)
    ax.set_xlabel('X', fontsize=FONT_SIZES['axis_label'])
    ax.set_ylabel('Y', fontsize=FONT_SIZES['axis_label'])
    ax.set_title('True Score Langevin Sampling', fontsize=FONT_SIZES['title'])
    ax.tick_params(axis='both', which='major', labelsize=FONT_SIZES['tick_label'])
    
    # Initialize scatter plot
    scat = ax.scatter([], [], s=POINT_SIZE, alpha=0.6, c=COLORS['training'], edgecolors='none')
    
    # Add step text
    step_text = ax.text(0.02, 0.98, '', transform=ax.transAxes, 
                       fontsize=FONT_SIZES['annotation'], verticalalignment='top',
                       bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
    
    def animate(frame):
        # Show normal trajectory frames
        if frame < len(trajectory):
            data = trajectory[frame].detach().cpu().numpy()
            scat.set_offsets(data)
            
            # Update step text
            step = step_indices[frame] if frame < len(step_indices) else num_steps
            step_text.set_text(f'Step: {step}')
        else:
            # Show final frame for pause frames
            data = trajectory[-1].detach().cpu().numpy()
            scat.set_offsets(data)
            step_text.set_text('Final Result')
            
        return [scat, step_text]
    
    # Create animation with extended frames (original + pause frames)
    total_frames = len(trajectory) + final_pause_frames
    anim = FuncAnimation(fig, animate, frames=total_frames, 
                        interval=interval, blit=True, repeat=True)
    
    # Save animation to figures directory
    os.makedirs('./figures', exist_ok=True)
    if save_name is None:
        save_name = 'true_langevin_sampling.gif'
    save_path = os.path.join('./figures', save_name)
    
    # Save with PillowWriter
    writer = PillowWriter(fps=20)
    anim.save(save_path, writer=writer)
    plt.close()
    
    print(f"Saved true Langevin sampling animation: {save_path}")
    return save_path


def create_comprehensive_visualization(model_type: str,
                                     model,
                                     dataset,
                                     config: Dict[str, Any],
                                     device: torch.device,
                                     losses: List[float],
                                     num_samples: int = 2000,  # Use more samples by default
                                     create_animation: bool = True,
                                     final_pause_frames: int = 20,
                                     **kwargs):
    """
    Create all visualizations for a trained model.
    
    Args:
        model_type: Type of model ('ddpm', 'sm', 'ncsm')
        model: Trained model
        dataset: Dataset
        config: Configuration dictionary
        device: Device to run on
        losses: Training losses
        num_samples: Number of samples to generate
        create_animation: Whether to create sampling animation
        final_pause_frames: Number of extra frames to show final result in animations
        **kwargs: Additional arguments for specific methods
    """
    print(f"\nCreating visualizations for {model_type.upper()}...")
    
    # Adjust sample count for NCSM to avoid memory issues
    if model_type == 'ncsm' and num_samples > config.get('n_eval_samples_ncsm', 200):
        print(f"📉 Reducing sample count from {num_samples} to {config.get('n_eval_samples_ncsm', 200)} for NCSM (memory optimization)")
        num_samples = config.get('n_eval_samples_ncsm', 200)
    
    # 1. Plot training losses
    plot_losses(losses, model_type, config, show=False)
    
    # 2. Generate samples and create comparison with fixed seed training data
    model.eval()
    with torch.no_grad():
        # Generate training data samples with same seed as training
        original_seed = dataset.seed
        training_samples = dataset.sample(num_samples, return_tensor=True, seed=config.get('data_seed', 42))
        
        # Generate samples based on model type - SAME samples used in animation
        if model_type == 'ddpm':
            from sampling import create_ddpm_sampler
            sampler = create_ddpm_sampler(model, device)
            generated_samples, sampling_info = sampler.sample(num_samples)
        elif model_type == 'sm':
            from sampling import create_learned_sampler
            bounds = dataset.get_bounds()
            sampler = create_learned_sampler(model, bounds, device, dataset)
            generated_samples, sampling_info = sampler.sample(num_samples, config['langevin_num_steps'], config['langevin_step_size'], store_trajectory=False)
        elif model_type == 'ncsm':
            from sampling import create_annealed_sampler_from_ncsn
            noise_levels = kwargs.get('noise_levels', _get_noise_levels_from_config(config))
            sampler = create_annealed_sampler_from_ncsn(model, noise_levels, 
                                                      dataset, device)
            generated_samples, sampling_info = sampler.sample(num_samples, config['langevin_steps_per_noise'], base_step_size=config['langevin_base_step_size'], store_trajectory=False)
    
    # 3. Create data comparison with fixed seed training data and PDF contours
    os.makedirs('./figures', exist_ok=True)
    plot_data_comparison(training_samples, generated_samples, model_type, 
                        save_path=os.path.join('./figures', f'{model_type}_data_comparison.png'),
                        show_plot=False, dataset=dataset)
    
    # 4. Save samples grid
    save_samples_grid(generated_samples, model_type, config)
    
    # 5. Create score visualizations for SM and NCSM
    create_score_visualizations(model_type, model, dataset, config, device, 
                              training_samples=training_samples, generated_samples=generated_samples, 
                              **kwargs)
    
    # 6. Create sampling animation (if requested) - using SAME samples as comparison
    if create_animation:
        try:
            # Generate samples WITH trajectory for animation - same count as comparison
            if model_type == 'ddpm':
                from sampling import create_ddpm_sampler
                animation_sampler = create_ddpm_sampler(model, device)
                animation_samples, animation_info = animation_sampler.sample(num_samples, store_trajectory=True, 
                                                          store_every=10, verbose=False)
                
                # Create animation using the trajectory
                fig, ax = plt.subplots(figsize=FIGURE_SIZE)
                ax.set_xlim(PLOT_BOUNDS[0], PLOT_BOUNDS[1])
                ax.set_ylim(PLOT_BOUNDS[2], PLOT_BOUNDS[3])
                ax.set_aspect('equal')
                ax.grid(True, alpha=0.3)
                ax.set_title('DDPM Sampling Process')
                
                trajectory = animation_info['trajectory']
                timesteps = animation_info['timesteps']
                
                # Initialize scatter plot with bigger points
                scat = ax.scatter([], [], s=POINT_SIZE, alpha=0.6, c=COLORS['ddpm'], edgecolors='none')
                
                # Add timestep text
                time_text = ax.text(0.02, 0.98, '', transform=ax.transAxes, 
                                   fontsize=FONT_SIZES['annotation'], verticalalignment='top',
                                   bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
                
                def animate(frame):
                    # Show normal trajectory frames
                    if frame < len(trajectory):
                        data = trajectory[frame].detach().cpu().numpy()
                        scat.set_offsets(data)
                        
                        # Update timestep text
                        if frame < len(timesteps):
                            t = timesteps[frame]
                            time_text.set_text(f'Timestep: {t}')
                    else:
                        # Show final frame for pause frames
                        data = trajectory[-1].detach().cpu().numpy()
                        scat.set_offsets(data)
                        time_text.set_text('Final Result')
                        
                    return [scat, time_text]
                
                # Create animation with extended frames (original + pause frames)
                total_frames = len(trajectory) + final_pause_frames
                anim = FuncAnimation(fig, animate, frames=total_frames, 
                                    interval=100, blit=True, repeat=True)
                
                # Save animation to figures directory
                os.makedirs('./figures', exist_ok=True)
                save_path = os.path.join('./figures', 'ddpm_sampling_animation.gif')
                
                # Save with PillowWriter
                writer = PillowWriter(fps=10)
                anim.save(save_path, writer=writer)
                plt.close()
                
                print(f"Saved DDPM sampling animation: {save_path}")
                
            elif model_type == 'sm':
                from sampling import create_learned_sampler
                bounds = dataset.get_bounds()
                animation_sampler = create_learned_sampler(model, bounds, device, dataset)
                animation_samples, animation_info = animation_sampler.sample(num_samples, 1000, 0.01, 
                                                      store_trajectory=True, store_every=10, verbose=False)
                
                # Create animation
                fig, ax = plt.subplots(figsize=FIGURE_SIZE)
                ax.set_xlim(PLOT_BOUNDS[0], PLOT_BOUNDS[1])
                ax.set_ylim(PLOT_BOUNDS[2], PLOT_BOUNDS[3])
                ax.set_aspect('equal')
                ax.grid(True, alpha=0.3)
                ax.set_title('Score Matching Langevin Sampling')
                
                trajectory = animation_info['trajectory']
                
                # Initialize scatter plot with bigger points
                scat = ax.scatter([], [], s=POINT_SIZE, alpha=0.6, c=COLORS['sm'], edgecolors='none')
                
                # Add step text
                step_text = ax.text(0.02, 0.98, '', transform=ax.transAxes, 
                                   fontsize=FONT_SIZES['annotation'], verticalalignment='top',
                                   bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
                
                def animate(frame):
                    # Show normal trajectory frames
                    if frame < len(trajectory):
                        data = trajectory[frame].detach().cpu().numpy()
                        scat.set_offsets(data)
                        
                        # Update step text
                        step = frame * 10  # store_every=10
                        step_text.set_text(f'Step: {step}')
                        
                    return [scat, step_text]
                
                # Create animation with extended frames (original + pause frames)
                total_frames = len(trajectory) + final_pause_frames
                anim = FuncAnimation(fig, animate, frames=total_frames, 
                                    interval=50, blit=True, repeat=True)
                
                # Save animation to figures directory
                os.makedirs('./figures', exist_ok=True)
                save_path = os.path.join('./figures', 'sm_sampling_animation.gif')
                
                # Save with PillowWriter
                writer = PillowWriter(fps=20)
                anim.save(save_path, writer=writer)
                plt.close()
                
                print(f"Saved SM sampling animation: {save_path}")
                
            elif model_type == 'ncsm':
                from sampling import create_annealed_sampler_from_ncsn
                noise_levels = kwargs.get('noise_levels', _get_noise_levels_from_config(config))
                animation_sampler = create_annealed_sampler_from_ncsn(model, noise_levels, dataset, device)
                store_every = 5  # Store every N steps for animation
                animation_samples, animation_info = animation_sampler.sample(num_samples, config['langevin_steps_per_noise'], 
                                                          base_step_size=config['langevin_base_step_size'],
                                                          store_trajectory=True, store_every=store_every, verbose=False)
                
                # Create animation
                fig, ax = plt.subplots(figsize=FIGURE_SIZE)
                ax.set_xlim(PLOT_BOUNDS[0], PLOT_BOUNDS[1])
                ax.set_ylim(PLOT_BOUNDS[2], PLOT_BOUNDS[3])
                ax.set_aspect('equal')
                ax.grid(True, alpha=0.3)
                ax.set_title('NCSM Annealed Langevin Sampling')
                
                trajectory = animation_info['trajectory']
                noise_levels_used = animation_info.get('noise_levels', noise_levels)
                
                # Initialize scatter plot with bigger points
                scat = ax.scatter([], [], s=POINT_SIZE, alpha=0.6, c=COLORS['ncsm'], edgecolors='none')
                
                # Add noise level text
                noise_text = ax.text(0.02, 0.98, '', transform=ax.transAxes, 
                                    fontsize=FONT_SIZES['annotation'], verticalalignment='top',
                                    bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
                
                def animate(frame):
                    if frame < len(trajectory):
                        data = trajectory[frame].detach().cpu().numpy()
                        scat.set_offsets(data)
                        
                        # Update noise level text - USE ACTUAL PARAMETERS
                        frames_per_noise = config['langevin_steps_per_noise'] // store_every
                        noise_idx = min(frame // frames_per_noise,
                                      len(noise_levels_used) - 1)
                        if noise_idx < len(noise_levels_used):
                            noise_level = noise_levels_used[noise_idx]
                            noise_text.set_text(f'Noise Level: {noise_level:.3f}')
                    else:
                        # Show final frame for pause frames
                        data = trajectory[-1].detach().cpu().numpy()
                        scat.set_offsets(data)
                        noise_text.set_text('Final Result')
                        
                    return [scat, noise_text]
                
                # Create animation with extended frames (original + pause frames)
                total_frames = len(trajectory) + final_pause_frames
                anim = FuncAnimation(fig, animate, frames=total_frames, 
                                    interval=50, blit=True, repeat=True)
                
                # Save animation to figures directory
                os.makedirs('./figures', exist_ok=True)
                save_path = os.path.join('./figures', 'ncsm_sampling_animation.gif')
                
                # Save with PillowWriter
                writer = PillowWriter(fps=20)
                anim.save(save_path, writer=writer)
                plt.close()
                
                print(f"Saved NCSM sampling animation: {save_path}")
                
        except Exception as e:
            print(f"Warning: Could not create sampling animation: {e}")
    
    print(f"Visualizations for {model_type.upper()} completed!")