code/scripts/infer_v16_visualize.py

#!/usr/bin/env python3
"""
V16 Inference with Visualization

Runs V16 per-lead model on Kaggle images and plots predictions as dots on the strips.
V16 uses signal-region-only input (T0:T1) and 500px crop height.
"""

import os
import sys
import argparse
import random
import numpy as np
import pandas as pd
from pathlib import Path
from tqdm import tqdm
from scipy.signal import savgol_filter

import torch
import torch.nn as nn
import torch.nn.functional as F

import cv2
import timm


# =============================================================================
# Constants (from train_v16_perlead.py)
# =============================================================================
TARGET_HEIGHT, TARGET_WIDTH = 1696, 4352
ZERO_MV = np.array([703.5, 987.5, 1271.5, 1531.5])
MV_TO_PIXEL = 78.5
T0, T1 = 235, 4161
X0, X1 = 0, 2176
Y0, Y1 = 0, 1696
OUTPUT_WIDTH = T1 - T0  # 3926 - V16 input width

# Per-row crop parameters (V16 uses larger crop)
CROP_HALF_HEIGHT = 250
ROW_HEIGHT = 500

# ECG amplitude limits (mV) - values beyond this are likely errors
ECG_MV_MIN, ECG_MV_MAX = -10.0, 10.0

# SNR threshold for "bad" predictions
LOW_SNR_THRESHOLD = 5.0  # dB

VALID_VARIANTS = ['0001', '0003', '0004', '0005', '0006', '0009', '0010', '0011', '0012']

# Validation sample IDs (holdout set - same as training)
# VAL_SAMPLE_IDS = [
#     '1006427285', '1006867983', '1012423188', '10140238', '1015663939',
#     '102150619', '1026034238', '1041099777', '104573050', '1048962695',
#     '1052007218', '1053922973', '1059602762', '1063816858', '106482869',
#     '1067371646', '1067975047', '1068062585', '1072767337', '1079294623',
#     '1084993373', '108599929'
# ]

VAL_SAMPLE_IDS = [
    '1006427285', '1006867983', '1012423188', '10140238', '1015663939',
    '102150619', '1026034238', '1041099777', '104573050', '1048962695',
    '1052007218', '1053922973', '1059602762', '1063816858', #'106482869',
    '1067371646', '1067975047', '1068062585', '1072767337', #'1079294623',
    '1084993373', '108599929'
]

# Lead layout (same as training)
LEAD_LAYOUT = [
    ['I', 'aVR', 'V1', 'V4'],
    ['II', 'aVL', 'V2', 'V5'],
    ['III', 'aVF', 'V3', 'V6'],
]


# =============================================================================
# Model Architecture (copy from train_v16_perlead.py)
# =============================================================================
class CoordConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, **kwargs):
        super().__init__()
        self.conv = nn.Conv2d(in_channels + 2, out_channels, kernel_size, **kwargs)
    
    def forward(self, x):
        B, C, H, W = x.shape
        yy = torch.linspace(-1, 1, H, device=x.device).view(1, 1, H, 1).expand(B, 1, H, W)
        xx = torch.linspace(-1, 1, W, device=x.device).view(1, 1, 1, W).expand(B, 1, H, W)
        x = torch.cat([x, yy, xx], dim=1)
        return self.conv(x)


class UNetDecoderBlock(nn.Module):
    def __init__(self, in_ch, skip_ch, out_ch):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch + skip_ch, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.GELU(),
            nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.GELU(),
        )
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
    
    def forward(self, x, skip=None):
        x = self.upsample(x)
        if skip is not None:
            if x.shape[2:] != skip.shape[2:]:
                x = F.interpolate(x, size=skip.shape[2:], mode='bilinear', align_corners=True)
            x = torch.cat([x, skip], dim=1)
        return self.conv(x)


class PerLeadNet(nn.Module):
    def __init__(self, encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=True):
        super().__init__()
        
        self.encoder = timm.create_model(
            encoder_name,
            pretrained=pretrained,
            features_only=True,
            out_indices=(0, 1, 2, 3),
        )
        enc_channels = self.encoder.feature_info.channels()
        
        decoder_dims = [256, 128, 64, 32]
        self.dec_blocks = nn.ModuleList()
        in_ch = enc_channels[-1]
        skip_channels = enc_channels[:-1][::-1] + [0]
        
        for skip_ch, out_ch in zip(skip_channels, decoder_dims):
            self.dec_blocks.append(UNetDecoderBlock(in_ch, skip_ch, out_ch))
            in_ch = out_ch
        
        self.final_up = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            nn.Conv2d(decoder_dims[-1], decoder_dims[-1], 3, padding=1, bias=False),
            nn.BatchNorm2d(decoder_dims[-1]),
            nn.GELU(),
        )
        
        self.height_attention = nn.Sequential(
            CoordConv2d(decoder_dims[-1], 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.GELU(),
            nn.Conv2d(64, 1, 1),
        )
        
        self.regression_head = nn.Sequential(
            nn.Conv1d(decoder_dims[-1], 128, 7, padding=3),
            nn.BatchNorm1d(128),
            nn.GELU(),
            nn.Conv1d(128, 64, 5, padding=2),
            nn.BatchNorm1d(64),
            nn.GELU(),
            nn.Conv1d(64, 1, 1),
        )
    
    def forward(self, x):
        B, C, H, W = x.shape
        
        features = self.encoder(x)
        
        d = features[-1]
        skips = features[:-1][::-1] + [None]
        
        for block, skip in zip(self.dec_blocks, skips):
            d = block(d, skip)
        
        d = self.final_up(d)
        
        if d.shape[3] != W:
            d = F.interpolate(d, size=(d.shape[2], W), mode='bilinear', align_corners=True)
        
        attn = self.height_attention(d)
        attn = F.softmax(attn, dim=2)
        
        d = (d * attn).sum(dim=2)
        
        out = self.regression_head(d)
        out = torch.sigmoid(out)
        
        return out.squeeze(1)


# =============================================================================
# Inference Enhancement Functions (from V10 high-scoring solution)
# =============================================================================
def apply_savgol_smoothing(signal_mv, window=7, polyorder=2):
    """Apply Savitzky-Golay smoothing to remove high-frequency noise."""
    if len(signal_mv) >= window:
        return savgol_filter(signal_mv, window_length=window, polyorder=polyorder)
    return signal_mv


def apply_einthoven_correction(pred_mv_rows, alpha=0.33):
    """
    Apply Einthoven's law correction on short lead segments.
    
    Einthoven's Law: II = I + III (in mV)
    
    For each time point, if there's a violation e = II - (I + III), distribute error:
    - I'  = I  + α*e
    - III'= III + α*e  
    - II' = II - α*e (but we use rhythm strip, so skip)
    
    pred_mv_rows: dict with row predictions in mV
    """
    segment_width = len(pred_mv_rows[0]) // 4
    
    # Row 0 segment 0 = Lead I
    # Row 1 segment 0 = Lead II short (0-2.5s)
    # Row 2 segment 0 = Lead III
    
    lead_I = pred_mv_rows[0][:segment_width].copy()
    lead_II_short = pred_mv_rows[1][:segment_width].copy()
    lead_III = pred_mv_rows[2][:segment_width].copy()
    
    # Compute Einthoven violation: e = II - (I + III)
    derived_II = lead_I + lead_III
    error = lead_II_short - derived_II
    
    # Correct - distribute error
    lead_I_corrected = lead_I + alpha * error
    lead_III_corrected = lead_III + alpha * error
    
    # Update rows (only segment 0 for leads I and III)
    pred_mv_rows[0][:segment_width] = lead_I_corrected
    pred_mv_rows[2][:segment_width] = lead_III_corrected
    
    return pred_mv_rows


def clamp_ecg_amplitude(signal_mv):
    """Clamp signal to reasonable ECG range."""
    return np.clip(signal_mv, ECG_MV_MIN, ECG_MV_MAX)


def interpolate_nan(signal_1d):
    """Interpolate NaN values from valid neighbors. Falls back to 0 if all NaN."""
    valid_mask = np.isfinite(signal_1d)
    if valid_mask.all():
        return signal_1d
    if not valid_mask.any():
        return np.zeros_like(signal_1d)
    
    x = np.arange(len(signal_1d))
    signal_1d[~valid_mask] = np.interp(x[~valid_mask], x[valid_mask], signal_1d[valid_mask])
    return signal_1d


# =============================================================================
# Inference and Visualization
# =============================================================================
def load_model(checkpoint_path, device):
    """Load trained V15 model."""
    model = PerLeadNet(encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=False)
    
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    model.load_state_dict(checkpoint['model'])
    model = model.to(device)
    model.eval()
    
    print(f"Loaded checkpoint from epoch {checkpoint['epoch']}")
    print(f"  SNR: {checkpoint.get('snr', 'N/A'):.2f} dB")
    print(f"  MAE: {checkpoint.get('mae', 'N/A'):.2f} px")
    
    return model


def crop_row(image, row_idx):
    """Crop a single row centered on its baseline, signal region only (T0:T1)."""
    baseline_y = int(ZERO_MV[row_idx])
    y_start = max(0, baseline_y - CROP_HALF_HEIGHT)
    y_end = min(TARGET_HEIGHT, baseline_y + CROP_HALF_HEIGHT)
    
    # V16: Crop x to signal region only (T0:T1)
    row_crop = image[y_start:y_end, T0:T1, :].copy()
    
    # Pad if necessary
    if row_crop.shape[0] < ROW_HEIGHT:
        pad_top = max(0, CROP_HALF_HEIGHT - baseline_y)
        pad_bottom = max(0, (baseline_y + CROP_HALF_HEIGHT) - TARGET_HEIGHT)
        row_crop = np.pad(row_crop, ((pad_top, pad_bottom), (0, 0), (0, 0)), mode='edge')
    
    return row_crop


def predict_row(model, row_crop, device):
    """Run inference on a single row crop."""
    # Prepare input
    image_tensor = torch.from_numpy(row_crop.astype(np.float32) / 255.0)
    image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0).to(device)
    
    with torch.no_grad():
        with torch.cuda.amp.autocast():
            output = model(image_tensor)
    
    # Convert from normalized [0, 1] to crop-relative pixels
    pred_y_crop = output[0].cpu().numpy() * ROW_HEIGHT
    
    return pred_y_crop


def convert_crop_to_full(pred_y_crop, row_idx):
    """
    Convert crop-relative y-coordinates to full image coordinates.
    V16: Output is OUTPUT_WIDTH (3926) values, maps to x positions T0:T1.
    """
    baseline_y = int(ZERO_MV[row_idx])
    y_start = max(0, baseline_y - CROP_HALF_HEIGHT)
    
    # Adjust for padding
    pad_top = max(0, CROP_HALF_HEIGHT - baseline_y)
    
    # Convert to full image y
    pred_y_full = pred_y_crop - pad_top + y_start
    
    return pred_y_full


def visualize_predictions(image, predictions, output_path):
    """
    Draw predictions as dots on the image.
    
    predictions: list of 4 arrays, one per row, each [OUTPUT_WIDTH] in full image y-coords
    V16: predictions are already OUTPUT_WIDTH (3926) values for x positions 0 to OUTPUT_WIDTH-1.
    Output is cropped to signal region.
    Plots EVERY predicted point.
    """
    # Crop image to signal region (x: T0 to T1)
    vis_image = image[:, T0:T1, :].copy()
    
    # Colors for each row (BGR)
    colors = [
        (255, 0, 0),    # Blue for row 0
        (0, 255, 0),    # Green for row 1
        (255, 0, 255),  # Magenta for row 2
        (0, 165, 255),  # Orange for row 3 (rhythm)
    ]
    
    for row_idx, pred_y in enumerate(predictions):
        color = colors[row_idx]
        
        # V16: pred_y has OUTPUT_WIDTH values, directly maps to vis_image x coords
        for x in range(len(pred_y)):
            y = int(np.clip(pred_y[x], 0, TARGET_HEIGHT - 1))
            # Use radius 1 for dense dots
            cv2.circle(vis_image, (x, y), 1, color, -1)
    
    cv2.imwrite(str(output_path), vis_image)


def compute_snr_per_lead(pred_y_full, df, row_idx, epsilon=1e-10):
    """
    Compute SNR in dB for each lead in a row.
    
    IMPORTANT: Resample predictions UP to match GT sample rate (typically 500 Hz).
    Kaggle evaluates SNR at the original sample rate, not pixel resolution.
    
    pred_y_full: [OUTPUT_WIDTH] predictions in full image y-coords
    df: DataFrame with ground truth signals in mV
    row_idx: 0-2 for short leads (4 segments), 3 for rhythm strip
    
    Returns dict of lead_name -> SNR in dB
    """
    baseline_y = ZERO_MV[row_idx]
    segment_width = OUTPUT_WIDTH // 4
    
    lead_snrs = {}
    
    if row_idx < 3:
        # Rows 0-2: 4 leads each
        lead_names = LEAD_LAYOUT[row_idx]
        for seg_idx, lead_name in enumerate(lead_names):
            if lead_name not in df.columns:
                lead_snrs[lead_name] = None
                continue
            
            # Get original GT signal in mV
            gt_mv = df[lead_name].dropna().values
            if len(gt_mv) == 0:
                lead_snrs[lead_name] = None
                continue
            
            # Special handling for Lead II (10s data in GT)
            # Lead II short strip is in row 1, segment 0 (first column)
            # Training uses first 25% of GT (0-2.5s), so inference must match
            # Lead II GT is typically 4x longer than other short leads (10s vs 2.5s)
            if lead_name == 'II':
                # Lead II in CSV contains 10s rhythm strip data
                # For short strip (2.5s), use first quarter
                # Compare with other short lead lengths to detect 10s data
                ref_len = len(df['I'].dropna().values) if 'I' in df.columns else len(gt_mv) // 4
                if len(gt_mv) > ref_len * 2:  # Lead II is at least 2x longer = it's 10s data
                    quarter_len = len(gt_mv) // 4
                    gt_mv = gt_mv[:quarter_len]  # First quarter (0-2.5s) - matches training
            
            # Get prediction segment in pixels
            seg_start = seg_idx * segment_width
            seg_end = (seg_idx + 1) * segment_width
            pred_y_seg = pred_y_full[seg_start:seg_end]
            
            # Convert prediction from pixels to mV
            pred_mv_pixels = (baseline_y - pred_y_seg) / MV_TO_PIXEL
            
            # Resample prediction UP to match GT length
            x_pred = np.linspace(0, 1, len(pred_mv_pixels))
            x_gt = np.linspace(0, 1, len(gt_mv))
            pred_mv_resampled = np.interp(x_gt, x_pred, pred_mv_pixels)
            
            # Compute SNR
            signal_power = (gt_mv ** 2).mean()
            noise_power = ((pred_mv_resampled - gt_mv) ** 2).mean()
            
            if noise_power < epsilon:
                lead_snrs[lead_name] = 50.0
            else:
                snr = 10 * np.log10(signal_power / (noise_power + epsilon))
                lead_snrs[lead_name] = float(snr)
    else:
        # Row 3: Full Lead II rhythm strip (10s)
        if 'II' not in df.columns:
            lead_snrs['II_rhythm'] = None
            return lead_snrs
        
        gt_mv = df['II'].dropna().values
        if len(gt_mv) == 0:
            lead_snrs['II_rhythm'] = None
            return lead_snrs
        
        # Convert full prediction from pixels to mV
        pred_mv_pixels = (baseline_y - pred_y_full) / MV_TO_PIXEL
        
        # Resample prediction UP to match GT length
        x_pred = np.linspace(0, 1, len(pred_mv_pixels))
        x_gt = np.linspace(0, 1, len(gt_mv))
        pred_mv_resampled = np.interp(x_gt, x_pred, pred_mv_pixels)
        
        # Compute SNR
        signal_power = (gt_mv ** 2).mean()
        noise_power = ((pred_mv_resampled - gt_mv) ** 2).mean()
        
        if noise_power < epsilon:
            lead_snrs['II_rhythm'] = 50.0
        else:
            snr = 10 * np.log10(signal_power / (noise_power + epsilon))
            lead_snrs['II_rhythm'] = float(snr)
    
    return lead_snrs


def process_image(model, image_path, csv_path, output_dir, device, 
                  apply_smoothing=True, apply_einthoven=True, negative_dir=None):
    """Process a single image, save visualization, and compute per-lead SNR.
    
    Args:
        apply_smoothing: Apply Savitzky-Golay smoothing to predictions
        apply_einthoven: Apply Einthoven's law correction
        negative_dir: If provided, save low-SNR predictions here
    
    Returns:
        Tuple of (all_lead_snrs dict, min_snr value, output_path)
    """
    # Load image
    image = cv2.imread(str(image_path), cv2.IMREAD_COLOR)
    if image is None:
        print(f"Failed to load: {image_path}")
        return None, None, None
    
    # Preprocess
    image = image[Y0:Y1, X0:X1]
    image = cv2.resize(image, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR)
    
    # Load ground truth
    df = pd.read_csv(csv_path)
    
    # Predict each row
    predictions = []
    pred_mv_rows = {}  # Store mV predictions for Einthoven correction
    
    for row_idx in range(4):
        row_crop = crop_row(image, row_idx)
        pred_y_crop = predict_row(model, row_crop, device)
        pred_y_full = convert_crop_to_full(pred_y_crop, row_idx)
        
        # Convert to mV for post-processing
        baseline_y = ZERO_MV[row_idx]
        pred_mv = (baseline_y - pred_y_full) / MV_TO_PIXEL
        
        # Apply smoothing
        if apply_smoothing:
            pred_mv = apply_savgol_smoothing(pred_mv, window=7, polyorder=2)
        
        # Clamp to reasonable range
        pred_mv = clamp_ecg_amplitude(pred_mv)
        
        # Handle any NaN values
        pred_mv = interpolate_nan(pred_mv.copy())
        
        pred_mv_rows[row_idx] = pred_mv
        predictions.append(pred_y_full)
    
    # Apply Einthoven correction (only for short leads in rows 0-2)
    if apply_einthoven:
        pred_mv_rows = apply_einthoven_correction(pred_mv_rows, alpha=0.33)
    
    # Convert corrected mV back to pixel coordinates for visualization
    corrected_predictions = []
    for row_idx in range(4):
        baseline_y = ZERO_MV[row_idx]
        pred_y_corrected = baseline_y - pred_mv_rows[row_idx] * MV_TO_PIXEL
        corrected_predictions.append(pred_y_corrected)
    
    # Compute SNR using corrected predictions
    all_lead_snrs = {}
    for row_idx in range(4):
        lead_snrs = compute_snr_per_lead(corrected_predictions[row_idx], df, row_idx)
        all_lead_snrs.update(lead_snrs)
    
    # Find minimum SNR for this image
    valid_snrs = [v for v in all_lead_snrs.values() if v is not None]
    min_snr = min(valid_snrs) if valid_snrs else 0.0
    
    # Visualize
    sample_id = image_path.parent.name
    variant = image_path.stem.split('-')[-1] if '-' in image_path.stem else '0000'
    output_path = output_dir / f"{sample_id}_{variant}.png"
    
    visualize_predictions(image, corrected_predictions, output_path)
    
    # Save to negative folder if SNR is low
    if negative_dir is not None and min_snr < LOW_SNR_THRESHOLD:
        neg_output_path = negative_dir / f"{sample_id}_{variant}_snr{min_snr:.1f}.png"
        visualize_predictions(image, corrected_predictions, neg_output_path)
    
    return all_lead_snrs, min_snr, output_path


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--checkpoint', type=str, 
                        default='/data/ecg-digitization/checkpoints/v16_perlead_latest.pth')
    parser.add_argument('--kaggle_data', type=str, 
                        default='/data/ecg-digitization/stage1_data/train')
    parser.add_argument('--output_dir', type=str, 
                        default=os.path.expanduser('~/tmp/pred/v16'))
    parser.add_argument('--num_samples', type=int, default=None,
                        help='Number of samples (default: all holdout samples)')
    parser.add_argument('--use_holdout', action='store_true', default=True,
                        help='Use holdout/validation set instead of random samples')
    parser.add_argument('--no_smoothing', action='store_true',
                        help='Disable Savitzky-Golay smoothing')
    parser.add_argument('--no_einthoven', action='store_true',
                        help='Disable Einthoven law correction')
    parser.add_argument('--seed', type=int, default=42)
    args = parser.parse_args()
    
    # Setup
    random.seed(args.seed)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    output_dir = Path(args.output_dir)
    negative_dir = Path(os.path.expanduser('~/tmp/pred/v16/negpreds'))
    
    # Delete old predictions
    if output_dir.exists():
        import shutil
        shutil.rmtree(output_dir)
        print(f"Deleted old predictions in {output_dir}")
    output_dir.mkdir(parents=True, exist_ok=True)
    negative_dir.mkdir(parents=True, exist_ok=True)
    
    print(f"{'='*70}")
    print(f"V16 Inference Visualization (with enhancements)")
    print(f"{'='*70}")
    print(f"Checkpoint: {args.checkpoint}")
    print(f"Output: {output_dir}")
    print(f"Negative predictions: {negative_dir}")
    print(f"Device: {device}")
    print(f"Smoothing: {'OFF' if args.no_smoothing else 'ON (Savgol w=7)'}")
    print(f"Einthoven correction: {'OFF' if args.no_einthoven else 'ON (alpha=0.33)'}")
    print(f"Low SNR threshold: {LOW_SNR_THRESHOLD} dB")
    print(f"{'='*70}")
    
    # Load model
    model = load_model(args.checkpoint, device)
    
    # Find all valid images with their CSV files
    kaggle_dir = Path(args.kaggle_data)
    val_sample_set = set(VAL_SAMPLE_IDS)
    all_samples = []
    
    for sample_dir in kaggle_dir.iterdir():
        if not sample_dir.is_dir():
            continue
        
        # Filter by holdout set if requested
        if args.use_holdout and sample_dir.name not in val_sample_set:
            continue
        
        # Find CSV file
        csv_files = list(sample_dir.glob('*.csv'))
        if len(csv_files) != 1:
            continue
        csv_path = csv_files[0]
        
        for img_path in sample_dir.glob('*.png'):
            variant = img_path.stem.split('-')[-1] if '-' in img_path.stem else '0000'
            if variant in VALID_VARIANTS:
                all_samples.append((img_path, csv_path))
    
    set_type = "holdout" if args.use_holdout else "all"
    print(f"Found {len(all_samples)} valid images in {set_type} set")
    
    # Select samples (holdout = all by default, random otherwise)
    if args.num_samples is not None and len(all_samples) > args.num_samples:
        selected = random.sample(all_samples, args.num_samples)
    else:
        selected = all_samples
    
    print(f"Processing {len(selected)} images...")
    
    # Collect per-lead SNRs and track low-SNR samples
    all_snrs = {lead: [] for lead in ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 
                                       'V1', 'V2', 'V3', 'V4', 'V5', 'V6', 'II_rhythm']}
    low_snr_samples = []  # Track samples with low SNR
    
    # Process each image
    for img_path, csv_path in tqdm(selected):
        try:
            lead_snrs, min_snr, output_path = process_image(
                model, img_path, csv_path, output_dir, device,
                apply_smoothing=not args.no_smoothing,
                apply_einthoven=not args.no_einthoven,
                negative_dir=negative_dir
            )
            if lead_snrs:
                for lead, snr in lead_snrs.items():
                    if snr is not None:
                        all_snrs[lead].append(snr)
                
                # Track low SNR samples
                if min_snr is not None and min_snr < LOW_SNR_THRESHOLD:
                    # Find which lead has the lowest SNR
                    worst_lead = min(lead_snrs, key=lambda k: lead_snrs[k] if lead_snrs[k] is not None else float('inf'))
                    low_snr_samples.append({
                        'sample': str(img_path.parent.name),
                        'variant': img_path.stem.split('-')[-1] if '-' in img_path.stem else '0000',
                        'min_snr': min_snr,
                        'worst_lead': worst_lead
                    })
        except Exception as e:
            print(f"Error processing {img_path}: {e}")
    
    # Print per-lead SNR statistics
    print(f"\n{'='*70}")
    print(f"Per-Lead SNR Statistics (dB)")
    print(f"{'='*70}")
    print(f"{'Lead':<12} {'Mean':>8} {'Std':>8} {'Min':>8} {'Max':>8} {'Count':>6}")
    print(f"{'-'*70}")
    
    total_snrs = []
    for lead in ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6', 'II_rhythm']:
        snrs = all_snrs[lead]
        if len(snrs) > 0:
            mean_snr = np.mean(snrs)
            std_snr = np.std(snrs)
            min_snr = np.min(snrs)
            max_snr = np.max(snrs)
            print(f"{lead:<12} {mean_snr:>8.2f} {std_snr:>8.2f} {min_snr:>8.2f} {max_snr:>8.2f} {len(snrs):>6}")
            total_snrs.extend(snrs)
        else:
            print(f"{lead:<12} {'N/A':>8} {'N/A':>8} {'N/A':>8} {'N/A':>8} {0:>6}")
    
    print(f"{'-'*70}")
    if len(total_snrs) > 0:
        print(f"{'OVERALL':<12} {np.mean(total_snrs):>8.2f} {np.std(total_snrs):>8.2f} "
              f"{np.min(total_snrs):>8.2f} {np.max(total_snrs):>8.2f} {len(total_snrs):>6}")
    print(f"{'='*70}")
    
    # Save SNR report to file
    snr_report_path = output_dir / 'snr_report.csv'
    with open(snr_report_path, 'w') as f:
        f.write("Lead,Mean_SNR,Std_SNR,Min_SNR,Max_SNR,Count\n")
        for lead in ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6', 'II_rhythm']:
            snrs = all_snrs[lead]
            if len(snrs) > 0:
                f.write(f"{lead},{np.mean(snrs):.2f},{np.std(snrs):.2f},{np.min(snrs):.2f},{np.max(snrs):.2f},{len(snrs)}\n")
            else:
                f.write(f"{lead},N/A,N/A,N/A,N/A,0\n")
        if len(total_snrs) > 0:
            f.write(f"OVERALL,{np.mean(total_snrs):.2f},{np.std(total_snrs):.2f},{np.min(total_snrs):.2f},{np.max(total_snrs):.2f},{len(total_snrs)}\n")
    
    print(f"\nSNR report saved to: {snr_report_path}")
    print(f"Done! Saved {len(selected)} visualizations to {output_dir}")
    
    # Print low-SNR summary
    if low_snr_samples:
        print(f"\n{'='*70}")
        print(f"Low SNR Samples (< {LOW_SNR_THRESHOLD} dB) - Saved to {negative_dir}")
        print(f"{'='*70}")
        print(f"{'Sample':<15} {'Variant':<8} {'Min SNR':>10} {'Worst Lead':<12}")
        print(f"{'-'*70}")
        
        # Sort by SNR (worst first)
        low_snr_samples.sort(key=lambda x: x['min_snr'])
        
        for item in low_snr_samples[:20]:  # Show top 20 worst
            print(f"{item['sample']:<15} {item['variant']:<8} {item['min_snr']:>10.2f} {item['worst_lead']:<12}")
        
        if len(low_snr_samples) > 20:
            print(f"... and {len(low_snr_samples) - 20} more")
        
        print(f"\nTotal low-SNR samples: {len(low_snr_samples)}")
        
        # Save low-SNR report
        low_snr_report_path = negative_dir / 'low_snr_report.csv'
        with open(low_snr_report_path, 'w') as f:
            f.write("sample,variant,min_snr,worst_lead\n")
            for item in low_snr_samples:
                f.write(f"{item['sample']},{item['variant']},{item['min_snr']:.2f},{item['worst_lead']}\n")
        print(f"Low-SNR report saved to: {low_snr_report_path}")
    else:
        print(f"\nNo samples with SNR < {LOW_SNR_THRESHOLD} dB")
    
    print(f"{'='*70}")


if __name__ == '__main__':
    main()