utils.py

import numpy as np
import scipy.signal
from scipy.stats import pearsonr
from scipy.signal import butter, filtfilt, resample
import torch
from torch.nn.utils.rnn import pad_sequence
import torch.nn as nn

from sklearn.metrics import mean_squared_error
# from dtaidistance import dtw
# from skimage.metrics import structural_similarity as ssim

def collate_fn_processed(batch):
    inearpcg_signals, target_signals, info = zip(*batch)
    inearpcg_signals = pad_sequence(inearpcg_signals, batch_first=True, padding_value=0)
    target_signals = pad_sequence(target_signals, batch_first=True, padding_value=0)
    return inearpcg_signals, target_signals, info

def collate_fn_info(batch):
    inearpcg, scg, info = zip(*batch)
    inearpcg = torch.stack(inearpcg)
    scg = torch.stack(scg)
    info = list(info)  # Now info is a list of [start_inearpcg_idx, ...] per sample
    return inearpcg, scg, info
    
def compute_similarity_score(groundtruth, prediction):
    scores = []
    for gt, pred in zip(groundtruth, prediction):
        score, _ = pearsonr(gt.flatten(), pred.flatten())
        scores.append(score)
    return sum(scores) / len(scores)  # Average similarity score

def compute_spectral_overlap_index(groundtruth, prediction, sampling_rate=1000, freq_range=(1, 50)):
    """
    Computes the Spectral Overlap Index (SOI) between ground truth and prediction signals.

    Args:
        groundtruth (numpy.ndarray): Ground truth signal (1D array).
        prediction (numpy.ndarray): Predicted signal (1D array).
        sampling_rate (int): Sampling rate of the signals in Hz. Default is 1000Hz.
        freq_range (tuple): Frequency range for SOI calculation (low_freq, high_freq).

    Returns:
        float: Spectral Overlap Index (SOI) value.
    """
    # Compute the Fourier Transform of both signals
    freqs, gt_psd = scipy.signal.welch(groundtruth, fs=sampling_rate, nperseg=32)
    _, pred_psd = scipy.signal.welch(prediction, fs=sampling_rate, nperseg=32)

    # Limit to the specified frequency range
    freq_mask = (freqs >= freq_range[0]) & (freqs <= freq_range[1])
    gt_psd = gt_psd[freq_mask]
    pred_psd = pred_psd[freq_mask]

    # Compute the Spectral Overlap Index
    soi = np.sum(np.minimum(gt_psd, pred_psd)) / np.sum(gt_psd)
    return soi

class TemplateLoss(nn.Module):
    def __init__(self, lambda_param=0.1):
        """
        Custom loss function for signal reconstruction.
        Args:
            lambda_param (float): Weight for the Maxpool term.
        """
        super(TemplateLoss, self).__init__()
        self.lambda_param = lambda_param

    def forward(self, template, maxpool_output):
        """
        Compute the loss.
        Args:
            template (torch.Tensor): The learned convolution filter (template T).
            maxpool_output (torch.Tensor): The output of the Maxpool layer.
        Returns:
            torch.Tensor: The computed loss.
        """
        # L2-Norm Regularization Term
        l2_norm = torch.norm(template, p=2)  # L2 norm of the template

        # Maxpool Output Term
        maxpool_term = torch.mean(1.0 / (maxpool_output + 1e-8))  # Avoid division by zero

        # Combined Loss
        loss = l2_norm + self.lambda_param * maxpool_term
        return loss
    

def butter_bandpass(lowcut, highcut, fs, order=4):
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype='band')
    return b, a

def bandpass_filter(data, lowcut, highcut, fs, order=4):
    b, a = butter_bandpass(lowcut, highcut, fs, order=order)
    y = filtfilt(b, a, data)
    return y

def compute_cross_correlation_similarity(pred, target):
    """
    Compute maximum cross-correlation coefficient (handles time offset).
    
    Returns:
        max_corr (float): Max correlation coefficient (0-1)
        best_lag (int): Optimal time shift in samples
    """
    # Normalize signals
    pred_norm = (pred - np.mean(pred)) / (np.std(pred) + 1e-8)
    target_norm = (target - np.mean(target)) / (np.std(target) + 1e-8)
    
    # Compute cross-correlation
    correlation = scipy.signal.correlate(target_norm, pred_norm, mode='same')
    correlation = correlation / len(pred)
    
    # Find max correlation and corresponding lag
    max_corr = np.max(np.abs(correlation))
    best_lag = np.argmax(np.abs(correlation)) - len(pred) // 2
    
    return max_corr, best_lag


def compute_dtw_distance(pred, target, window=50):
    """
    Dynamic Time Warping distance (handles both time and amplitude offset).
    Lower is better.
    
    Args:
        window (int): Sakoe-Chiba window constraint (limits warping)
    
    Returns:
        dtw_dist (float): DTW distance
        normalized_dtw (float): Normalized by sequence length (0-1 scale)
    """
    try:
        from dtaidistance import dtw
        dtw_dist = dtw.distance(pred, target, window=window)
        normalized_dtw = dtw_dist / len(pred)
        return dtw_dist, normalized_dtw
    except ImportError:
        print("Warning: dtaidistance not installed. Run: pip install dtaidistance")
        return None, None


def compute_phase_corrected_mse(pred, target, max_shift=100):
    """
    MSE after aligning signals via cross-correlation.
    
    Args:
        max_shift (int): Maximum allowed shift in samples
    
    Returns:
        aligned_mse (float): MSE after optimal alignment
        best_shift (int): Optimal shift applied
    """
    best_mse = float('inf')
    best_shift = 0
    
    for shift in range(-max_shift, max_shift + 1):
        if shift < 0:
            pred_shifted = pred[-shift:]
            target_shifted = target[:shift]
        elif shift > 0:
            pred_shifted = pred[:-shift]
            target_shifted = target[shift:]
        else:
            pred_shifted = pred
            target_shifted = target
        
        if len(pred_shifted) > 0:
            mse = mean_squared_error(target_shifted, pred_shifted)
            if mse < best_mse:
                best_mse = mse
                best_shift = shift
    
    return best_mse, best_shift


def compute_envelope_similarity(pred, target):
    """
    Compare signal envelopes (amplitude modulation, ignoring fine details).
    Useful for cardiac signals where overall shape matters more than exact phase.
    
    Returns:
        envelope_corr (float): Pearson correlation of envelopes (0-1)
    """
    from scipy.signal import hilbert
    
    # Extract envelopes using Hilbert transform
    pred_envelope = np.abs(hilbert(pred))
    target_envelope = np.abs(hilbert(target))
    
    # Compute correlation
    envelope_corr, _ = pearsonr(pred_envelope, target_envelope)
    
    return envelope_corr


def compute_structural_similarity(pred, target, window_size=11):
    """
    Structural Similarity Index (SSIM) adapted for 1D signals.
    Originally designed for images but works for time series.
    
    Returns:
        ssim_score (float): SSIM score (0-1, higher is better)
    """
    
    # SSIM requires at least a 7x7 image, but we have 1D signals
    # Solution: Tile the signal vertically to create a pseudo-2D image
    # This allows SSIM to work while maintaining the temporal structure
    
    min_height = 7  # Minimum height required by SSIM
    signal_len = len(pred)
    
    # Create 2D representation by repeating the signal vertically
    pred_2d = np.tile(pred, (min_height, 1))
    target_2d = np.tile(target, (min_height, 1))
    
    # Calculate appropriate window size (must be odd and at least 3)
    max_win_size = min(signal_len, min_height, window_size)
    if max_win_size % 2 == 0:
        max_win_size -= 1
    max_win_size = max(3, max_win_size)
    
    # Compute SSIM
    data_range = max(target_2d.max() - target_2d.min(), 1e-8)
    ssim_score = ssim(target_2d, pred_2d, data_range=data_range, win_size=max_win_size)
    
    return ssim_score


def compute_frequency_domain_similarity(pred, target, fs=500):
    """
    Compare power spectral densities (frequency content).
    Robust to time shifts and DC offsets.
    
    Returns:
        freq_corr (float): Correlation of PSDs (0-1)
        spectral_distance (float): L2 distance between normalized PSDs
    """
    from scipy.signal import welch
    
    # Compute power spectral density
    f_pred, psd_pred = welch(pred, fs=fs, nperseg=min(256, len(pred)//4))
    f_target, psd_target = welch(target, fs=fs, nperseg=min(256, len(target)//4))
    
    # Normalize PSDs
    psd_pred_norm = psd_pred / (np.sum(psd_pred) + 1e-8)
    psd_target_norm = psd_target / (np.sum(psd_target) + 1e-8)
    
    # Correlation
    freq_corr, _ = pearsonr(psd_pred_norm, psd_target_norm)
    
    # L2 distance
    spectral_distance = np.linalg.norm(psd_pred_norm - psd_target_norm)
    
    return freq_corr, spectral_distance


def compute_all_offset_robust_metrics(pred, target, fs=500):
    """
    Compute all offset-robust metrics at once.
    
    Returns:
        dict: Dictionary with all computed metrics
    """
    metrics = {}
    
    # 1. Cross-correlation (handles time offset)
    max_corr, best_lag = compute_cross_correlation_similarity(pred, target)
    metrics['cross_corr'] = max_corr
    metrics['time_lag'] = best_lag
    metrics['time_lag_ms'] = (best_lag / fs) * 1000  # Convert to ms
    
    # 2. Phase-corrected MSE
    aligned_mse, best_shift = compute_phase_corrected_mse(pred, target, max_shift=100)
    metrics['aligned_mse'] = aligned_mse
    metrics['best_shift'] = best_shift
    
    # 3. DTW distance (if available)
    dtw_dist, norm_dtw = compute_dtw_distance(pred, target, window=50)
    if dtw_dist is not None:
        metrics['dtw_distance'] = dtw_dist
        metrics['normalized_dtw'] = norm_dtw
    
    # 4. Envelope similarity
    envelope_corr = compute_envelope_similarity(pred, target)
    metrics['envelope_corr'] = envelope_corr
    
    # 5. SSIM
    ssim_score = compute_structural_similarity(pred, target)
    metrics['ssim'] = ssim_score
    
    # 6. Frequency domain
    freq_corr, spec_dist = compute_frequency_domain_similarity(pred, target, fs=fs)
    metrics['freq_corr'] = freq_corr
    metrics['spectral_distance'] = spec_dist
    
    # 7. Original metrics (for comparison)
    metrics['pearson_corr'] = pearsonr(pred, target)[0]
    metrics['raw_mse'] = mean_squared_error(target, pred)
    
    return metrics