main/evaluate_visualize.py

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns
import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from main.GEMS import MaskedAutoencoderRaman
from sklearn.metrics import confusion_matrix, classification_report, roc_curve, auc
from sklearn.ensemble import IsolationForest
from sklearn.preprocessing import label_binarize
import torch
from mpl_toolkits.axes_grid1.inset_locator import mark_inset
from sklearn.metrics.pairwise import cosine_distances
from scipy import signal
from sklearn.manifold import TSNE

from sklearn.decomposition import PCA
plt.rcParams['font.sans-serif'] = ['WenQuanYi Micro Hei', 'SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

size = {"single": 84, "double": 170}

def mm_to_inches(mm):
    return mm / 25.4

def generate_snr_report(snr_results, save_dir, status):
    methods = ['peak_to_noise', 'rms', 'mad', 'baseline_corrected', 'multi_peak']
    spectrum_types = ['pure_snr', 'reconstructed_snr']
    stats = {}
    for spectrum_type in spectrum_types:
        stats[spectrum_type] = {}
        for method in methods:
            snr_values = []
            for result in snr_results:
                snr_db = result[spectrum_type].get(method, {}).get('snr_db', np.nan)
                if not np.isnan(snr_db) and snr_db != float('inf'):
                    snr_values.append(snr_db)
            
            if snr_values:
                stats[spectrum_type][method] = {
                    'mean': np.mean(snr_values),
                    'std': np.std(snr_values),
                    'median': np.median(snr_values),
                    'min': np.min(snr_values),
                    'max': np.max(snr_values),
                    'count': len(snr_values)
                }
            else:
                stats[spectrum_type][method] = {
                    'mean': np.nan, 'std': np.nan, 'median': np.nan,
                    'min': np.nan, 'max': np.nan, 'count': 0
                }
    if status == 'finetune':
        report_filename = 'fine_tune_snr_report.txt'
    elif status == 'pretrain':
        report_filename = 'pretrain_snr_report.txt'
    elif status == 'downtask':
        report_filename = 'downstream_snr_report.txt'

    report_path = os.path.join(save_dir, report_filename)

    with open(report_path, 'w', encoding='utf-8') as f:
        f.write("=" * 60 + "\n")
        f.write("raman spectra (SNR) analysis report\n")
        f.write("=" * 60 + "\n\n")
        
        f.write(f"Number of samples analyzed: {len(snr_results)}\n")
        if status == 'fine_tune':
            f.write(f"Training phase: {'Fine-tuning'}\n\n")
        elif status == 'pretrain':
            f.write(f"Training phase: {'Pre-training'}\n\n")
        elif status == 'downtask':
            f.write(f"Training phase: {'Downstream task'}\n\n")
        # Comparison table
        f.write("Statistics of various SNR calculation methods (dB):\n")
        f.write("-" * 80 + "\n")
        f.write(f"{'Method':<20} {'Spectrum Type':<15} {'Mean':<8} {'Std Dev':<8} {'Median':<8} {'Min':<8} {'Max':<8}\n")
        f.write("-" * 80 + "\n")
        
        for method in methods:
            for i, spectrum_type in enumerate(spectrum_types):
                type_name = {'pure_snr': 'Pure', 'reconstructed_snr': 'Reconstructed'}[spectrum_type]
                stat = stats[spectrum_type][method]
                
                method_name = method if i == 0 else ""
                f.write(f"{method_name:<20} {type_name:<15} {stat['mean']:<8.2f} {stat['std']:<8.2f} "
                       f"{stat['median']:<8.2f} {stat['min']:<8.2f} {stat['max']:<8.2f}\n")
            f.write("-" * 80 + "\n")
        
        # SNR improvement analysis
        f.write("\nSNR improvement analysis:\n")
        f.write("-" * 40 + "\n")
        
        for method in methods:
            orig_mean = stats['pure_snr'][method]['mean']
            recon_mean = stats['reconstructed_snr'][method]['mean']
            
            if not any(np.isnan([orig_mean, recon_mean])):
                recon_improvement = recon_mean - orig_mean
                
                f.write(f"{method}:\n")
                f.write(f"  Reconstruction relative to original: {recon_improvement:+.2f} dB\n")
    
    print(f"SNR analysis report saved to {report_path}")
    
    # Generate SNR comparison charts
    plot_snr_comparison(stats, save_dir, status)

def plot_snr_comparison(stats, save_dir, status):
    methods = ['peak_to_noise', 'rms', 'mad', 'baseline_corrected', 'multi_peak']
    spectrum_types = ['pure_snr', 'reconstructed_snr']
    type_labels = ['Pure', 'Reconstructed']
    colors = ['blue', 'red']
    
    fig, axes = plt.subplots(2, 3, figsize=(18, 12))
    axes = axes.flatten()
    
    for i, method in enumerate(methods):
        ax = axes[i]
        
        means = []
        stds = []
        labels = []
        
        for j, spectrum_type in enumerate(spectrum_types):
            stat = stats[spectrum_type][method]
            if not np.isnan(stat['mean']):
                means.append(stat['mean'])
                stds.append(stat['std'])
                labels.append(type_labels[j])
        
        if means:
            x = np.arange(len(labels))
            bars = ax.bar(x, means, yerr=stds, capsize=5, alpha=0.7, 
                         color=[colors[spectrum_types.index(st + '_snr')] for st in 
                               ['pure', 'reconstructed'] if st + '_snr' in 
                               [spectrum_types[k] for k in range(len(labels))]])
            
            ax.set_title(f'{method.replace("_", " ").title()} SNR', fontsize=12)
            ax.set_ylabel('SNR (dB)')
            ax.set_xticks(x)
            ax.set_xticklabels(labels, rotation=45)
            ax.grid(True, alpha=0.3)
        
            for bar, mean, std in zip(bars, means, stds):
                height = bar.get_height()
                ax.text(bar.get_x() + bar.get_width()/2., height + std + 0.5,
                       f'{mean:.1f}', ha='center', va='bottom', fontsize=10)
    
    if len(methods) < len(axes):
        for i in range(len(methods), len(axes)):
            fig.delaxes(axes[i])
    
    plt.tight_layout()
    
    if status == 'finetune':
        chart_filename = 'fine_tune_snr_comparison.png'
    elif status == 'pretrain':
        chart_filename = 'pretrain_snr_comparison.png'
    elif status == 'downtask':
        chart_filename = 'downstream_snr_comparison.png'
    chart_path = os.path.join(save_dir, chart_filename)
    plt.savefig(chart_path, dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"SNR Comparison chart saved to {chart_path}")


def visualize_transformed_and_reconstructed(model, status, test_dataset, wavenumbers, device, save_dir, num_samples=100):
    model.eval()
    os.makedirs(save_dir, exist_ok=True)
    indices = np.random.choice(len(test_dataset), num_samples, replace=False)
    all_snr_results = []
    rows, cols = 4, 4
    n_plots = rows * cols
    plt.figure(figsize=(16, 12))
    plot_indices = indices[:n_plots]
    fig, axes = plt.subplots(rows, cols, figsize=(16, 12))
    axes = axes.flatten()

    for ax in axes[len(plot_indices):]:
        ax.axis('off')

    for i, idx in enumerate(plot_indices):
        ax = axes[i]
        data_item = test_dataset[idx]
        augumented_spectra = data_item[1].unsqueeze(0).to(device) 
        processed_spectra = data_item[0].unsqueeze(0).to(device)  
        mask_ratio = 0.5
        with torch.no_grad():
            reconstructed, embedding, mask, loss = model(processed_spectra, mask_ratio=mask_ratio, tgt=processed_spectra)

        reconstructed = reconstructed.view(augumented_spectra.size(0), -1)  # (1, signal_length)

        original_np = processed_spectra.cpu().squeeze().numpy()
        reconstructed_np = reconstructed.cpu().squeeze().numpy()
        processed_np = augumented_spectra.cpu().squeeze().numpy()
        original_snr = comprehensive_snr_analysis(original_np)
        reconstructed_snr = comprehensive_snr_analysis(reconstructed_np)
        sample_results = {
            'sample_idx': idx,
            'mask_ratio': mask_ratio,
            'pure_snr': original_snr,
            'reconstructed_snr': reconstructed_snr
        }
        all_snr_results.append(sample_results)
        color1 = plt.cm.tab20c.colors[0]
        color2 = plt.cm.tab20c.colors[4]

        ax.plot(wavenumbers[400:1800], original_np[400:1800], label='Processed', linewidth=1.5, alpha=0.4, color=color2)
        ax.plot(wavenumbers[400:1800], reconstructed_np[400:1800], label=f'Reconstructed (mask={mask_ratio:.2f})', linewidth=1, alpha=0.9, color=color1)

        orig_snr_peak = original_snr.get('peak_to_noise', {}).get('snr_db', np.nan)
        recon_snr_peak = reconstructed_snr.get('peak_to_noise', {}).get('snr_db', np.nan)

        title = f'Sample {i + 1} - Pure: {orig_snr_peak:.1f} dB | Recon: {recon_snr_peak:.1f} dB'
        ax.set_title(title, fontsize=9)
        ax.set_ylabel('Intensity', fontsize=8)
        ax.grid(True, alpha=0.3)
        ax.legend(fontsize=7)
        
        if i // cols == rows - 1:
            ax.set_xlabel('Raman Shift (cm-1)', fontsize=8)
        else:
            ax.set_xticklabels([])

    plt.tight_layout()
    
    if status == 'finetune':
        save_path = os.path.join(save_dir, 'fine_tune_spectrum_snr_comparison.png')
    elif status == 'pretrain':
        save_path = os.path.join(save_dir, 'spectrum_snr_comparison.png')
    elif status == 'downtask':
        save_path = os.path.join(save_dir, 'downstream_spectrum_snr_comparison.png')
    # plt.show()
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()
    
    generate_snr_report(all_snr_results, save_dir, status)

    print(f"spectral SNR comparison saved to {save_path}")
    
    return all_snr_results

def load_and_visualize_mae_model(model_path, status, test_dataset, device, save_dir, input_length, wavenumbers, patch_num=100,
                                 embedding_dim=128, num_heads=16, num_layers=12):
    
    if status == "fine_tune":
        model_path = os.path.join(model_path, 'Fine_tuned.pth')
    elif status == "pretrain":
        model_path = os.path.join(model_path, 'Pretexted.pth')
    elif status == "downtask":
        model_path = model_path
    print(f"loading model: {model_path}")

    mae_model = MaskedAutoencoderRaman(
        input_length=input_length,
        patch_num=patch_num,
        embed_dim=embedding_dim,
        depth=num_layers,
        num_heads=num_heads,
        decoder_embed_dim=embedding_dim // 2,
        decoder_depth=4,
        decoder_num_heads=num_heads // 2
    ).to(device)

    checkpoint = torch.load(model_path, map_location=device)

    # Load model weights based on checkpoint structure
    if 'model_state_dict' in checkpoint:
        mae_model.load_state_dict(checkpoint['model_state_dict'])
        print("✅ Successfully loaded pretrained model weights")
    elif 'state_dict' in checkpoint:
        mae_model.load_state_dict(checkpoint['state_dict'])
        print("✅ Successfully loaded pretrained model weights")
    else:
        # If the checkpoint directly contains model weights
        mae_model.load_state_dict(checkpoint)
        print("✅ Successfully loaded pretrained model weights")
    mae_model.eval()
    print("Generating visualization results...")
    visualize_transformed_and_reconstructed(mae_model, status, test_dataset, wavenumbers, device, save_dir)

    return mae_model


def calculate_snr_methods(spectrum, method='peak_to_noise'):
    if isinstance(spectrum, torch.Tensor):
        spectrum = spectrum.detach().cpu().numpy()
    
    spectrum = spectrum.flatten()
    
    if method == 'peak_to_noise':
        peaks, properties = signal.find_peaks(spectrum, height=np.mean(spectrum) + 2*np.std(spectrum))
        
        if len(peaks) > 0:
            max_peak_idx = peaks[np.argmax(spectrum[peaks])]
            signal_intensity = spectrum[max_peak_idx]
            baseline_mask = spectrum < np.percentile(spectrum, 25)
            if np.sum(baseline_mask) > 10:
                noise_level = np.std(spectrum[baseline_mask])
            else:
                noise_level = np.std(spectrum) * 0.1 
            
            if noise_level > 0:
                snr_linear = signal_intensity / noise_level
                snr_db = 20 * np.log10(snr_linear)
            else:
                snr_db = float('inf')
            
            details = {
                'signal_intensity': signal_intensity,
                'noise_level': noise_level,
                'peak_position': max_peak_idx,
                'num_peaks': len(peaks)
            }
        else:
            snr_db = 0
            details = {'error': 'No peaks found'}
    
    elif method == 'rms':
        signal_rms = np.sqrt(np.mean(spectrum**2))
        spectrum_smooth = signal.savgol_filter(spectrum, 
                                             window_length=min(51, len(spectrum)//10*2+1), 
                                             polyorder=3)
        noise = spectrum - spectrum_smooth
        noise_rms = np.sqrt(np.mean(noise**2))
        if noise_rms > 0:
            snr_linear = signal_rms / noise_rms
            snr_db = 20 * np.log10(snr_linear)
        else:
            snr_db = float('inf')
        
        details = {
            'signal_rms': signal_rms,
            'noise_rms': noise_rms
        }
    
    elif method == 'mad':
        median_intensity = np.median(spectrum)
        mad = np.median(np.abs(spectrum - median_intensity))
        signal_intensity = np.max(spectrum)
        noise_level = 1.4826 * mad
        
        if noise_level > 0:
            snr_linear = (signal_intensity - median_intensity) / noise_level
            snr_db = 20 * np.log10(snr_linear)
        else:
            snr_db = float('inf')
        
        details = {
            'signal_intensity': signal_intensity,
            'median_intensity': median_intensity,
            'mad': mad,
            'noise_level': noise_level
        }
    
    elif method == 'baseline_corrected':
        x = np.arange(len(spectrum))
        baseline_points = []
        window_size = len(spectrum) // 20
        
        for i in range(0, len(spectrum), window_size):
            end_idx = min(i + window_size, len(spectrum))
            window = spectrum[i:end_idx]
            baseline_points.append(np.percentile(window, 5))  # 5th percentile as baseline
    
        baseline_x = np.linspace(0, len(spectrum)-1, len(baseline_points))
        baseline = np.interp(x, baseline_x, baseline_points)
        corrected_spectrum = spectrum - baseline
        signal_power = np.mean(corrected_spectrum[corrected_spectrum > 0]**2)
        noise_power = np.mean(corrected_spectrum[corrected_spectrum <= np.percentile(corrected_spectrum, 20)]**2)
        
        if noise_power > 0:
            snr_linear = signal_power / noise_power
            snr_db = 10 * np.log10(snr_linear) 
        else:
            snr_db = float('inf')
        
        details = {
            'signal_power': signal_power,
            'noise_power': noise_power,
            'baseline_corrected': True
        }
    
    elif method == 'multi_peak':
        peaks, properties = signal.find_peaks(spectrum, 
                                            height=np.mean(spectrum) + np.std(spectrum),
                                            distance=len(spectrum)//50)
        
        if len(peaks) >= 2:
            top_peaks = peaks[np.argsort(spectrum[peaks])[-3:]]
            
            snr_values = []
            for peak_idx in top_peaks:
                start_idx = max(0, peak_idx - 20)
                end_idx = min(len(spectrum), peak_idx + 20)
                local_region = spectrum[start_idx:end_idx]
                
                signal_intensity = spectrum[peak_idx]
                local_baseline = np.percentile(local_region, 10)
                local_noise = np.std(local_region[local_region < local_baseline + np.std(local_region)])
                
                if local_noise > 0:
                    local_snr = (signal_intensity - local_baseline) / local_noise
                    snr_values.append(20 * np.log10(local_snr))
            
            snr_db = np.mean(snr_values) if snr_values else 0
            
            details = {
                'num_peaks_analyzed': len(top_peaks),
                'individual_snrs': snr_values,
                'peak_positions': top_peaks.tolist()
            }
        else:
            return calculate_snr_methods(spectrum, method='peak_to_noise')
    
    else:
        raise ValueError(f"Unknown SNR calculation method: {method}")
    
    return snr_db, details

def comprehensive_snr_analysis(spectrum):
    methods = ['peak_to_noise', 'rms', 'mad', 'baseline_corrected', 'multi_peak']
    snr_results = {}
    
    for method in methods:
        try:
            snr_db, details = calculate_snr_methods(spectrum, method=method)
            snr_results[method] = {
                'snr_db': snr_db,
                'details': details
            }
        except Exception as e:
            snr_results[method] = {
                'snr_db': np.nan,
                'details': {'error': str(e)}
            }
    
    return snr_results

def remove_class_outliers(embeddings, labels, 
                           class_threshold=2.5,
                           enable_global_filter=True,
                           contamination=0.02):
    
    if enable_global_filter:
        print("Running Global Isolation Forest...")
        iso = IsolationForest(contamination=contamination, random_state=42, n_jobs=-1)
        global_mask = iso.fit_predict(embeddings) == 1
        
        embeddings = embeddings[global_mask]
        labels = labels[global_mask]
        print(f"Global filtering removed {np.sum(~global_mask)} samples.")

    cleaned_embeds, cleaned_labels = [], []

    for cls in np.unique(labels):
        mask = labels == cls
        cls_embed = embeddings[mask]
        if len(cls_embed) < 5:
            cleaned_embeds.append(cls_embed)
            cleaned_labels.append(labels[mask])
            continue

        center = cls_embed.mean(axis=0, keepdims=True)
        dist = cosine_distances(cls_embed, center).ravel()
        median_dist = np.median(dist)
        mad = np.median(np.abs(dist - median_dist))
        
        if mad == 0:
            keep = np.ones(len(dist), dtype=bool)
        else:
            mod_z_score = 0.6745 * (dist - median_dist) / mad
            keep = mod_z_score < class_threshold
            
        cleaned_embeds.append(cls_embed[keep])
        cleaned_labels.append(labels[mask][keep])
        
    return np.vstack(cleaned_embeds), np.concatenate(cleaned_labels)

def plot_tsne_embeddings(embeddings, labels, class_names=None, title=None, save_path=None, 
                        perplexity=30, max_iter=1000, 
                        figsize=(mm_to_inches(84), mm_to_inches(70)), 
                        outlier_z=2.5):
    
    embeddings, labels = remove_class_outliers(embeddings, labels, class_threshold=outlier_z)
    print(f"processing t-SNE... (perplexity={perplexity}, max_iter={max_iter})...")
    
    tsne = TSNE(
        n_components=2, 
        perplexity=min(perplexity, len(embeddings)//4),  
        max_iter=max_iter,
        random_state=42,
        learning_rate=400.0,  
        early_exaggeration=20.0,  
        init='pca',
        n_iter_without_progress=300,  
        method='barnes_hut',
        angle=0.3  
    )
    
    if isinstance(embeddings, torch.Tensor):
        embeddings = embeddings.cpu().numpy()
    if isinstance(labels, torch.Tensor):
        labels = labels.cpu().numpy()
    
    from sklearn.preprocessing import StandardScaler
    scaler = StandardScaler()
    embeddings_scaled = scaler.fit_transform(embeddings)

    n_samples, n_features = embeddings_scaled.shape
    n_components = min(50, n_samples, n_features)
    
    if n_features > 50 and n_samples > n_components:
        pca = PCA(n_components=n_components, random_state=42)
        embeddings_proc = pca.fit_transform(embeddings_scaled)
    else:
        embeddings_proc = embeddings_scaled

    effective_perplexity = min(max(perplexity, 5), max(5, len(embeddings_proc) // 10))

    tsne = TSNE(
        n_components=2,
        perplexity=effective_perplexity,
        max_iter=max(max_iter, 2000),
        random_state=42,
        learning_rate=800.0,
        early_exaggeration=36.0,
        init='pca',
        n_iter_without_progress=500,
        method='barnes_hut',
        angle=0.3
    )

    try:
        embeddings_2d = tsne.fit_transform(embeddings_proc)
        print("Completed t-SNE dimensionality reduction")
    except Exception as e:
        print(f"t-SNE dimensionality reduction failed: {e}")
        pca_fallback = PCA(n_components=2, random_state=42)
        embeddings_2d = pca_fallback.fit_transform(embeddings_proc)
        print("PCA dimensionality reduction completed")
    
    plt.rcParams.update({
        'font.size': 6,              
        'axes.labelsize': 6,         
        'axes.titlesize': 7,        
        'xtick.labelsize': 5,        
        'ytick.labelsize': 5,
        'legend.fontsize': 5,        
        'lines.linewidth': 0.4,      
        'axes.linewidth': 0.4,      
        'grid.linewidth': 0.3,       
        'xtick.major.width': 0.4,
        'ytick.major.width': 0.4,
        'font.family': 'sans-serif'
    })
    fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)
    
    unique_labels = np.unique(labels)
    colors = plt.cm.tab20(np.linspace(0, 1, len(unique_labels)))
    if class_names is None:
        class_names = [f"class_{i}" for i in unique_labels]

    legend_handles = []
    legend_labels_list = []

    for i, label in enumerate(unique_labels):
        mask = labels == label
        if isinstance(label, (int, np.integer)):
            if class_names is not None and 0 <= label < len(class_names):
                class_name = class_names[label]
            else:
                class_name = f"class_{label}"
        else:
            # If label is not an integer (e.g. string), use it directly
            class_name = str(label)
        
        sc = ax.scatter(
            embeddings_2d[mask, 0], 
            embeddings_2d[mask, 1],
            c=[colors[i]], 
            label=class_name, 
            alpha=0.8,
            s=3,
            edgecolors='none',
        )
        legend_handles.append(sc)
        legend_labels_list.append(class_name)
    
    for i, label in enumerate(unique_labels):
        mask = labels == label
        center_x = np.mean(embeddings_2d[mask, 0])
        center_y = np.mean(embeddings_2d[mask, 1])
        
        if label < len(class_names):
            class_name = class_names[label]
        else:
            class_name = f"class_{label}"
        
    if title: ax.set_title(title, pad=3)
    ax.set_xlabel('Dim 1', labelpad=1)
    ax.set_ylabel('Dim 2', labelpad=1)
    ax.tick_params(axis='both', which='major', pad=1, length=2)
    ax.grid(True, alpha=0.3, linestyle='--', linewidth=0.3)
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight', pad_inches=0.02)
        print(f"t-SNE saved: {save_path}")
    # plt.show()
    plt.close(fig)

    if save_path:
        dir_name, file_name = os.path.split(save_path)
        name_root, ext = os.path.splitext(file_name)
        legend_save_path = os.path.join(dir_name, f"{name_root}_legend{ext}")
        fig_leg = plt.figure(figsize=(3, 3)) 
        ax_leg = fig_leg.add_subplot(111)
        ax_leg.axis('off') 
        n_classes = len(unique_labels)
        n_cols = 4 if n_classes > 12 else (3 if n_classes > 6 else 1)
        leg = ax_leg.legend(
            legend_handles, 
            legend_labels_list,
            loc='center',       
            ncol=n_cols,             
            fontsize=7,      
            markerscale=3.0,    
            handletextpad=0.5, 
            columnspacing=1.0  
        )
        fig_leg.savefig(legend_save_path, dpi=300, bbox_inches='tight', pad_inches=0.02)
        print(f"Legend saved separately: {legend_save_path}")
        plt.close(fig_leg)
    
    # plt.show()
    
    return embeddings_2d

def plot_umap_embeddings(embeddings, labels, class_names=None, save_path=None,
                        n_neighbors=15, min_dist=0.1, figsize=(mm_to_inches(42), mm_to_inches(35))):
    from umap.umap_ import UMAP
    
    import matplotlib.pyplot as plt
    import numpy as np
    
    print(f"Starting UMAP dimensionality reduction (n_neighbors={n_neighbors}, min_dist={min_dist})...")
    
    reducer = UMAP(
        n_components=2,
        n_neighbors=n_neighbors,
        min_dist=min_dist,
        random_state=42
    )
    
    if isinstance(embeddings, torch.Tensor):
        embeddings = embeddings.cpu().numpy()
    if isinstance(labels, torch.Tensor):
        labels = labels.cpu().numpy()
    
    try:
        embeddings_2d = reducer.fit_transform(embeddings)
        print("UMAP dimensionality reduction completed")
    except Exception as e:
        print(f"UMAP dimensionality reduction failed: {e}")
        return None
    
    # Create visualization (same plotting logic as t-SNE)
    plt.rcParams.update({
        'font.size': 6,              # Global base font size
        'axes.labelsize': 6,         # Axis labels
        'xtick.labelsize': 5,        # Tick label size
        'ytick.labelsize': 5,
        'font.family': 'sans-serif',
        'lines.linewidth': 0.4,
        'axes.linewidth': 0.4,       # Thinner axis frame lines
        'grid.linewidth': 0.3
    })
  
    fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)
    unique_labels = np.unique(labels)
    colors = plt.cm.tab20(np.linspace(0, 1, len(unique_labels)))
    legend_handles = []
    legend_labels_list = []
    for i, label in enumerate(unique_labels):
        mask = labels == label
        if hasattr(class_names, '__getitem__') and label < len(class_names):
            c_name = class_names[label]
        else:
            c_name = f"Class {label}"
        
        sc = ax.scatter(
            embeddings_2d[mask, 0], 
            embeddings_2d[mask, 1],
            color=colors[i], 
            label=c_name,
            alpha=0.8,
            s=3,             
            edgecolors='none' 
        )
        legend_handles.append(sc)
        legend_labels_list.append(c_name)
    
    ax.set_title('UMAP', pad=3)
    ax.set_xlabel('Component 1', labelpad=1)
    ax.set_ylabel('Component 2', labelpad=1)
    ax.tick_params(axis='both', which='major', pad=1, length=2)
    ax.grid(True, alpha=0.3, linestyle='--', linewidth=0.3)
    
    if save_path:
        umap_save_path = save_path.replace('.png', '_umap.png')
        plt.savefig(umap_save_path, dpi=300, bbox_inches='tight')
        print(f"UMAP image saved as: {umap_save_path}")
    plt.close(fig)
    
    return embeddings_2d

def plot_confusion_matrix(y_true, y_pred, class_names=None, normalize=None,
                          title='Confusion Matrix', 
                          figsize=(mm_to_inches(84), mm_to_inches(70)), 
                          cmap='Blues',
                          save_path=None, 
                          fontsize=6):

    cm = confusion_matrix(y_true, y_pred)
    if normalize == 'true':
        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
        fmt = '.1%'
        # title = title + ' (normalized by true labels)'
    elif normalize == 'pred':
        cm = cm.astype('float') / cm.sum(axis=0)[np.newaxis, :]
        fmt = '.1%'
        # title = title + ' (normalized by predicted labels)'
    elif normalize == 'all':
        cm = cm.astype('float') / cm.sum()
        fmt = '.1%'
        # title = title + ' (normalized globally)'
    else:
        fmt = 'd'

    n_classes = cm.shape[0]
    if class_names is None:
        class_names = [f"C{i}" for i in range(n_classes)]
    elif len(class_names) < n_classes:
        class_names = list(class_names) + [f"C{i}" for i in range(len(class_names), n_classes)]
    
    fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)
    annot_size = fontsize if n_classes <= 10 else fontsize - 1
    sns.heatmap(cm, annot=True, fmt=fmt, cmap=cmap,
                xticklabels=class_names, yticklabels=class_names,
                cbar=True, square=True,
                linewidths=0.3, linecolor='white', 
                cbar_kws={"shrink": 0.7, "aspect": 15, "fraction": 0.05, "pad": 0.02},
                annot_kws={"size": annot_size, "weight": 'normal'}, 
                ax=ax)

    ax.set_xlabel('Predicted Label', fontsize=fontsize, labelpad=4)
    ax.set_ylabel('True Label', fontsize=fontsize, labelpad=4)
    max_label_len = max([len(str(n)) for n in class_names])

    if n_classes > 10 or max_label_len > 5:
        rotation_angle = 45 
        ha_mode = 'right'
    else:
        rotation_angle = 0
        ha_mode = 'center'
    
    ax.tick_params(axis='both', which='major', labelsize=fontsize, length=2, pad=2)
    plt.setp(ax.get_xticklabels(), rotation=rotation_angle, ha=ha_mode, rotation_mode="anchor")
    plt.setp(ax.get_yticklabels(), rotation=0)

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight', pad_inches=0.02)
        print(f"Confusion matrix saved as: {save_path}")

    plt.show()
    plt.close()

def plot_separation_heatmap(embeddings, labels, class_names=None,
                            metric='euclidean', 
                            figsize=(mm_to_inches(size['double']), mm_to_inches(size['double']*0.9)),
                            title='Class Separation', 
                            cmap='viridis',
                            save_path=None,
                            fontsize=6): 

    unique_labels = np.unique(labels)
    n_classes = len(unique_labels)
    
    if class_names is None:
        display_names = [f"C{i}" for i in unique_labels] 
    else:
        display_names = []
        for label in unique_labels:
            if label < len(class_names):
                display_names.append(class_names[label])
            else:
                display_names.append(f"C{label}")

    separation_matrix = np.zeros((n_classes, n_classes))
    centroids = []
    for label in unique_labels:
        mask = labels == label
        centroids.append(np.mean(embeddings[mask], axis=0))

    centroids = np.array(centroids)
    for i in range(n_classes):
        for j in range(n_classes):
            if metric == 'cosine':
                separation_matrix[i, j] = 1 - np.dot(centroids[i], centroids[j]) / (
                        np.linalg.norm(centroids[i]) * np.linalg.norm(centroids[j]) + 1e-8)
            else:
                separation_matrix[i, j] = np.linalg.norm(centroids[i] - centroids[j])

    fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)

    annot_size = fontsize if n_classes <= 10 else fontsize - 1.5
    
    fmt = '.2f' if metric == 'cosine' else '.1f'


    sns.heatmap(separation_matrix, annot=True, fmt=fmt, cmap=cmap,
                xticklabels=display_names, yticklabels=display_names,
                cbar=True, square=True,
                linewidths=0.3, linecolor='white', 
                cbar_kws={"shrink": 0.7, "aspect": 15, "fraction": 0.05, "pad": 0.02},
                annot_kws={"size": annot_size, "weight": 'normal'}, 
                ax=ax)

    ax.set_xlabel('Class Label', fontsize=fontsize+1, labelpad=4)
    ax.set_ylabel('Class Label', fontsize=fontsize+1, labelpad=4)
    
    if title:
        ax.set_title(title, fontsize=fontsize+2, pad=6, fontweight='bold')

    cbar = ax.collections[0].colorbar
    cbar.ax.tick_params(labelsize=fontsize-1)
    cbar_label = 'Dist.' if metric == 'euclidean' else 'Cos. Dist.'
    cbar.set_label(cbar_label, fontsize=fontsize, labelpad=4)

    max_label_len = max([len(str(n)) for n in display_names])
    if n_classes > 10 or max_label_len > 4:
        rotation_angle = 45 
        ha_mode = 'right'
    else:
        rotation_angle = 0
        ha_mode = 'center'

    ax.tick_params(axis='both', which='major', labelsize=fontsize, length=2, pad=2)
    plt.setp(ax.get_xticklabels(), rotation=rotation_angle, ha=ha_mode, rotation_mode="anchor")
    plt.setp(ax.get_yticklabels(), rotation=0)
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight', pad_inches=0.02)
        print(f"Separation heatmap saved as: {save_path}")

    plt.show()
    plt.close()


def plot_classification_metrics(y_true, y_pred, class_names=None,
                                title=None,  
                                figsize=(mm_to_inches(size['double']), mm_to_inches(size['double']*0.9)),
                                save_path=None):
    report = classification_report(y_true, y_pred, output_dict=True)
    unique_labels = np.unique(y_true)
    n_classes = len(unique_labels)

    if class_names is None:
        class_names = [f"Class {i}" for i in unique_labels]
    elif len(class_names) < n_classes:
        class_names = list(class_names) + [f"Class {i}" for i in range(len(class_names), n_classes)]
    
    display_names = []
    for label in unique_labels:
        idx = int(label) if isinstance(label, (int, float, np.integer)) else list(unique_labels).index(label)
        if idx < len(class_names):
            display_names.append(class_names[idx])
        else:
            display_names.append(f"{label}")

    metrics_data = {
        'precision': [report[str(label)]['precision'] for label in unique_labels],
        'recall': [report[str(label)]['recall'] for label in unique_labels],
        'f1-score': [report[str(label)]['f1-score'] for label in unique_labels]
    }

    plt.rcParams.update({
        'font.size': 6,             
        'axes.labelsize': 7,       
        'xtick.labelsize': 6,      
        'ytick.labelsize': 6,
        'legend.fontsize': 5,      
        'font.family': 'sans-serif',
        'lines.linewidth': 0.5,
        'axes.linewidth': 0.5
    })

    fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)

    colors = ['#4e79a7', '#f28e2b', '#76b7b2'] 
    
    bar_width = 0.25
    x = np.arange(n_classes)

    rects1 = ax.bar(x - bar_width, metrics_data['precision'], width=bar_width, color=colors[0], label='Precision', zorder=3)
    rects2 = ax.bar(x, metrics_data['recall'], width=bar_width, color=colors[1], label='Recall', zorder=3)
    rects3 = ax.bar(x + bar_width, metrics_data['f1-score'], width=bar_width, color=colors[2], label='F1 Score', zorder=3)

    def autolabel(rects):
        for rect in rects:
            height = rect.get_height()
            if height > 0: 
                ax.text(rect.get_x() + rect.get_width() / 2., height + 0.02,
                        f'{height:.2f}',
                        ha='center', va='bottom', 
                        rotation=90,      
                        fontsize=4.5)     

    autolabel(rects1)
    autolabel(rects2)
    autolabel(rects3)

    ax.set_ylabel('Score')

    if title:
        ax.set_title(title, fontsize=7, pad=4)


    ax.set_xticks(x)
    max_len = max([len(str(n)) for n in display_names])
    rot = 0 if max_len < 4 else (30 if max_len < 8 else 45)
    ax.set_xticklabels(display_names, rotation=rot, ha='right' if rot > 0 else 'center')

    ax.set_yticks(np.arange(0, 1.2, 0.2))
    ax.set_ylim(0, 1.25) 

    ax.grid(True, axis='y', linestyle='--', alpha=0.5, zorder=0)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)

    ax.legend(loc='upper center', bbox_to_anchor=(0.5, 1.0), 
              ncol=3, frameon=False, handletextpad=0.3)

    avg_precision = report['macro avg']['precision']
    avg_recall = report['macro avg']['recall']
    avg_f1 = report['macro avg']['f1-score']
    
    stats_text = (f"Macro Avg:\n"
                  f"P: {avg_precision:.2f}\n"
                  f"R: {avg_recall:.2f}\n"
                  f"F1: {avg_f1:.2f}")

    ax.text(0.98, 0.95, stats_text, transform=ax.transAxes,
            ha='right', va='top', fontsize=5,
            bbox=dict(boxstyle='round,pad=0.3', facecolor='white', alpha=0.8, edgecolor='gray', linewidth=0.3))

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight', pad_inches=0.02)
        print(f"✅ Chart saved as: {save_path}")
    else:
        plt.show()
    plt.close()


def plot_multi_roc_curve(y_true, y_score, class_names=None, average="macro",
                         title=None, 
                         figsize=(mm_to_inches(84), mm_to_inches(78)), 
                         save_path=None,
                         zoom_view=True):

    y_true = np.array(y_true, dtype=int)
    n_classes = y_score.shape[1]
    
    if class_names is None:
        class_names = [f"Class {i}" for i in range(n_classes)]
    
    y_true_bin = label_binarize(y_true, classes=range(n_classes))

    fpr = dict()
    tpr = dict()
    roc_auc = dict()

    for i in range(n_classes):
        fpr[i], tpr[i], _ = roc_curve(y_true_bin[:, i], y_score[:, i])
        roc_auc[i] = auc(fpr[i], tpr[i])

    # Micro-average
    fpr["micro"], tpr["micro"], _ = roc_curve(y_true_bin.ravel(), y_score.ravel())
    roc_auc["micro"] = auc(fpr["micro"], tpr["micro"])

    # Macro-average
    all_fpr = np.unique(np.concatenate([fpr[i] for i in range(n_classes)]))
    mean_tpr = np.zeros_like(all_fpr)
    for i in range(n_classes):
        mean_tpr += np.interp(all_fpr, fpr[i], tpr[i])
    mean_tpr /= n_classes
    fpr["macro"] = all_fpr
    tpr["macro"] = mean_tpr
    roc_auc["macro"] = auc(fpr["macro"], tpr["macro"])

    plt.rcParams.update({
        'font.size': 6,
        'axes.labelsize': 7,
        'xtick.labelsize': 6,
        'ytick.labelsize': 6,
        'lines.linewidth': 0.6,
        'axes.linewidth': 0.5,
        'grid.linewidth': 0.3,
        'font.family': 'sans-serif'
    })

    fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)

    cmap = plt.cm.tab20 if n_classes > 10 else plt.cm.tab10
    colors = cmap(np.linspace(0, 1, n_classes))

    legend_handles = []
    legend_labels = []

    ax.plot([0, 1], [0, 1], 'k--', lw=0.5, alpha=0.5, label='Random')

    for i, color in zip(range(n_classes), colors):
        label_str = f'{class_names[i]} ({roc_auc[i]:.2f})'
        l, = ax.plot(fpr[i], tpr[i], color=color, lw=0.6, alpha=0.6, label=label_str)
        legend_handles.append(l)
        legend_labels.append(label_str)

    l_micro, = ax.plot(fpr["micro"], tpr["micro"], color='deeppink', linestyle=':', lw=1.2, 
            label=f'Micro-avg ({roc_auc["micro"]:.2f})')
    l_macro, = ax.plot(fpr["macro"], tpr["macro"], color='navy', linestyle='--', lw=1.2, 
            label=f'Macro-avg ({roc_auc["macro"]:.2f})')

    legend_handles = [l_micro, l_macro] + legend_handles
    legend_labels = [f'Micro-avg ({roc_auc["micro"]:.2f})', f'Macro-avg ({roc_auc["macro"]:.2f})'] + legend_labels

    if zoom_view:

        axins = ax.inset_axes([0.45, 0.12, 0.48, 0.45])
        
        for i, color in zip(range(n_classes), colors):
            axins.plot(fpr[i], tpr[i], color=color, lw=0.8, alpha=0.8) 
        
        axins.plot(fpr["micro"], tpr["micro"], color='deeppink', linestyle=':', lw=1.2)
        axins.plot(fpr["macro"], tpr["macro"], color='navy', linestyle='--', lw=1.2)

        x1, x2, y1, y2 = 0.0, 0.1, 0.9, 1.01 
        axins.set_xlim(x1, x2)
        axins.set_ylim(y1, y2)
        

        axins.set_xticklabels([]) 
        axins.set_yticklabels([]) 
        axins.tick_params(axis='both', which='both', length=2) 
        axins.grid(True, linestyle='--', alpha=0.3)
        
        mark_inset(ax, axins, loc1=2, loc2=4, fc="none", ec="0.4", lw=0.5, linestyle='--')

    # --- 5. Axis labels and settings ---
    ax.set_xlabel('False Positive Rate (FPR)', labelpad=2)
    ax.set_ylabel('True Positive Rate (TPR)', labelpad=2)
    if title:
        ax.set_title(title, fontsize=7, pad=4)
    
    ax.grid(True, linestyle='--', alpha=0.4)
    ax.set_xlim([0.0, 1.0])
    ax.set_ylim([0.0, 1.02])

    if n_classes <= 5:
        ax.legend(loc='lower right', fontsize=5, frameon=False) 
        if zoom_view:
             ax.legend(loc='center right', fontsize=5, frameon=False, bbox_to_anchor=(1, 0.5))
    else:
        pass 

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight', pad_inches=0.02)
        print(f"ROC curve saved: {save_path}")
    
    plt.close(fig)

    if save_path and n_classes > 5:
        dir_name, file_name = os.path.split(save_path)
        name_root, ext = os.path.splitext(file_name)
        legend_save_path = os.path.join(dir_name, f"{name_root}_legend{ext}")

        fig_leg = plt.figure(figsize=(3, 3))
        ax_leg = fig_leg.add_subplot(111)
        ax_leg.axis('off')
        
        n_items = len(legend_labels)
        n_cols = 3 if n_items > 9 else 2

        ax_leg.legend(
            legend_handles, 
            legend_labels,
            loc='center',
            ncol=n_cols,
            frameon=False,
            fontsize=7,
            handlelength=1.5,
            columnspacing=1.0
        )

        fig_leg.savefig(legend_save_path, dpi=300, bbox_inches='tight', pad_inches=0.02)
        print(f"ROC Legend saved separately: {legend_save_path}")
        plt.close(fig_leg)


def plot_training_history(train_losses, val_losses, train_accuracies, val_accuracies,
                          figsize=(mm_to_inches(84), mm_to_inches(60)),
                          save_path=None):
    
    plt.rcParams.update({
        'font.size': 6,
        'axes.labelsize': 7,
        'xtick.labelsize': 6,
        'ytick.labelsize': 6,
        'legend.fontsize': 5,
        'lines.linewidth': 0.8,   
        'axes.linewidth': 0.5,
        'grid.linewidth': 0.3,
        'font.family': 'sans-serif'
    })
    fig, ax1 = plt.subplots(figsize=figsize, constrained_layout=True)

    if len(train_losses) == 0 and len(val_losses) == 0 and len(train_accuracies) == 0 and len(val_accuracies) == 0:
        print("⚠️ No training history to plot.")
        plt.close(fig)
        return

    color_loss = 'tab:blue'
    ax1.set_xlabel('Epoch', labelpad=2)
    ax1.set_ylabel('Loss', color=color_loss, labelpad=2)

    lines = []

    if len(train_losses) > 0:
        epochs_train_loss = range(1, len(train_losses) + 1)
        l1, = ax1.plot(epochs_train_loss, train_losses, color=color_loss, linestyle='-', alpha=0.8, label='Train Loss')
        lines.append(l1)
    if len(val_losses) > 0:
        epochs_val_loss = range(1, len(val_losses) + 1)
        l2, = ax1.plot(epochs_val_loss, val_losses, color=color_loss, linestyle='--', alpha=0.6, label='Val Loss')
        lines.append(l2)
    
    ax1.tick_params(axis='y', labelcolor=color_loss, pad=1, length=2)
    ax1.tick_params(axis='x', pad=1, length=2)
    ax1.grid(True, linestyle='--', alpha=0.3)

    ax2 = ax1.twinx()  
    color_acc = 'tab:red'
    ax2.set_ylabel('Accuracy', color=color_acc, labelpad=2)

    if len(train_accuracies) > 0:
        epochs_train_acc = range(1, len(train_accuracies) + 1)
        l3, = ax2.plot(epochs_train_acc, train_accuracies, color=color_acc, linestyle='-', alpha=0.8, label='Train Acc')
        lines.append(l3)
    if len(val_accuracies) > 0:
        epochs_val_acc = range(1, len(val_accuracies) + 1)
        l4, = ax2.plot(epochs_val_acc, val_accuracies, color=color_acc, linestyle='--', alpha=0.6, label='Val Acc')
        lines.append(l4)
    
    ax2.tick_params(axis='y', labelcolor=color_acc, pad=1, length=2)
    ax2.set_ylim([0, 1.05]) 

    if lines:
        labels = [l.get_label() for l in lines]
        ax1.legend(lines, labels, loc='center right', frameon=False)

    len_set = {len(train_losses), len(val_losses), len(train_accuracies), len(val_accuracies)}
    len_set.discard(0)
    if len(len_set) > 1:
        print(
            "⚠️ History length mismatch detected: "
            f"train_losses={len(train_losses)}, val_losses={len(val_losses)}, "
            f"train_accuracies={len(train_accuracies)}, val_accuracies={len(val_accuracies)}. "
            "Plotted each curve with its own epoch range."
        )

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight', pad_inches=0.02)
        print(f"Training history saved as: {save_path}")

    # plt.show()
    plt.close()


def visualize_model_performance(classifier, test_loader, device, class_names=None,
                                save_dir=None):

    if save_dir and not os.path.exists(save_dir):
        os.makedirs(save_dir)
        print(f"Created directory: {save_dir}")
    all_embeddings = []
    all_labels = []
    all_preds = []
    all_probs = []
    classifier.eval()
    with torch.no_grad():
        for batch in test_loader:
            if len(batch) == 3:
                inputs, _, labels = batch
            else:
                inputs, labels = batch
            inputs = inputs.to(device)
            labels = labels.to(device)
            logits, embeddings = classifier(inputs)
            probs = torch.softmax(logits, dim=1)
            preds = torch.argmax(logits, dim=1)
            all_embeddings.append(embeddings.cpu().numpy())
            all_labels.append(labels.cpu().numpy())
            all_preds.append(preds.cpu().numpy())
            all_probs.append(probs.cpu().numpy())
    all_embeddings = np.vstack(all_embeddings)
    all_labels = np.concatenate(all_labels)
    all_preds = np.concatenate(all_preds)
    all_probs = np.vstack(all_probs)
    n_classes = all_probs.shape[1]
    if class_names is None:
        class_names = [f"class{i}" for i in range(n_classes)]

    # 1. t-SNE
    tsne_path = os.path.join(save_dir, "tsne_visualization.png") if save_dir else None
    embeddings_2d = plot_tsne_embeddings(all_embeddings, all_labels, class_names=class_names,
                                         title='t-SNE', save_path=tsne_path)
    # Optional: UMAP visualization
    # umap_path = os.path.join(save_dir, "umap_visualization.png") if save_dir else None
    # plot_umap_embeddings(all_embeddings, all_labels, class_names=class_names,
    #                     save_path=umap_path)

    # 2. confusion matrix
    # cm_path = os.path.join(save_dir, "confusion_matrix.png") if save_dir else None
    # plot_confusion_matrix(all_labels, all_preds, class_names=class_names,
    #                       title='Confusion Matrix', save_path=cm_path)

    cm_norm_path = os.path.join(save_dir, "confusion_matrix_normalized.png") if save_dir else None
    plot_confusion_matrix(all_labels, all_preds, class_names=class_names, normalize='true',
                          title='Normalized Confusion Matrix', save_path=cm_norm_path)
    # 3. Class Separation Heatmap
    sep_path = os.path.join(save_dir, "class_separation_heatmap.png") if save_dir else None
    plot_separation_heatmap(all_embeddings, all_labels, class_names=class_names,
                            title='Class Separation Heatmap', save_path=sep_path)
    # 4. Classification Metrics Plot
    metrics_path = os.path.join(save_dir, "classification_metrics.png") if save_dir else None
    plot_classification_metrics(all_labels, all_preds, class_names=class_names,
                                title='Classification Metrics by Class', save_path=metrics_path)
    # 5. ROC Curves
    roc_path = os.path.join(save_dir, "roc_curves.png") if save_dir else None
    plot_multi_roc_curve(all_labels, all_probs, class_names=class_names,
                         title='Multi-class ROC Curves', save_path=roc_path)
    print("\nClassification Report:")
    print(classification_report(all_labels, all_preds, target_names=class_names))

    # Save classification report
    if save_dir:
        report_path = os.path.join(save_dir, "classification_report.txt")
        with open(report_path, 'w') as f:
            f.write(classification_report(all_labels, all_preds, target_names=class_names))
            print(f"Classification report saved to: {report_path}")

    return {
        'embeddings_2d': embeddings_2d,
        'true_labels': all_labels,
        'pred_labels': all_preds,
        'probabilities': all_probs,
        'class_names': class_names
    }