#!/usr/bin/env python3
"""
V18 Inference with aVR/aVL/aVF Baseline Fix

Runs V18 (V16 frozen + Refiner with Cross-Row Attention) on Kaggle images.
Includes hardcoded baseline correction for better aVR/aVL/aVF accuracy.
"""

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, find_peaks
from scipy.ndimage import gaussian_filter1d

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

import cv2
import timm


# =============================================================================
# Constants
# =============================================================================
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

# Per-row crop parameters
CROP_HALF_HEIGHT = 250
ROW_HEIGHT = 500

# Model input dimensions
INPUT_HEIGHT = ROW_HEIGHT  # 500
INPUT_WIDTH = OUTPUT_WIDTH  # 3926

# 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'],
]

# Hardcoded baseline offsets (mV) - computed from 977 images with V18 model
# These are median(prediction - ground_truth) values
# V18 has essentially zero bias (~0.0001 mV), much smaller than V20 (~0.007 mV)
BASELINE_OFFSETS = {
    'I': 0.0001, 'II': 0.0010, 'III': 0.0006,
    'aVR': -0.0003, 'aVL': -0.0001, 'aVF': 0.0002,
    'V1': -0.0002, 'V2': -0.0002, 'V3': 0.0004,
    'V4': -0.0005, 'V5': 0.0001, 'V6': 0.0002,
}

# Leads where V18 helps (positive delta from experiments)
# Based on results: I +1.20, III +1.61, aVR +0.67, aVF +0.97, V1 +0.55, V6 +0.47, II_rhythm +0.74
# Leads V18 hurts: II -0.31, aVL -0.37, V2 -0.42, V3 0.00, V4 -0.09, V5 -0.29
V18_BENEFICIAL_LEADS = {'I', 'III', 'aVR', 'aVF', 'V1', 'V6', 'II_rhythm'}

# Remote VM configuration
REMOTE_HOST = os.environ.get('REMOTE_TRAIN_HOST', 'azureuser@172.212.222.231')
REMOTE_CHECKPOINT_DIR = '/data/ecg-digitization/checkpoints'
REMOTE_V16_CHECKPOINT = 'v16_perlead_best_snr.pth'
REMOTE_V18_CHECKPOINT_PATTERN = 'v18_refiner'


# =============================================================================
# SCP Checkpoint from Remote VM
# =============================================================================
def scp_checkpoint(remote_host, remote_path, local_path):
    """SCP a single checkpoint file from remote VM."""
    local_path = Path(local_path)
    local_path.parent.mkdir(parents=True, exist_ok=True)
    
    print(f"  Fetching: {remote_host}:{remote_path}")
    
    try:
        # Check if we already have this file (by size comparison)
        if local_path.exists():
            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 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 - {result.stderr}")
            return None
        
        print(f"  ✓ Downloaded ({local_path.stat().st_size / 1024 / 1024:.1f} MB)")
        return local_path
        
    except Exception as e:
        print(f"ERROR: {e}")
        return None


def scp_latest_checkpoint(remote_host, remote_dir, pattern, local_dir):
    """SCP the latest checkpoint matching pattern from remote VM."""
    local_dir = Path(local_dir)
    local_dir.mkdir(parents=True, exist_ok=True)
    
    # Find latest checkpoint on remote
    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"  No checkpoint found matching '{pattern}'")
            return None
        
        remote_path = result.stdout.strip()
        filename = os.path.basename(remote_path)
        local_path = local_dir / filename
        
        return scp_checkpoint(remote_host, remote_path, local_path)
        
    except Exception as e:
        print(f"ERROR: {e}")
        return None


# =============================================================================
# V16 Model Architecture (frozen prior)
# =============================================================================
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 PerLeadNetV16(nn.Module):
    """V16 Per-Lead ECG Network (frozen, used as prior)."""
    
    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)


# =============================================================================
# V18 Refiner Architecture
# =============================================================================
class CrossRowAttention(nn.Module):
    """Cross-Row Multi-Head Self-Attention."""
    def __init__(self, embed_dim, num_heads=8, dropout=0.1):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scale = self.head_dim ** -0.5
        
        self.qkv = nn.Linear(embed_dim, embed_dim * 3)
        self.proj = nn.Linear(embed_dim, embed_dim)
        self.dropout = nn.Dropout(dropout)
        self.norm = nn.LayerNorm(embed_dim)
    
    def forward(self, x):
        B, num_rows, C, W = x.shape
        x = x.permute(0, 3, 1, 2).reshape(B * W, num_rows, C)
        
        residual = x
        x = self.norm(x)
        
        qkv = self.qkv(x).reshape(B * W, num_rows, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.dropout(attn)
        
        out = (attn @ v).transpose(1, 2).reshape(B * W, num_rows, C)
        out = self.proj(out)
        out = self.dropout(out)
        out = out + residual
        
        out = out.reshape(B, W, num_rows, C).permute(0, 2, 3, 1)
        return out


class CrossRowTransformerBlock(nn.Module):
    """Full transformer block with cross-row attention and FFN."""
    def __init__(self, embed_dim, num_heads=8, mlp_ratio=4.0, dropout=0.1):
        super().__init__()
        
        self.attn = CrossRowAttention(embed_dim, num_heads, dropout)
        
        self.norm = nn.LayerNorm(embed_dim)
        hidden_dim = int(embed_dim * mlp_ratio)
        self.ffn = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, embed_dim),
            nn.Dropout(dropout),
        )
    
    def forward(self, x):
        x = self.attn(x)
        
        B, num_rows, C, W = x.shape
        residual = x
        x = x.permute(0, 3, 1, 2).reshape(B * W, num_rows, C)
        x = self.norm(x)
        x = self.ffn(x) + residual.permute(0, 3, 1, 2).reshape(B * W, num_rows, C)
        x = x.reshape(B, W, num_rows, C).permute(0, 2, 3, 1)
        
        return x


class RefinerEncoder(nn.Module):
    """Lightweight encoder for refinement (4 input channels: RGB + guide)."""
    def __init__(self, encoder_name='efficientnet_b0', pretrained=True):
        super().__init__()
        
        self.encoder = timm.create_model(
            encoder_name,
            pretrained=pretrained,
            features_only=True,
            out_indices=(1, 2, 3),
            in_chans=4,
        )
        self.channels = self.encoder.feature_info.channels()
    
    def forward(self, x):
        return self.encoder(x)


class V18RefinerNet(nn.Module):
    """V18 Refiner Network: Takes V16 predictions and refines them."""
    
    def __init__(self, 
                 encoder_name='efficientnet_b0',
                 pretrained=True,
                 cross_row_layers=3,
                 cross_row_dim=128,
                 num_heads=4):
        super().__init__()
        
        self.row_encoder = RefinerEncoder(encoder_name, pretrained)
        enc_channels = self.row_encoder.channels
        
        self.feature_proj = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, None)),
            nn.Flatten(1, 2),
        )
        self.channel_proj = nn.Conv1d(enc_channels[-1], cross_row_dim, 1)
        
        self.cross_row_blocks = nn.ModuleList([
            CrossRowTransformerBlock(
                embed_dim=cross_row_dim,
                num_heads=num_heads,
                mlp_ratio=2.0,
                dropout=0.1
            )
            for _ in range(cross_row_layers)
        ])
        
        self.residual_head = nn.Sequential(
            nn.Conv1d(cross_row_dim, 64, 5, padding=2),
            nn.BatchNorm1d(64),
            nn.GELU(),
            nn.Conv1d(64, 32, 3, padding=1),
            nn.BatchNorm1d(32),
            nn.GELU(),
            nn.Conv1d(32, 1, 1),
            nn.Tanh(),
        )
        
        self.residual_scale = nn.Parameter(torch.tensor(0.1))
        
    def create_guide_channel(self, v16_pred, height, sigma=15.0):
        """Create a Gaussian guide channel from V16 prediction."""
        B, W = v16_pred.shape
        device = v16_pred.device
        
        y_pred = v16_pred * height
        y_grid = torch.arange(height, device=device, dtype=torch.float32)
        y_grid = y_grid.view(1, height, 1)
        y_pred = y_pred.unsqueeze(1)
        
        guide = torch.exp(-0.5 * ((y_grid - y_pred) / sigma) ** 2)
        guide = guide.unsqueeze(1)
        
        return guide
    
    def forward(self, images, v16_preds):
        B, num_rows, C, H, W = images.shape
        
        all_features = []
        
        for row_idx in range(num_rows):
            row_image = images[:, row_idx]
            row_v16 = v16_preds[:, row_idx]
            
            guide = self.create_guide_channel(row_v16, H)
            row_input = torch.cat([row_image, guide], dim=1)
            
            enc_features = self.row_encoder(row_input)
            row_feat = enc_features[-1]
            
            row_feat = self.feature_proj(row_feat)
            row_feat = F.interpolate(row_feat, size=W, mode='linear', align_corners=True)
            row_feat = self.channel_proj(row_feat)
            
            all_features.append(row_feat)
        
        features = torch.stack(all_features, dim=1)
        
        for block in self.cross_row_blocks:
            features = block(features)
        
        residuals = []
        for row_idx in range(num_rows):
            row_feat = features[:, row_idx]
            residual = self.residual_head(row_feat)
            residuals.append(residual.squeeze(1))
        
        residuals = torch.stack(residuals, dim=1)
        scaled_residuals = residuals * self.residual_scale * 0.1
        
        refined = torch.clamp(v16_preds + scaled_residuals, 0, 1)
        
        return refined, scaled_residuals


class V18Pipeline(nn.Module):
    """Full V18 pipeline: V16 (frozen) + Refiner."""
    
    def __init__(self, v16_checkpoint_path, refiner_encoder='efficientnet_b0'):
        super().__init__()
        
        # Load frozen V16
        print(f"Loading V16 from: {v16_checkpoint_path}")
        self.v16 = PerLeadNetV16(pretrained=False)
        checkpoint = torch.load(v16_checkpoint_path, map_location='cpu', weights_only=False)
        state_dict = checkpoint['model']
        if list(state_dict.keys())[0].startswith('module.'):
            state_dict = {k[7:]: v for k, v in state_dict.items()}
        self.v16.load_state_dict(state_dict)
        
        for param in self.v16.parameters():
            param.requires_grad = False
        self.v16.eval()
        
        print(f"V16 loaded (epoch {checkpoint.get('epoch', '?')}, SNR {checkpoint.get('snr', 0):.2f} dB)")
        
        self.refiner = V18RefinerNet(
            encoder_name=refiner_encoder,
            pretrained=False,  # Will load from checkpoint
            cross_row_layers=3,
            cross_row_dim=128,
            num_heads=4,
        )
    
    def forward(self, images):
        B, num_rows, C, H, W = images.shape
        
        with torch.no_grad():
            v16_preds = []
            for row_idx in range(num_rows):
                row_pred = self.v16(images[:, row_idx])
                v16_preds.append(row_pred)
            v16_preds = torch.stack(v16_preds, dim=1)
        
        refined, residuals = self.refiner(images, v16_preds)
        
        return v16_preds, refined, residuals


# =============================================================================
# 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."""
    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 apply_hardcoded_baseline_correction(pred_mv_rows):
    """
    Apply hardcoded baseline correction for each lead.
    
    Uses pre-computed median offsets from 977 images. These are tiny
    (~0.007 mV) so the effect is minimal, but included for completeness.
    
    This is production-ready - no GT required.
    """
    segment_width = len(pred_mv_rows[0]) // 4
    
    for row_idx in range(3):  # Only for rows 0-2 (not rhythm strip)
        lead_names = LEAD_LAYOUT[row_idx]
        for seg_idx, lead_name in enumerate(lead_names):
            offset = BASELINE_OFFSETS.get(lead_name, 0.0)
            
            seg_start = seg_idx * segment_width
            seg_end = (seg_idx + 1) * segment_width
            pred_mv_rows[row_idx][seg_start:seg_end] -= offset
    
    # Rhythm strip uses Lead II offset
    pred_mv_rows[3] -= BASELINE_OFFSETS.get('II', 0.0)
    
    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


def apply_qrs_sharpening(signal_mv, 
                          peak_prominence=0.3,
                          sharpening_radius=5,
                          sharpening_strength=0.5,
                          min_peak_distance=20):
    """
    Apply QRS sharpening to enhance R-peaks and S-waves.
    
    The idea: V18 refinement might slightly blur sharp QRS complexes.
    This detects peaks/troughs and enhances them by reducing local smoothing.
    
    Args:
        signal_mv: Signal in mV
        peak_prominence: Minimum prominence to detect as QRS peak (mV)
        sharpening_radius: Pixels around peak to sharpen
        sharpening_strength: How much to enhance (0=none, 1=full unsmoothing)
        min_peak_distance: Minimum distance between peaks (samples)
    
    Returns:
        Sharpened signal in mV
    """
    signal = signal_mv.copy()
    
    # Compute first derivative to find rapid changes (QRS hallmark)
    derivative = np.gradient(signal)
    abs_derivative = np.abs(derivative)
    
    # Smooth the derivative to find QRS regions
    smoothed_deriv = gaussian_filter1d(abs_derivative, sigma=3)
    
    # Find R-peaks (positive peaks in signal)
    r_peaks, r_props = find_peaks(
        signal, 
        prominence=peak_prominence,
        distance=min_peak_distance
    )
    
    # Find S-waves / Q-waves (negative peaks = troughs)
    s_peaks, s_props = find_peaks(
        -signal, 
        prominence=peak_prominence,
        distance=min_peak_distance
    )
    
    # Combine all QRS-related peaks
    all_peaks = np.concatenate([r_peaks, s_peaks])
    
    if len(all_peaks) == 0:
        return signal
    
    # Create a "sharpening mask" - higher near peaks
    sharpening_mask = np.zeros_like(signal)
    
    for peak_idx in all_peaks:
        # Create Gaussian weight around each peak
        start = max(0, peak_idx - sharpening_radius * 3)
        end = min(len(signal), peak_idx + sharpening_radius * 3)
        
        for i in range(start, end):
            dist = abs(i - peak_idx)
            weight = np.exp(-0.5 * (dist / sharpening_radius) ** 2)
            sharpening_mask[i] = max(sharpening_mask[i], weight)
    
    # Create "unsmoothed" version using smaller Savgol window (preserves peaks better)
    if len(signal) >= 5:
        sharp_signal = savgol_filter(signal, window_length=5, polyorder=2)
    else:
        sharp_signal = signal
    
    # Blend: near peaks use sharper version, elsewhere use original
    # sharpened = original * (1 - mask * strength) + sharp * (mask * strength)
    blend_weight = sharpening_mask * sharpening_strength
    output = signal * (1 - blend_weight) + sharp_signal * blend_weight
    
    # Alternative approach: enhance peak amplitudes slightly
    for peak_idx in r_peaks:
        if 0 <= peak_idx < len(output):
            # Find local max in small window (the true peak might be off by 1-2 samples)
            window_start = max(0, peak_idx - 2)
            window_end = min(len(output), peak_idx + 3)
            local_max_idx = window_start + np.argmax(signal[window_start:window_end])
            
            # Enhance this peak slightly (push it higher)
            enhancement = signal[local_max_idx] * 0.02 * sharpening_strength
            output[local_max_idx] += enhancement
    
    for peak_idx in s_peaks:
        if 0 <= peak_idx < len(output):
            # Find local min
            window_start = max(0, peak_idx - 2)
            window_end = min(len(output), peak_idx + 3)
            local_min_idx = window_start + np.argmin(signal[window_start:window_end])
            
            # Enhance this trough (push it lower)
            enhancement = signal[local_min_idx] * 0.02 * sharpening_strength
            output[local_min_idx] += enhancement  # signal is negative here, so this makes it more negative
    
    return output


# =============================================================================
# Inference and Visualization
# =============================================================================
def load_model(v16_checkpoint, v18_checkpoint, device):
    """Load V18 pipeline (V16 + refiner)."""
    model = V18Pipeline(v16_checkpoint, refiner_encoder='efficientnet_b0')
    
    # Load refiner checkpoint
    checkpoint = torch.load(v18_checkpoint, map_location='cpu', weights_only=False)
    model.refiner.load_state_dict(checkpoint['refiner'])
    
    model = model.to(device)
    model.eval()
    
    print(f"V18 Refiner loaded from epoch {checkpoint['epoch']}")
    print(f"  Refined SNR: {checkpoint.get('snr', 'N/A'):.2f} dB")
    print(f"  Residual scale: {model.refiner.residual_scale.item():.4f}")
    
    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)
    
    row_crop = image[y_start:y_end, T0:T1, :].copy()
    
    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 convert_crop_to_full(pred_y_crop, row_idx):
    """Convert crop-relative y-coordinates to full image coordinates."""
    baseline_y = int(ZERO_MV[row_idx])
    y_start = max(0, baseline_y - CROP_HALF_HEIGHT)
    pad_top = max(0, CROP_HALF_HEIGHT - baseline_y)
    pred_y_full = pred_y_crop - pad_top + y_start
    return pred_y_full


def visualize_predictions(image, v16_predictions, v18_predictions, output_path):
    """
    Draw both V16 and V18 predictions as dots on the image.
    
    V16 = cyan (baseline)
    V18 = red (refined)
    """
    vis_image = image[:, T0:T1, :].copy()
    
    # Colors (BGR)
    v16_colors = [
        (255, 255, 0),   # Cyan for row 0 (V16)
        (255, 255, 0),   # Cyan for row 1 (V16)
        (255, 255, 0),   # Cyan for row 2 (V16)
        (255, 255, 0),   # Cyan for row 3 (V16)
    ]
    
    v18_colors = [
        (0, 0, 255),     # Red for row 0 (V18)
        (0, 255, 0),     # Green for row 1 (V18)
        (255, 0, 255),   # Magenta for row 2 (V18)
        (0, 165, 255),   # Orange for row 3 (V18)
    ]
    
    # Draw V16 first (background)
    for row_idx, pred_y in enumerate(v16_predictions):
        color = v16_colors[row_idx]
        for x in range(0, len(pred_y), 2):  # Every other point for V16
            y = int(np.clip(pred_y[x], 0, TARGET_HEIGHT - 1))
            cv2.circle(vis_image, (x, y), 1, color, -1)
    
    # Draw V18 on top (foreground)
    for row_idx, pred_y in enumerate(v18_predictions):
        color = v18_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."""
    baseline_y = ZERO_MV[row_idx]
    segment_width = OUTPUT_WIDTH // 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
            
            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


@torch.no_grad()
def process_image(model, image_path, csv_path, output_dir, device,
                  apply_smoothing=True, apply_einthoven=True, 
                  apply_qrs_sharp=False, sharpening_strength=0.5,
                  selective_v18=False, 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, 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)
    
    # Prepare row crops
    row_crops = []
    for row_idx in range(4):
        row_crop = crop_row(image, row_idx)
        row_crops.append(row_crop)
    
    # Stack and convert to tensor: [1, 4, 3, H, W]
    row_crops_np = np.stack(row_crops, axis=0)  # [4, H, W, 3]
    row_crops_tensor = torch.from_numpy(row_crops_np.astype(np.float32) / 255.0)
    row_crops_tensor = row_crops_tensor.permute(0, 3, 1, 2)  # [4, 3, H, W]
    row_crops_tensor = row_crops_tensor.unsqueeze(0).to(device)  # [1, 4, 3, H, W]
    
    # Run inference
    with torch.cuda.amp.autocast():
        v16_preds, v18_preds, residuals = model(row_crops_tensor)
    
    # Convert to numpy
    v16_preds = v16_preds[0].cpu().numpy()  # [4, W]
    v18_preds = v18_preds[0].cpu().numpy()  # [4, W]
    residuals = residuals[0].cpu().numpy()  # [4, W]
    
    # Convert to full image coordinates and mV
    v16_predictions = []
    v18_predictions = []
    v18_mv_rows = {}
    
    for row_idx in range(4):
        # V16
        v16_y_crop = v16_preds[row_idx] * ROW_HEIGHT
        v16_y_full = convert_crop_to_full(v16_y_crop, row_idx)
        v16_predictions.append(v16_y_full)
        
        # V18
        v18_y_crop = v18_preds[row_idx] * ROW_HEIGHT
        v18_y_full = convert_crop_to_full(v18_y_crop, row_idx)
        
        # Convert to mV for post-processing
        baseline_y = ZERO_MV[row_idx]
        pred_mv = (baseline_y - v18_y_full) / MV_TO_PIXEL
        
        if apply_smoothing:
            pred_mv = apply_savgol_smoothing(pred_mv, window=7, polyorder=2)
        
        # Apply QRS sharpening AFTER smoothing (recovers peaks that smoothing might blur)
        if apply_qrs_sharp:
            pred_mv = apply_qrs_sharpening(
                pred_mv, 
                peak_prominence=0.3,
                sharpening_radius=5,
                sharpening_strength=sharpening_strength
            )
        
        pred_mv = clamp_ecg_amplitude(pred_mv)
        pred_mv = interpolate_nan(pred_mv.copy())
        
        v18_mv_rows[row_idx] = pred_mv
        v18_predictions.append(v18_y_full)
    
    # Apply Einthoven correction to V18
    if apply_einthoven:
        v18_mv_rows = apply_einthoven_correction(v18_mv_rows, alpha=0.33)
    
    # Apply V18-specific baseline correction (tiny ~0.0001 mV offsets)
    v18_mv_rows = apply_hardcoded_baseline_correction(v18_mv_rows)
    
    # Convert corrected mV back to pixel coordinates
    corrected_v18_predictions = []
    for row_idx in range(4):
        baseline_y = ZERO_MV[row_idx]
        pred_y_corrected = baseline_y - v18_mv_rows[row_idx] * MV_TO_PIXEL
        corrected_v18_predictions.append(pred_y_corrected)
    
    # Selective V18: use V18 only for leads where it helps, V16 for others
    if selective_v18:
        segment_width = OUTPUT_WIDTH // 4
        final_predictions = []
        for row_idx in range(4):
            if row_idx == 3:
                # Row 3 is II_rhythm (full width)
                if 'II_rhythm' in V18_BENEFICIAL_LEADS:
                    final_predictions.append(corrected_v18_predictions[row_idx].copy())
                else:
                    final_predictions.append(v16_predictions[row_idx].copy())
            else:
                # Rows 0-2 have 4 leads each
                row_pred = v16_predictions[row_idx].copy()  # Start with V16
                for col_idx in range(4):
                    lead_name = LEAD_LAYOUT[row_idx][col_idx]
                    seg_start = col_idx * segment_width
                    seg_end = (col_idx + 1) * segment_width
                    if lead_name in V18_BENEFICIAL_LEADS:
                        # Use V18 for this lead segment
                        row_pred[seg_start:seg_end] = corrected_v18_predictions[row_idx][seg_start:seg_end]
                final_predictions.append(row_pred)
        corrected_v18_predictions = final_predictions
    
    # Compute SNR for V16 and V18
    v16_snrs = {}
    v18_snrs = {}
    for row_idx in range(4):
        v16_lead_snrs = compute_snr_per_lead(v16_predictions[row_idx], df, row_idx)
        v18_lead_snrs = compute_snr_per_lead(corrected_v18_predictions[row_idx], df, row_idx)
        v16_snrs.update(v16_lead_snrs)
        v18_snrs.update(v18_lead_snrs)
    
    # Find minimum SNR
    v16_valid_snrs = [v for v in v16_snrs.values() if v is not None]
    v18_valid_snrs = [v for v in v18_snrs.values() if v is not None]
    v16_min_snr = min(v16_valid_snrs) if v16_valid_snrs else 0.0
    v18_min_snr = min(v18_valid_snrs) if v18_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, v16_predictions, corrected_v18_predictions, output_path)
    
    # Save to negative folder if V18 SNR is low
    if negative_dir is not None and v18_min_snr < LOW_SNR_THRESHOLD:
        neg_output_path = negative_dir / f"{sample_id}_{variant}_snr{v18_min_snr:.1f}.png"
        visualize_predictions(image, v16_predictions, corrected_v18_predictions, neg_output_path)
    
    return v16_snrs, v18_snrs, v18_min_snr, output_path


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--v16_checkpoint', type=str, default=None,
                        help='Local V16 checkpoint path. If not provided, will SCP from remote.')
    parser.add_argument('--v18_checkpoint', type=str, default=None,
                        help='Local V18 refiner checkpoint. If not provided, will SCP from remote.')
    parser.add_argument('--remote_host', type=str, default=REMOTE_HOST,
                        help='Remote SSH host for SCP')
    parser.add_argument('--remote_dir', type=str, default=REMOTE_CHECKPOINT_DIR,
                        help='Remote checkpoint directory')
    parser.add_argument('--local_cache', type=str, 
                        default='/tmp/v18_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/v18'))
    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('--qrs_sharpening', action='store_true',
                        help='Enable QRS peak sharpening (enhances R-peaks and S-waves)')
    parser.add_argument('--sharpening_strength', type=float, default=0.5,
                        help='QRS sharpening strength (0-1, default 0.5)')
    parser.add_argument('--selective_v18', action='store_true',
                        help='Apply V18 only to leads where it helps (I,III,aVR,aVF,V1,V6,II_rhythm)')
    parser.add_argument('--no_scp', action='store_true',
                        help='Skip SCP, use local checkpoints 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/v18/negpreds'))
    
    print(f"{'='*70}")
    print(f"V18 Inference: V16 (frozen) + Cross-Row Refiner")
    print(f"{'='*70}")
    
    # Get checkpoints
    v16_checkpoint = args.v16_checkpoint
    v18_checkpoint = args.v18_checkpoint
    local_cache = Path(args.local_cache)
    
    if not args.no_scp:
        print(f"\nFetching checkpoints from {args.remote_host}...")
        
        # V16 checkpoint
        if v16_checkpoint is None:
            v16_remote = f"{args.remote_dir}/{REMOTE_V16_CHECKPOINT}"
            v16_checkpoint = scp_checkpoint(
                args.remote_host, v16_remote, local_cache / REMOTE_V16_CHECKPOINT
            )
        
        # V18 checkpoint (latest)
        if v18_checkpoint is None:
            v18_checkpoint = scp_latest_checkpoint(
                args.remote_host, args.remote_dir, REMOTE_V18_CHECKPOINT_PATTERN, local_cache
            )
    
    # Fallback to local cache
    if v16_checkpoint is None:
        v16_checkpoint = local_cache / REMOTE_V16_CHECKPOINT
    if v18_checkpoint is None:
        v18_checkpoints = sorted(local_cache.glob(f'{REMOTE_V18_CHECKPOINT_PATTERN}*.pth'))
        if v18_checkpoints:
            v18_checkpoint = v18_checkpoints[-1]
    
    # Validate checkpoints exist
    if not Path(v16_checkpoint).exists():
        print(f"ERROR: V16 checkpoint not found: {v16_checkpoint}")
        return
    if not v18_checkpoint or not Path(v18_checkpoint).exists():
        print(f"ERROR: V18 checkpoint not found")
        print(f"  Looked for: {REMOTE_V18_CHECKPOINT_PATTERN}*.pth in {local_cache}")
        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"\nV16 Checkpoint: {v16_checkpoint}")
    print(f"V18 Checkpoint: {v18_checkpoint}")
    print(f"Output: {output_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"QRS Sharpening: {'ON (strength=' + str(args.sharpening_strength) + ')' if args.qrs_sharpening else 'OFF'}")
    print(f"Selective V18: {'ON (I,III,aVR,aVF,V1,V6,II_rhythm only)' if args.selective_v18 else 'OFF (all leads)'}")
    print(f"{'='*70}")
    
    # Load model
    model = load_model(v16_checkpoint, v18_checkpoint, device)
    
    # Find all valid images
    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
    v16_all_snrs = {lead: [] for lead in ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 
                                           'V1', 'V2', 'V3', 'V4', 'V5', 'V6', 'II_rhythm']}
    v18_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:
            v16_snrs, v18_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,
                apply_qrs_sharp=args.qrs_sharpening,
                sharpening_strength=args.sharpening_strength,
                selective_v18=args.selective_v18,
                negative_dir=negative_dir
            )
            if v16_snrs and v18_snrs:
                for lead in v16_snrs:
                    if v16_snrs[lead] is not None:
                        v16_all_snrs[lead].append(v16_snrs[lead])
                    if v18_snrs[lead] is not None:
                        v18_all_snrs[lead].append(v18_snrs[lead])
                
                if min_snr is not None and min_snr < LOW_SNR_THRESHOLD:
                    worst_lead = min(v18_snrs, key=lambda k: v18_snrs[k] if v18_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}")
            import traceback
            traceback.print_exc()
    
    # Print per-lead SNR comparison
    print(f"\n{'='*80}")
    print(f"Per-Lead SNR Comparison (dB): V16 → V18")
    print(f"{'='*80}")
    print(f"{'Lead':<12} {'V16 Mean':>10} {'V18 Mean':>10} {'Δ':>8} {'Count':>6}")
    print(f"{'-'*80}")
    
    v16_total_snrs = []
    v18_total_snrs = []
    
    for lead in ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6', 'II_rhythm']:
        v16_snrs = v16_all_snrs[lead]
        v18_snrs = v18_all_snrs[lead]
        
        if len(v16_snrs) > 0 and len(v18_snrs) > 0:
            v16_mean = np.mean(v16_snrs)
            v18_mean = np.mean(v18_snrs)
            delta = v18_mean - v16_mean
            delta_str = f"{delta:+.2f}" if delta >= 0 else f"{delta:.2f}"
            print(f"{lead:<12} {v16_mean:>10.2f} {v18_mean:>10.2f} {delta_str:>8} {len(v18_snrs):>6}")
            v16_total_snrs.extend(v16_snrs)
            v18_total_snrs.extend(v18_snrs)
        else:
            print(f"{lead:<12} {'N/A':>10} {'N/A':>10} {'N/A':>8} {0:>6}")
    
    print(f"{'-'*80}")
    if len(v16_total_snrs) > 0:
        v16_overall = np.mean(v16_total_snrs)
        v18_overall = np.mean(v18_total_snrs)
        delta = v18_overall - v16_overall
        delta_str = f"{delta:+.2f}" if delta >= 0 else f"{delta:.2f}"
        print(f"{'OVERALL':<12} {v16_overall:>10.2f} {v18_overall:>10.2f} {delta_str:>8} {len(v18_total_snrs):>6}")
    print(f"{'='*80}")
    
    # Save SNR report
    snr_report_path = output_dir / 'snr_comparison.csv'
    with open(snr_report_path, 'w') as f:
        f.write("Lead,V16_Mean,V18_Mean,Delta,Count\n")
        for lead in ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6', 'II_rhythm']:
            v16_snrs = v16_all_snrs[lead]
            v18_snrs = v18_all_snrs[lead]
            if len(v16_snrs) > 0:
                v16_mean = np.mean(v16_snrs)
                v18_mean = np.mean(v18_snrs)
                delta = v18_mean - v16_mean
                f.write(f"{lead},{v16_mean:.2f},{v18_mean:.2f},{delta:+.2f},{len(v18_snrs)}\n")
        if len(v16_total_snrs) > 0:
            f.write(f"OVERALL,{np.mean(v16_total_snrs):.2f},{np.mean(v18_total_snrs):.2f},"
                    f"{np.mean(v18_total_snrs) - np.mean(v16_total_snrs):+.2f},{len(v18_total_snrs)}\n")
    
    print(f"\nSNR comparison 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)")
        print(f"{'='*70}")
        low_snr_samples.sort(key=lambda x: x['min_snr'])
        for item in low_snr_samples[:10]:
            print(f"  {item['sample']}_{item['variant']}: {item['min_snr']:.2f} dB ({item['worst_lead']})")
        if len(low_snr_samples) > 10:
            print(f"  ... and {len(low_snr_samples) - 10} more")
    
    print(f"{'='*70}")


if __name__ == '__main__':
    main()