#!/usr/bin/env python3
"""
Find Best V18+V19 Ensemble Ratio

Runs inference ONCE for both models, saves predictions, then sweeps
through different blend ratios to find the optimal ensemble weights.
"""

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

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

import cv2
import timm

try:
    from torchvision.ops import DeformConv2d
    HAS_DEFORM_CONV = True
except ImportError:
    HAS_DEFORM_CONV = False


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

CROP_HALF_HEIGHT = 250
ROW_HEIGHT = 500
ECG_MV_MIN, ECG_MV_MAX = -10.0, 10.0

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

LEAD_LAYOUT = [
    ['I', 'aVR', 'V1', 'V4'],
    ['II', 'aVL', 'V2', 'V5'],
    ['III', 'aVF', 'V3', 'V6'],
]

LOCAL_CHECKPOINT_DIR = '/data/ecg-digitization/checkpoints'


# =============================================================================
# V16 Model
# =============================================================================
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):
    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)
        return torch.sigmoid(out).squeeze(1)


# =============================================================================
# V18 Refiner
# =============================================================================
class CrossRowAttention(nn.Module):
    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).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) + residual
        return out.reshape(B, W, num_rows, C).permute(0, 2, 3, 1)


class CrossRowTransformerBlock(nn.Module):
    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.ffn(self.norm(x)) + residual.permute(0, 3, 1, 2).reshape(B * W, num_rows, C)
        return x.reshape(B, W, num_rows, C).permute(0, 2, 3, 1)


class RefinerEncoder(nn.Module):
    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):
    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(cross_row_dim, num_heads, 2.0, 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):
        B, W = v16_pred.shape
        y_pred = v16_pred * height
        y_grid = torch.arange(height, device=v16_pred.device, dtype=torch.float32).view(1, height, 1)
        return torch.exp(-0.5 * ((y_grid - y_pred.unsqueeze(1)) / sigma) ** 2).unsqueeze(1)
    
    def forward(self, images, v16_preds):
        B, num_rows, C, H, W = images.shape
        all_features = []
        for row_idx in range(num_rows):
            guide = self.create_guide_channel(v16_preds[:, row_idx], H)
            row_input = torch.cat([images[:, row_idx], guide], dim=1)
            row_feat = self.row_encoder(row_input)[-1]
            row_feat = self.feature_proj(row_feat)
            row_feat = F.interpolate(row_feat, size=W, mode='linear', align_corners=True)
            all_features.append(self.channel_proj(row_feat))
        features = torch.stack(all_features, dim=1)
        for block in self.cross_row_blocks:
            features = block(features)
        residuals = torch.stack([self.residual_head(features[:, i]).squeeze(1) for i in range(num_rows)], dim=1)
        scaled_residuals = residuals * self.residual_scale * 0.1
        return torch.clamp(v16_preds + scaled_residuals, 0, 1), scaled_residuals


# =============================================================================
# V19 Model
# =============================================================================
class DeformableConvBlock(nn.Module):
    def __init__(self, in_ch, out_ch, kernel_size=3, stride=1, padding=1):
        super().__init__()
        if HAS_DEFORM_CONV:
            self.offset_conv = nn.Sequential(nn.Conv2d(in_ch, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.Conv2d(64, 2 * kernel_size * kernel_size, 3, padding=1))
            nn.init.zeros_(self.offset_conv[-1].weight)
            nn.init.zeros_(self.offset_conv[-1].bias)
            self.deform_conv = DeformConv2d(in_ch, out_ch, kernel_size, stride=stride, padding=padding)
        else:
            self.conv = nn.Conv2d(in_ch, out_ch, kernel_size, stride=stride, padding=padding)
        self.norm = nn.BatchNorm2d(out_ch)
        self.act = nn.GELU()
    
    def forward(self, x):
        if HAS_DEFORM_CONV:
            out = self.deform_conv(x, self.offset_conv(x))
        else:
            out = self.conv(x)
        return self.act(self.norm(out))


class BiLSTMHead(nn.Module):
    def __init__(self, input_dim, hidden_dim=128, num_layers=2, dropout=0.1):
        super().__init__()
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True, bidirectional=True, dropout=dropout if num_layers > 1 else 0)
        self.output_proj = nn.Sequential(nn.Linear(hidden_dim * 2, hidden_dim), nn.LayerNorm(hidden_dim), nn.GELU())
        self.output_dim = hidden_dim
    def forward(self, x):
        return self.output_proj(self.lstm(x.permute(0, 2, 1))[0])


class AuxiliaryHeads(nn.Module):
    def __init__(self, feature_dim):
        super().__init__()
        self.grid_head = nn.Sequential(nn.Conv2d(32, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(inplace=True), nn.Conv2d(16, 1, 1), nn.Sigmoid())
        self.gradient_head = nn.Sequential(nn.Linear(feature_dim, 64), nn.GELU(), nn.Linear(64, 1), nn.Tanh())
        self.uncertainty_head = nn.Sequential(nn.Linear(feature_dim, 64), nn.GELU(), nn.Linear(64, 1))
    def forward(self, features_2d, features_1d):
        return self.grid_head(features_2d), self.gradient_head(features_1d).squeeze(-1), self.uncertainty_head(features_1d).squeeze(-1)


class UNetDecoderBlockV19(nn.Module):
    def __init__(self, in_ch, skip_ch, out_ch, use_deform=False):
        super().__init__()
        if use_deform and HAS_DEFORM_CONV:
            self.conv1 = DeformableConvBlock(in_ch + skip_ch, out_ch)
        else:
            self.conv1 = nn.Sequential(nn.Conv2d(in_ch + skip_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.GELU())
        self.conv2 = nn.Sequential(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.conv2(self.conv1(x))


class PerLeadNetV19(nn.Module):
    def __init__(self, encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=True):
        super().__init__()
        self.encoder = timm.create_model(encoder_name, pretrained=pretrained, features_only=True, out_indices=(0, 1, 2, 3))
        enc_channels = self.encoder.feature_info.channels()
        decoder_dims = [256, 128, 64, 32]
        self.dec_blocks = nn.ModuleList()
        in_ch = enc_channels[-1]
        skip_channels = enc_channels[:-1][::-1] + [0]
        for i, (skip_ch, out_ch) in enumerate(zip(skip_channels, decoder_dims)):
            self.dec_blocks.append(UNetDecoderBlockV19(in_ch, skip_ch, out_ch, i >= 2))
            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.bilstm = BiLSTMHead(decoder_dims[-1], 128, 2, 0.1)
        self.regression_head = nn.Sequential(nn.Linear(self.bilstm.output_dim, 64), nn.GELU(), nn.Linear(64, 1), nn.Sigmoid())
        self.aux_heads = AuxiliaryHeads(self.bilstm.output_dim)
    
    def forward(self, x, return_aux=False):
        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)
        features_2d = d
        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 = F.softmax(self.height_attention(d), dim=2)
        pooled = (d * attn).sum(dim=2)
        temporal_features = self.bilstm(pooled)
        y_pred = self.regression_head(temporal_features).squeeze(-1)
        if return_aux:
            return y_pred, {'grid': self.aux_heads.grid_head(features_2d), 'gradient': self.aux_heads.gradient_head(temporal_features).squeeze(-1), 'log_var': self.aux_heads.uncertainty_head(temporal_features).squeeze(-1)}
        return y_pred


# =============================================================================
# Helper Functions
# =============================================================================
def crop_row(image, row_idx):
    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 apply_savgol_smoothing(signal_mv, window=7, polyorder=2):
    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):
    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
    pred_mv_rows[0][:segment_width] = lead_I + alpha * error
    pred_mv_rows[2][:segment_width] = lead_III + alpha * error
    return pred_mv_rows


def compute_snr(pred_mv_rows, df, epsilon=1e-10):
    """Compute overall SNR for a prediction."""
    segment_width = OUTPUT_WIDTH // 4
    all_snrs = []
    
    for row_idx in range(3):
        lead_names = LEAD_LAYOUT[row_idx]
        for seg_idx, lead_name in enumerate(lead_names):
            if lead_name not in df.columns:
                continue
            gt_mv = df[lead_name].dropna().values
            if len(gt_mv) == 0:
                continue
            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]
            seg_start = seg_idx * segment_width
            seg_end = (seg_idx + 1) * segment_width
            pred_mv = pred_mv_rows[row_idx][seg_start:seg_end]
            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:
                all_snrs.append(10 * np.log10(signal_power / noise_power))
    
    # Lead II rhythm strip
    if 'II' in df.columns:
        gt_mv = df['II'].dropna().values
        if len(gt_mv) > 0:
            pred_mv = pred_mv_rows[3]
            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:
                all_snrs.append(10 * np.log10(signal_power / noise_power))
    
    return np.mean(all_snrs) if all_snrs else 0.0


def load_models(v16_ckpt, v18_ckpt, v19_ckpt, device):
    # V16
    v16_model = PerLeadNetV16(pretrained=False)
    ckpt = torch.load(v16_ckpt, map_location=device, weights_only=False)
    state = ckpt['model']
    if any(k.startswith('module.') for k in state.keys()):
        state = {k.replace('module.', ''): v for k, v in state.items()}
    v16_model.load_state_dict(state)
    v16_model = v16_model.to(device).eval()
    print(f"V16: epoch {ckpt.get('epoch', '?')}, SNR {ckpt.get('snr', 0):.2f} dB")
    
    # V18 refiner
    v18_refiner = V18RefinerNet(pretrained=False)
    ckpt = torch.load(v18_ckpt, map_location=device, weights_only=False)
    v18_refiner.load_state_dict(ckpt['refiner'])
    v18_refiner = v18_refiner.to(device).eval()
    print(f"V18: epoch {ckpt.get('epoch', '?')}, SNR {ckpt.get('snr', 0):.2f} dB")
    
    # V19
    v19_model = PerLeadNetV19(pretrained=False)
    ckpt = torch.load(v19_ckpt, map_location=device, weights_only=False)
    state = ckpt['model']
    if any(k.startswith('module.') for k in state.keys()):
        state = {k.replace('module.', ''): v for k, v in state.items()}
    v19_model.load_state_dict(state)
    v19_model = v19_model.to(device).eval()
    print(f"V19: epoch {ckpt.get('epoch', '?')}, SNR {ckpt.get('snr', ckpt.get('best_snr', 0)):.2f} dB")
    
    return v16_model, v18_refiner, v19_model


@torch.no_grad()
def run_inference(v16_model, v18_refiner, v19_model, samples, device):
    """Run inference once for all samples, return raw mV predictions."""
    all_predictions = []
    
    for img_path, csv_path in tqdm(samples, desc="Running inference"):
        image = cv2.imread(str(img_path), cv2.IMREAD_COLOR)
        if image is None:
            continue
        
        image = image[Y0:Y1, X0:X1]
        image = cv2.resize(image, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR)
        
        df = pd.read_csv(csv_path)
        row_crops = [crop_row(image, i) for i in range(4)]
        
        # V18 prediction
        row_tensors = [torch.from_numpy(c.astype(np.float32) / 255.0).permute(2, 0, 1) for c in row_crops]
        images = torch.stack(row_tensors, dim=0).unsqueeze(0).to(device)
        
        with torch.amp.autocast('cuda', dtype=torch.float16):
            v16_preds = torch.stack([v16_model(images[:, i]) for i in range(4)], dim=1)
            refined, _ = v18_refiner(images, v16_preds)
        
        v18_mv = {}
        refined_np = refined[0].cpu().numpy()
        for row_idx in range(4):
            pred_y_crop = refined_np[row_idx] * ROW_HEIGHT
            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
            v18_mv[row_idx] = (ZERO_MV[row_idx] - pred_y_full) / MV_TO_PIXEL
        
        # V19 prediction
        v19_mv = {}
        for row_idx in range(4):
            row_tensor = torch.from_numpy(row_crops[row_idx].astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0).to(device)
            with torch.amp.autocast('cuda'):
                output = v19_model(row_tensor, return_aux=False)
            pred_y_crop = output[0].cpu().numpy() * ROW_HEIGHT
            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
            v19_mv[row_idx] = (ZERO_MV[row_idx] - pred_y_full) / MV_TO_PIXEL
        
        all_predictions.append({
            'img_path': str(img_path),
            'csv_path': str(csv_path),
            'v18_mv': v18_mv,
            'v19_mv': v19_mv,
            'df': df,
        })
    
    return all_predictions


def evaluate_ratio(predictions, v18_weight, v19_weight, apply_smoothing=True, apply_einthoven=True):
    """Evaluate a specific ensemble ratio."""
    snrs = []
    
    for pred in predictions:
        # Blend predictions
        ensemble = {}
        for row_idx in range(4):
            ensemble[row_idx] = v18_weight * pred['v18_mv'][row_idx] + v19_weight * pred['v19_mv'][row_idx]
        
        # Post-processing
        for row_idx in range(4):
            if apply_smoothing:
                ensemble[row_idx] = apply_savgol_smoothing(ensemble[row_idx], window=7, polyorder=2)
            ensemble[row_idx] = np.clip(ensemble[row_idx], ECG_MV_MIN, ECG_MV_MAX)
        
        if apply_einthoven:
            ensemble = apply_einthoven_correction(ensemble, alpha=0.33)
        
        # Compute SNR
        snr = compute_snr(ensemble, pred['df'])
        if snr > 0:
            snrs.append(snr)
    
    return np.mean(snrs) if snrs else 0.0


def main():
    parser = argparse.ArgumentParser(description='Find best V18+V19 ensemble ratio')
    parser.add_argument('--v16_checkpoint', type=str, default=f'{LOCAL_CHECKPOINT_DIR}/v16_perlead_epoch020.pth')
    parser.add_argument('--v18_checkpoint', type=str, default=f'{LOCAL_CHECKPOINT_DIR}/v18_refiner_best.pth')
    parser.add_argument('--v19_checkpoint', type=str, default=f'{LOCAL_CHECKPOINT_DIR}/v19_enhanced_epoch014.pth')
    parser.add_argument('--kaggle_data', type=str, default='/data/ecg-digitization/stage1_data/train')
    parser.add_argument('--num_samples', type=int, default=2000)
    parser.add_argument('--cache_file', type=str, default='/tmp/ensemble_predictions.pkl')
    parser.add_argument('--seed', type=int, default=42)
    args = parser.parse_args()
    
    random.seed(args.seed)
    np.random.seed(args.seed)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    print(f"{'='*70}")
    print(f"Finding Best V18+V19 Ensemble Ratio")
    print(f"{'='*70}")
    
    # Check if we have cached predictions
    cache_path = Path(args.cache_file)
    if cache_path.exists():
        print(f"Loading cached predictions from {cache_path}...")
        with open(cache_path, 'rb') as f:
            predictions = pickle.load(f)
        print(f"Loaded {len(predictions)} cached predictions")
    else:
        # Load models
        print("\nLoading models...")
        v16_model, v18_refiner, v19_model = load_models(
            args.v16_checkpoint, args.v18_checkpoint, args.v19_checkpoint, device
        )
        
        # Find samples
        kaggle_dir = Path(args.kaggle_data)
        all_samples = []
        
        for sample_dir in kaggle_dir.iterdir():
            if not sample_dir.is_dir():
                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")
        
        if len(all_samples) > args.num_samples:
            samples = random.sample(all_samples, args.num_samples)
        else:
            samples = all_samples
        
        print(f"Running inference on {len(samples)} samples...")
        predictions = run_inference(v16_model, v18_refiner, v19_model, samples, device)
        
        # Cache predictions
        print(f"Caching predictions to {cache_path}...")
        with open(cache_path, 'wb') as f:
            pickle.dump(predictions, f)
    
    # Sweep ratios
    print(f"\n{'='*70}")
    print(f"Sweeping V18:V19 ratios...")
    print(f"{'='*70}")
    
    ratios = [
        (0.0, 1.0),   # V19 only
        (0.1, 0.9),
        (0.2, 0.8),
        (0.3, 0.7),
        (0.4, 0.6),
        (0.5, 0.5),   # Equal
        (0.6, 0.4),
        (0.7, 0.3),
        (0.8, 0.2),
        (0.9, 0.1),
        (1.0, 0.0),   # V18 only
    ]
    
    results = []
    print(f"\n{'V18%':>6} {'V19%':>6} {'Mean SNR':>10}")
    print("-" * 30)
    
    for v18_w, v19_w in ratios:
        snr = evaluate_ratio(predictions, v18_w, v19_w)
        results.append((v18_w, v19_w, snr))
        print(f"{v18_w*100:>5.0f}% {v19_w*100:>5.0f}% {snr:>10.2f} dB")
    
    # Find best
    best = max(results, key=lambda x: x[2])
    
    print(f"\n{'='*70}")
    print(f"BEST RATIO: V18={best[0]*100:.0f}%, V19={best[1]*100:.0f}%  →  SNR={best[2]:.2f} dB")
    print(f"{'='*70}")
    
    # Fine-grained search around best
    print(f"\nFine-tuning around best ratio...")
    fine_ratios = []
    for delta in [-0.15, -0.10, -0.05, 0.0, 0.05, 0.10, 0.15]:
        v18_w = max(0, min(1, best[0] + delta))
        v19_w = 1 - v18_w
        fine_ratios.append((v18_w, v19_w))
    
    fine_results = []
    print(f"\n{'V18%':>6} {'V19%':>6} {'Mean SNR':>10}")
    print("-" * 30)
    
    for v18_w, v19_w in fine_ratios:
        snr = evaluate_ratio(predictions, v18_w, v19_w)
        fine_results.append((v18_w, v19_w, snr))
        print(f"{v18_w*100:>5.0f}% {v19_w*100:>5.0f}% {snr:>10.2f} dB")
    
    final_best = max(fine_results, key=lambda x: x[2])
    
    print(f"\n{'='*70}")
    print(f"FINAL BEST: V18={final_best[0]*100:.0f}%, V19={final_best[1]*100:.0f}%  →  SNR={final_best[2]:.2f} dB")
    print(f"{'='*70}")
    
    # Compare to individual models
    v18_only = evaluate_ratio(predictions, 1.0, 0.0)
    v19_only = evaluate_ratio(predictions, 0.0, 1.0)
    
    print(f"\nComparison:")
    print(f"  V18 only:     {v18_only:.2f} dB")
    print(f"  V19 only:     {v19_only:.2f} dB")
    print(f"  Best ensemble: {final_best[2]:.2f} dB  (+{final_best[2] - max(v18_only, v19_only):.2f} dB)")


if __name__ == '__main__':
    main()