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

Runs V20 Integral Regression + TCN model on Kaggle images.
Includes per-segment 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

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

# Integral Regression parameters
NUM_BINS = 128

# 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 = [
    ['I', 'aVR', 'V1', 'V4'],
    ['II', 'aVL', 'V2', 'V5'],
    ['III', 'aVF', 'V3', 'V6'],
]

# Hardcoded baseline offsets (mV) - computed from 977 images (V20 epoch 14)
# These are median(prediction - ground_truth) values
# Positive means model predicts slightly higher than GT
BASELINE_OFFSETS = {
    'I': 0.0052, 'II': 0.0071, 'III': 0.0049,
    'aVR': 0.0067, 'aVL': 0.0059, 'aVF': 0.0067,
    'V1': 0.0066, 'V2': 0.0061, 'V3': 0.0067,
    'V4': 0.0068, 'V5': 0.0062, 'V6': 0.0065,
}

# Checkpoint configuration (V20 trains locally)
LOCAL_CHECKPOINT_DIR = '/data/ecg-digitization/checkpoints'
# Remote VM configuration (for V19 which trains remotely)
REMOTE_HOST = os.environ.get('REMOTE_TRAIN_HOST', 'azureuser@172.212.222.231')
REMOTE_CHECKPOINT_DIR = '/data/ecg-digitization/checkpoints'


# =============================================================================
# Model Architecture (from train_v20_integral.py)
# =============================================================================
class CoordConv2d(nn.Module):
    """Adds coordinate channels to convolution."""
    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 TCNBlock(nn.Module):
    """
    Temporal Convolutional Network block with dilated convolutions.
    Uses residual connection for stable training.
    """
    def __init__(self, in_channels, out_channels, kernel_size=3, dilation=1, dropout=0.1):
        super().__init__()
        padding = (kernel_size - 1) * dilation // 2  # Same padding
        
        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size,
                               padding=padding, dilation=dilation)
        self.bn1 = nn.BatchNorm1d(out_channels)
        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size,
                               padding=padding, dilation=dilation)
        self.bn2 = nn.BatchNorm1d(out_channels)
        self.dropout = nn.Dropout(dropout)
        self.activation = nn.GELU()
        
        # Residual connection
        self.residual = nn.Conv1d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity()
    
    def forward(self, x):
        """x: (B, C, W)"""
        residual = self.residual(x)
        
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.activation(out)
        out = self.dropout(out)
        
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.activation(out + residual)
        
        return out


class TCNRefiner(nn.Module):
    """
    Multi-scale TCN for temporal refinement.
    Uses exponentially increasing dilations for large receptive field.
    """
    def __init__(self, in_channels, hidden_channels=64, num_layers=4, kernel_size=5, dropout=0.1):
        super().__init__()
        
        layers = []
        dilations = [2**i for i in range(num_layers)]  # [1, 2, 4, 8]
        
        # First layer
        layers.append(TCNBlock(in_channels, hidden_channels, kernel_size, dilations[0], dropout))
        
        # Hidden layers with increasing dilation
        for dilation in dilations[1:]:
            layers.append(TCNBlock(hidden_channels, hidden_channels, kernel_size, dilation, dropout))
        
        self.tcn = nn.Sequential(*layers)
        
        # Output projection back to input channels
        self.output_proj = nn.Conv1d(hidden_channels, in_channels, 1)
    
    def forward(self, x):
        """x: (B, C, W) - refines features temporally"""
        residual = x
        out = self.tcn(x)
        out = self.output_proj(out)
        return out + residual  # Residual connection


class UNetDecoderBlock(nn.Module):
    """U-Net style decoder block with skip connections."""
    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 IntegralRegressionHead(nn.Module):
    """
    Integral regression head using soft-argmax.
    
    Instead of directly regressing y, we:
    1. Predict a heatmap distribution over height bins
    2. Use soft-argmax to get expected y-value
    
    This provides smoother gradients and handles uncertainty.
    """
    def __init__(self, in_channels, num_bins=NUM_BINS, temperature=1.0):
        super().__init__()
        self.num_bins = num_bins
        self.temperature = temperature
        
        # Register bin centers as buffer (not parameter)
        bin_centers = torch.linspace(0, 1, num_bins)  # Normalized [0, 1]
        self.register_buffer('bin_centers', bin_centers)
        
        # Heatmap prediction head
        self.heatmap_conv = nn.Sequential(
            CoordConv2d(in_channels, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.GELU(),
            nn.Conv2d(128, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.GELU(),
            nn.Conv2d(64, num_bins, 1),  # (B, num_bins, H, W)
        )
        
        # Height attention to pool to (B, num_bins, W)
        self.height_attention = nn.Sequential(
            nn.Conv2d(in_channels, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.GELU(),
            nn.Conv2d(64, 1, 1),
        )
    
    def forward(self, x, temperature=None):
        """
        x: (B, C, H, W) decoder features
        Returns:
            coords: (B, W) predicted y-coordinates in [0, 1]
            heatmap: (B, num_bins, W) probability distribution
        """
        B, C, H, W = x.shape
        temp = temperature if temperature is not None else self.temperature
        
        # Predict heatmap logits
        heatmap_2d = self.heatmap_conv(x)  # (B, num_bins, H, W)
        
        # Height attention for pooling
        attn = self.height_attention(x)  # (B, 1, H, W)
        attn = F.softmax(attn, dim=2)  # Softmax over height
        
        # Pool over height using attention
        heatmap = (heatmap_2d * attn).sum(dim=2)  # (B, num_bins, W)
        
        # Apply temperature and softmax to get distribution
        heatmap_prob = F.softmax(heatmap / temp, dim=1)  # (B, num_bins, W)
        
        # Soft-argmax: expected y-value
        # bin_centers: (num_bins,) -> (1, num_bins, 1)
        coords = (heatmap_prob * self.bin_centers.view(1, -1, 1)).sum(dim=1)  # (B, W)
        
        return coords, heatmap_prob


class IntegralRegressionNet(nn.Module):
    """
    V20: Integral Regression Network with TCN for ECG digitization.
    
    ConvNeXt encoder + U-Net decoder + TCN refiner + Integral Regression.
    
    Input: (B, 3, INPUT_HEIGHT, INPUT_WIDTH) = (B, 3, 500, 3926)
    Output: 
        coords: (B, INPUT_WIDTH) - y-coordinate for each x position [0, 1]
        heatmap: (B, NUM_BINS, INPUT_WIDTH) - probability distribution
    """
    
    def __init__(self, encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=True, 
                 num_bins=NUM_BINS, temperature=1.0, use_tcn=True, tcn_layers=4, use_checkpoint=False):
        super().__init__()
        self.num_bins = num_bins
        self.use_tcn = use_tcn
        self.use_checkpoint = use_checkpoint
        
        # Encoder
        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()
        
        # U-Net Decoder
        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
        
        # Final upsampling
        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(),
        )
        
        # Height pooling to get 1D features for TCN
        self.height_pool = nn.Sequential(
            nn.Conv2d(decoder_dims[-1], 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.GELU(),
            nn.Conv2d(64, 1, 1),  # Attention weights
        )
        
        # TCN refiner for temporal smoothing (operates on 1D features)
        if use_tcn:
            self.tcn_refiner = TCNRefiner(
                in_channels=decoder_dims[-1],
                hidden_channels=64,
                num_layers=tcn_layers,
                kernel_size=5,
                dropout=0.1
            )
        
        # Integral regression head
        self.integral_head = IntegralRegressionHead(
            decoder_dims[-1], num_bins=num_bins, temperature=temperature
        )
    
    def forward(self, x, temperature=None):
        B, C, H, W = x.shape
        
        # Encoder
        features = self.encoder(x)
        
        # Decoder with skip connections
        d = features[-1]
        skips = features[:-1][::-1] + [None]
        
        for block, skip in zip(self.dec_blocks, skips):
            d = block(d, skip)
        
        # Final upsampling
        d = self.final_up(d)  # (B, 32, H', W')
        
        # Resize to match input width
        if d.shape[3] != W:
            d = F.interpolate(d, size=(d.shape[2], W), mode='bilinear', align_corners=True)
        
        # Apply TCN on height-pooled features
        if self.use_tcn:
            # Height attention pooling
            attn = self.height_pool(d)  # (B, 1, H', W)
            attn = F.softmax(attn, dim=2)
            d_1d = (d * attn).sum(dim=2)  # (B, C, W)
            
            # TCN refinement
            d_1d = self.tcn_refiner(d_1d)  # (B, C, W)
            
            # Expand back to 2D for integral head
            d = d_1d.unsqueeze(2).expand(-1, -1, d.shape[2], -1)  # (B, C, H', W)
        
        # Integral regression
        coords, heatmap = self.integral_head(d, temperature)
        
        return coords, heatmap


# =============================================================================
# Post-Processing Functions
# =============================================================================
def apply_savgol_smoothing(signal_mv, window=7, polyorder=2):
    """Apply Savitzky-Golay smoothing."""
    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."""
    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."""
    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_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


# =============================================================================
# SCP Checkpoint Helper
# =============================================================================
def scp_checkpoint(remote_host, remote_path, local_path):
    """SCP a checkpoint from remote to local."""
    local_path = Path(local_path)
    local_path.parent.mkdir(parents=True, exist_ok=True)
    
    cmd = ['scp', '-o', 'StrictHostKeyChecking=no', 
           f'{remote_host}:{remote_path}', str(local_path)]
    print(f"  Downloading: {remote_path}")
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    if result.returncode != 0:
        print(f"  SCP failed: {result.stderr}")
        return None
    return str(local_path)


# =============================================================================
# Inference Functions
# =============================================================================
def load_model(checkpoint_path, device, num_bins=NUM_BINS, use_tcn=True, tcn_layers=4):
    """Load V20 model."""
    model = IntegralRegressionNet(
        encoder_name='convnext_base.fb_in22k_ft_in1k', 
        pretrained=False,
        num_bins=num_bins,
        temperature=1.0,
        use_tcn=use_tcn,
        tcn_layers=tcn_layers
    )
    
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    
    # Handle DDP-wrapped model
    state_dict = checkpoint['model']
    if any(k.startswith('module.') for k in state_dict.keys()):
        state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
    
    model.load_state_dict(state_dict)
    model = model.to(device)
    model.eval()
    
    epoch = checkpoint.get('epoch', 'N/A')
    snr = checkpoint.get('snr', checkpoint.get('best_snr', 'N/A'))
    temp = checkpoint.get('best_temp', 1.0)
    
    print(f"Loaded V20 from epoch {epoch}")
    if isinstance(snr, (int, float)):
        print(f"  Best SNR: {snr:.2f} dB at temperature {temp}")
    
    return model


def crop_row(image, row_idx):
    """Crop a single row centered on its baseline."""
    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 predict_row(model, row_crop, device, temperature=1.0):
    """Run inference on a single row crop."""
    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.amp.autocast('cuda:1'):
            coords, heatmap = model(image_tensor, temperature=temperature)
    
    # coords is in [0, 1], convert to pixels in crop space
    pred_y_crop = coords[0].cpu().numpy() * ROW_HEIGHT
    
    return pred_y_crop, heatmap[0].cpu().numpy()


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, predictions, output_path, heatmaps=None):
    """Draw predictions on the image."""
    vis_image = image[:, T0:T1, :].copy()
    colors = [(255, 0, 0), (0, 255, 0), (255, 0, 255), (0, 165, 255)]
    
    for row_idx, pred_y in enumerate(predictions):
        color = colors[row_idx]
        for x in range(len(pred_y)):
            y = int(round(pred_y[x]))
            if 0 <= y < vis_image.shape[0]:
                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
            
            # Special handling for Lead II (10s data in GT, but short strip shows 2.5s)
            # Use Lead I length as reference to ensure exact alignment
            if lead_name == 'II':
                ref_lead = 'I' if 'I' in df.columns else 'III'
                if ref_lead in df.columns:
                    ref_len = len(df[ref_lead].dropna().values)
                    if len(gt_mv) > ref_len:
                        gt_mv = gt_mv[:ref_len]  # Use exact same length as 2.5s leads
            
            seg_start = seg_idx * segment_width
            seg_end = (seg_idx + 1) * segment_width
            pred_segment = pred_y_full[seg_start:seg_end]
            pred_mv = (baseline_y - pred_segment) / MV_TO_PIXEL
            
            x_pred = np.linspace(0, 1, len(pred_mv))
            x_gt = np.linspace(0, 1, len(gt_mv))
            pred_mv_resampled = np.interp(x_gt, x_pred, pred_mv)
            
            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:
                lead_snrs[lead_name] = 10 * np.log10(signal_power / noise_power)
    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 = (baseline_y - pred_y_full) / MV_TO_PIXEL
        x_pred = np.linspace(0, 1, len(pred_mv))
        x_gt = np.linspace(0, 1, len(gt_mv))
        pred_mv_resampled = np.interp(x_gt, x_pred, pred_mv)
        
        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:
            lead_snrs['II_rhythm'] = 10 * np.log10(signal_power / noise_power)
    
    return lead_snrs


def process_image(model, image_path, csv_path, output_dir, device, temperature=1.0,
                  apply_smoothing=True, apply_einthoven=True, apply_baseline_fix=True,
                  negative_dir=None):
    """Process a single image and compute SNR."""
    image = cv2.imread(str(image_path), cv2.IMREAD_COLOR)
    if image is None:
        print(f"Failed to load: {image_path}")
        return None, None, None
    
    image = image[Y0:Y1, X0:X1]
    image = cv2.resize(image, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR)
    
    df = pd.read_csv(csv_path)
    
    predictions = []
    pred_mv_rows = {}
    all_heatmaps = []
    
    for row_idx in range(4):
        row_crop = crop_row(image, row_idx)
        pred_y_crop, heatmap = predict_row(model, row_crop, device, temperature=temperature)
        pred_y_full = convert_crop_to_full(pred_y_crop, row_idx)
        
        baseline_y = ZERO_MV[row_idx]
        pred_mv = (baseline_y - pred_y_full) / MV_TO_PIXEL
        
        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)
        all_heatmaps.append(heatmap)
    
    if apply_einthoven:
        pred_mv_rows = apply_einthoven_correction(pred_mv_rows, alpha=0.33)
    
    # Apply hardcoded baseline correction (tiny effect, ~0.007 mV)
    if apply_baseline_fix:
        pred_mv_rows = apply_hardcoded_baseline_correction(pred_mv_rows)
    
    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)
    
    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
    
    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, all_heatmaps)
    
    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, all_heatmaps)
    
    return all_lead_snrs, min_snr, output_path


def main():
    parser = argparse.ArgumentParser(description='V20 Inference with Integral Regression + TCN')
    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 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=LOCAL_CHECKPOINT_DIR,
                        help='Local checkpoint directory')
    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/v20'))
    parser.add_argument('--num_samples', type=int, default=None)
    parser.add_argument('--use_holdout', action='store_true', default=True)
    parser.add_argument('--no_smoothing', action='store_true')
    parser.add_argument('--no_einthoven', action='store_true')
    parser.add_argument('--no_scp', action='store_true', default=True,
                        help='Skip SCP, use local checkpoint (default: True for V20)')
    parser.add_argument('--temperature', type=float, default=1.0,
                        help='Temperature for soft-argmax (lower = sharper)')
    parser.add_argument('--num_bins', type=int, default=NUM_BINS,
                        help='Number of bins for integral regression')
    parser.add_argument('--use_tcn', action='store_true', default=True,
                        help='Use TCN refiner')
    parser.add_argument('--no_tcn', action='store_true',
                        help='Disable TCN refiner')
    parser.add_argument('--tcn_layers', type=int, default=4,
                        help='Number of TCN layers')
    parser.add_argument('--no_baseline_fix', action='store_true',
                        help='Disable hardcoded baseline correction')
    parser.add_argument('--seed', type=int, default=42)
    args = parser.parse_args()
    
    random.seed(args.seed)
    device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
    output_dir = Path(args.output_dir)
    negative_dir = Path(os.path.expanduser('~/tmp/pred/v20/negpreds'))
    
    use_tcn = args.use_tcn and not args.no_tcn
    
    print(f"{'='*70}")
    print(f"V20 Inference: Integral Regression + TCN (aVR/aVL/aVF Baseline Fix)")
    print(f"{'='*70}")
    
    # Get checkpoint
    checkpoint = args.checkpoint
    local_cache = Path(args.local_cache)
    
    if not args.no_scp and checkpoint is None:
        print(f"\nFetching checkpoint from {args.remote_host}...")
        local_cache.mkdir(parents=True, exist_ok=True)
        
        # Try latest first, then best
        for name in ['v20_integral_latest.pth', 'v20_integral_best_snr.pth']:
            remote_path = f"{args.remote_dir}/{name}"
            local_path = local_cache / name
            checkpoint = scp_checkpoint(args.remote_host, remote_path, local_path)
            if checkpoint is not None:
                break
        
        if checkpoint is None:
            print("Failed to download checkpoint. Use --checkpoint to specify local path.")
            sys.exit(1)
    elif checkpoint is None:
        # Look for local checkpoint
        for name in ['v20_integral_latest.pth', 'v20_integral_best_snr.pth']:
            local_path = local_cache / name
            if local_path.exists():
                checkpoint = str(local_path)
                break
        
        if checkpoint is None:
            checkpoint = '/data/ecg-digitization/checkpoints/v20_integral_latest.pth'
    
    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"\nCheckpoint: {checkpoint}")
    print(f"Output: {output_dir}")
    print(f"Device: {device}")
    print(f"Temperature: {args.temperature}")
    print(f"TCN: {'ON' if use_tcn else 'OFF'} ({args.tcn_layers} layers)")
    print(f"Num bins: {args.num_bins}")
    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"Per-segment baseline fix: {'OFF' if args.no_baseline_fix else 'ON'}")
    print(f"{'='*70}")
    
    # Load model
    model = load_model(checkpoint, device, 
                       num_bins=args.num_bins, 
                       use_tcn=use_tcn, 
                       tcn_layers=args.tcn_layers)
    
    # Find 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))
    
    print(f"Found {len(all_samples)} valid images in holdout 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...")
    
    all_snrs = {lead: [] for lead in ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 
                                       'V1', 'V2', 'V3', 'V4', 'V5', 'V6', 'II_rhythm']}
    
    for img_path, csv_path in tqdm(selected):
        try:
            lead_snrs, min_snr, _ = process_image(
                model, img_path, csv_path, output_dir, device,
                temperature=args.temperature,
                apply_smoothing=not args.no_smoothing,
                apply_einthoven=not args.no_einthoven,
                apply_baseline_fix=not args.no_baseline_fix,
                negative_dir=negative_dir
            )
            if lead_snrs:
                for lead, snr in lead_snrs.items():
                    if snr is not None:
                        all_snrs[lead].append(snr)
        except Exception as e:
            print(f"Error processing {img_path}: {e}")
            import traceback
            traceback.print_exc()
    
    # Print results
    print(f"\n{'='*70}")
    print(f"Per-Lead SNR Statistics (dB) - Temperature={args.temperature}")
    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:
            print(f"{lead:<12} {np.mean(snrs):>8.2f} {np.std(snrs):>8.2f} "
                  f"{np.min(snrs):>8.2f} {np.max(snrs):>8.2f} {len(snrs):>6}")
            total_snrs.extend(snrs)
    
    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}")
    
    print(f"\nDone! Saved {len(selected)} visualizations to {output_dir}")


if __name__ == '__main__':
    main()