code/scripts/infer_v16_1_visualize.py

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

Runs V16.1 per-lead model (2x upscaled) on Kaggle images and plots predictions.
Automatically SCPs the latest checkpoint from remote training VM before inference.

Key differences from V16:
- Input upscaled 2x: 1000 × 7852 (from 500 × 3926)
- Lanczos4 interpolation for crisp lines
- Predictions are downsampled back to base resolution for visualization/SNR
"""

import os
import sys
import argparse
import random
import subprocess
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_1_perlead.py - 2x upscaled)
# =============================================================================
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_BASE = T1 - T0  # 3926 - base signal width

# V16.1: 2x upscaling
UPSCALE_FACTOR = 2
OUTPUT_WIDTH = OUTPUT_WIDTH_BASE * UPSCALE_FACTOR  # 7852

# Per-row crop parameters (at base resolution, before upscale)
CROP_HALF_HEIGHT_BASE = 250
ROW_HEIGHT_BASE = 500

# After 2x upscale
CROP_HALF_HEIGHT = CROP_HALF_HEIGHT_BASE * UPSCALE_FACTOR  # 500
ROW_HEIGHT = ROW_HEIGHT_BASE * UPSCALE_FACTOR  # 1000

# Model input dimensions (after 2x upscale)
INPUT_HEIGHT = ROW_HEIGHT  # 1000
INPUT_WIDTH = OUTPUT_WIDTH  # 7852

# MV_TO_PIXEL also scales 2x
MV_TO_PIXEL_UPSCALED = MV_TO_PIXEL * UPSCALE_FACTOR  # 157

# ECG amplitude limits (mV)
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)
# 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
LEAD_LAYOUT = [
    ['I', 'aVR', 'V1', 'V4'],
    ['II', 'aVL', 'V2', 'V5'],
    ['III', 'aVF', 'V3', 'V6'],
]

# Remote VM configuration for SCP
REMOTE_HOST = os.environ.get('REMOTE_TRAIN_HOST', 'azureuser@172.212.222.231')
REMOTE_CHECKPOINT_DIR = '/data/ecg-digitization/checkpoints'
REMOTE_CHECKPOINT_PATTERN = 'v16_1_perlead_best_snr'  # Fetch best SNR checkpoint


# =============================================================================
# SCP Checkpoint from Remote VM
# =============================================================================
def scp_latest_checkpoint(remote_host, remote_dir, pattern, local_dir):
    """
    SCP the latest checkpoint matching pattern from remote VM.
    
    Args:
        remote_host: SSH host (can be alias from ~/.ssh/config or user@ip)
        remote_dir: Remote directory containing checkpoints
        pattern: Pattern to match checkpoint files
        local_dir: Local directory to save checkpoint
    
    Returns:
        Path to local checkpoint file, or None if failed
    """
    local_dir = Path(local_dir)
    local_dir.mkdir(parents=True, exist_ok=True)
    
    print(f"{'='*70}")
    print(f"Fetching latest checkpoint from remote VM...")
    print(f"  Remote: {remote_host}:{remote_dir}")
    print(f"  Pattern: {pattern}")
    print(f"{'='*70}")
    
    # Find latest checkpoint on remote
    # List files matching pattern and sort by modification time
    find_cmd = f"ssh {remote_host} 'ls -t {remote_dir}/{pattern}*.pth 2>/dev/null | head -1'"
    
    try:
        result = subprocess.run(find_cmd, shell=True, capture_output=True, text=True, timeout=30)
        if result.returncode != 0 or not result.stdout.strip():
            print(f"ERROR: Could not find checkpoint matching '{pattern}' on remote")
            print(f"  stderr: {result.stderr}")
            return None
        
        remote_path = result.stdout.strip()
        filename = os.path.basename(remote_path)
        local_path = local_dir / filename
        
        print(f"  Found: {remote_path}")
        
        # Check if we already have this file (by size comparison)
        if local_path.exists():
            # Get remote file size
            size_cmd = f"ssh {remote_host} 'stat -c %s {remote_path}'"
            size_result = subprocess.run(size_cmd, shell=True, capture_output=True, text=True, timeout=10)
            if size_result.returncode == 0:
                remote_size = int(size_result.stdout.strip())
                local_size = local_path.stat().st_size
                if remote_size == local_size:
                    print(f"  Local copy already up-to-date: {local_path}")
                    return local_path
        
        # SCP the file
        print(f"  Downloading to: {local_path}")
        scp_cmd = f"scp {remote_host}:{remote_path} {local_path}"
        result = subprocess.run(scp_cmd, shell=True, capture_output=True, text=True, timeout=300)
        
        if result.returncode != 0:
            print(f"ERROR: SCP failed")
            print(f"  stderr: {result.stderr}")
            return None
        
        print(f"  ✓ Downloaded successfully ({local_path.stat().st_size / 1024 / 1024:.1f} MB)")
        return local_path
        
    except subprocess.TimeoutExpired:
        print("ERROR: SSH/SCP command timed out")
        return None
    except Exception as e:
        print(f"ERROR: {e}")
        return None


# =============================================================================
# Model Architecture (from train_v16_1_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):
    """
    Per-lead regression network for 2x upscaled input.
    
    V16.1 Input: (B, 3, 1000, 7852)
    Output: (B, 7852) - y-coordinate for each x position [0, 1]
    """
    
    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
# =============================================================================
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)
    """
    segment_width = len(pred_mv_rows[0]) // 4
    
    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()
    
    derived_II = lead_I + lead_III
    error = lead_II_short - derived_II
    
    lead_I_corrected = lead_I + alpha * error
    lead_III_corrected = lead_III + alpha * error
    
    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."""
    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 V16.1 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 (base resolution)")
    
    return model


def crop_row(image, row_idx):
    """
    Crop a single row centered on its baseline, signal region only (T0:T1).
    Returns 2x upscaled crop using Lanczos interpolation.
    """
    baseline_y = int(ZERO_MV[row_idx])
    y_start = max(0, baseline_y - CROP_HALF_HEIGHT_BASE)
    y_end = min(TARGET_HEIGHT, baseline_y + CROP_HALF_HEIGHT_BASE)
    
    # Crop x to signal region (T0:T1) at base resolution
    row_crop = image[y_start:y_end, T0:T1, :].copy()
    
    # Pad if necessary (at base resolution)
    if row_crop.shape[0] < ROW_HEIGHT_BASE:
        pad_top = max(0, CROP_HALF_HEIGHT_BASE - baseline_y)
        pad_bottom = max(0, (baseline_y + CROP_HALF_HEIGHT_BASE) - TARGET_HEIGHT)
        row_crop = np.pad(row_crop, ((pad_top, pad_bottom), (0, 0), (0, 0)), mode='edge')
    
    # 2x Upscale with Lanczos for crisp lines
    row_crop = cv2.resize(row_crop, (OUTPUT_WIDTH, ROW_HEIGHT), 
                          interpolation=cv2.INTER_LANCZOS4)
    
    return row_crop


def predict_row(model, row_crop, device):
    """Run inference on a single 2x upscaled 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 2x crop-relative pixels
    pred_y_crop_2x = output[0].cpu().numpy() * ROW_HEIGHT
    
    return pred_y_crop_2x


def convert_crop_to_full(pred_y_crop_2x, row_idx):
    """
    Convert 2x crop-relative y-coordinates to base resolution full image coordinates.
    
    pred_y_crop_2x: [OUTPUT_WIDTH] predictions at 2x resolution
    Returns: [OUTPUT_WIDTH_BASE] predictions in full image y-coords (base resolution)
    """
    baseline_y = int(ZERO_MV[row_idx])
    y_start = max(0, baseline_y - CROP_HALF_HEIGHT_BASE)
    
    # Adjust for padding
    pad_top = max(0, CROP_HALF_HEIGHT_BASE - baseline_y)
    
    # Convert from 2x crop-relative to 2x full image y
    pred_y_full_2x = (pred_y_crop_2x - pad_top * UPSCALE_FACTOR) / UPSCALE_FACTOR + y_start
    
    # Downsample from 2x resolution to base resolution
    x_2x = np.linspace(0, 1, len(pred_y_full_2x))
    x_base = np.linspace(0, 1, OUTPUT_WIDTH_BASE)
    pred_y_full_base = np.interp(x_base, x_2x, pred_y_full_2x)
    
    return pred_y_full_base


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_BASE] in full image y-coords
    Output is cropped to signal region.
    """
    # 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]
        
        for x in range(len(pred_y)):
            y = int(np.clip(pred_y[x], 0, TARGET_HEIGHT - 1))
            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.
    
    pred_y_full: [OUTPUT_WIDTH_BASE] predictions in full image y-coords (base resolution)
    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_BASE // 4
    
    lead_snrs = {}
    
    if row_idx < 3:
        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
            
            gt_mv = df[lead_name].dropna().values
            if len(gt_mv) == 0:
                lead_snrs[lead_name] = None
                continue
            
            # Lead II handling (10s data in GT)
            if lead_name == 'II':
                ref_len = len(df['I'].dropna().values) if 'I' in df.columns else len(gt_mv) // 4
                if len(gt_mv) > ref_len * 2:
                    quarter_len = len(gt_mv) // 4
                    gt_mv = gt_mv[:quarter_len]
            
            seg_start = seg_idx * segment_width
            seg_end = (seg_idx + 1) * segment_width
            pred_y_seg = pred_y_full[seg_start:seg_end]
            
            pred_mv_pixels = (baseline_y - pred_y_seg) / MV_TO_PIXEL
            
            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)
            
            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:
        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
        
        pred_mv_pixels = (baseline_y - pred_y_full) / MV_TO_PIXEL
        
        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)
        
        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."""
    # 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 (with 2x upscaling internally)
    predictions = []
    pred_mv_rows = {}
    
    for row_idx in range(4):
        row_crop = crop_row(image, row_idx)  # Returns 2x upscaled crop
        pred_y_crop_2x = predict_row(model, row_crop, device)  # Predict at 2x
        pred_y_full = convert_crop_to_full(pred_y_crop_2x, row_idx)  # Convert to base resolution
        
        # 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)
        
        pred_mv = clamp_ecg_amplitude(pred_mv)
        pred_mv = interpolate_nan(pred_mv.copy())
        
        pred_mv_rows[row_idx] = pred_mv
        predictions.append(pred_y_full)
    
    # Apply Einthoven correction
    if apply_einthoven:
        pred_mv_rows = apply_einthoven_correction(pred_mv_rows, alpha=0.33)
    
    # Convert corrected mV back to pixel coordinates
    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
    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)
    
    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)
    
    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=None,
                        help='Local checkpoint path. If not provided, will SCP from remote.')
    parser.add_argument('--remote_host', type=str, default=REMOTE_HOST,
                        help='Remote SSH host for SCP (default: gpu-vm)')
    parser.add_argument('--remote_dir', type=str, default=REMOTE_CHECKPOINT_DIR,
                        help='Remote checkpoint directory')
    parser.add_argument('--local_cache', type=str, 
                        default='/tmp/v16_1_checkpoints',
                        help='Local directory to cache downloaded checkpoints')
    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_1'))
    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('--no_scp', action='store_true',
                        help='Skip SCP, use local checkpoint only')
    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_1/negpreds'))
    
    # Get checkpoint
    checkpoint_path = None
    
    if args.checkpoint:
        checkpoint_path = Path(args.checkpoint)
        if not checkpoint_path.exists():
            print(f"ERROR: Checkpoint not found: {checkpoint_path}")
            return
    elif not args.no_scp:
        # SCP latest from remote
        checkpoint_path = scp_latest_checkpoint(
            args.remote_host,
            args.remote_dir,
            REMOTE_CHECKPOINT_PATTERN,
            args.local_cache
        )
        if checkpoint_path is None:
            print("ERROR: Failed to fetch checkpoint from remote")
            return
    else:
        # Look for local checkpoint
        local_cache = Path(args.local_cache)
        if local_cache.exists():
            checkpoints = sorted(local_cache.glob(f'{REMOTE_CHECKPOINT_PATTERN}*.pth'))
            if checkpoints:
                checkpoint_path = checkpoints[-1]
                print(f"Using local checkpoint: {checkpoint_path}")
        
        if checkpoint_path is None:
            print("ERROR: No checkpoint found. Provide --checkpoint or remove --no_scp")
            return
    
    # 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.1 Inference Visualization (2x Upscaled)")
    print(f"{'='*70}")
    print(f"Checkpoint: {checkpoint_path}")
    print(f"Output: {output_dir}")
    print(f"Negative predictions: {negative_dir}")
    print(f"Device: {device}")
    print(f"Input resolution: {INPUT_HEIGHT} x {INPUT_WIDTH} (2x upscaled)")
    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(checkpoint_path, 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
        
        if args.use_holdout and sample_dir.name not in val_sample_set:
            continue
        
        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")
    
    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
    all_snrs = {lead: [] for lead in ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 
                                       'V1', 'V2', 'V3', 'V4', 'V5', 'V6', 'II_rhythm']}
    low_snr_samples = []
    
    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)
                
                if min_snr is not None and min_snr < LOW_SNR_THRESHOLD:
                    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
    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}")
        
        low_snr_samples.sort(key=lambda x: x['min_snr'])
        
        for item in low_snr_samples[:20]:
            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)}")
        
        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()