code/scripts/find_ensemble_ratio.py

#!/usr/bin/env python3
"""
Find optimal ensemble ratio between V10.1 model and Friend's Net3 model.
Evaluates on random training samples with all variants.
Uses pre-processed stage1 images from /data/ecg-digitization/stage1_data/train/
"""

import os
import sys
import numpy as np
import pandas as pd
import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as T
import timm
from pathlib import Path
from tqdm import tqdm
from scipy.signal import savgol_filter
import scipy.signal
import scipy.optimize
import random
from typing import Tuple

# Add baseline path
BASELINE_PATH = '/data/ecg-digitization/hengck23-submit-physionet/hengck23-submit-physionet'
sys.path.insert(0, BASELINE_PATH)

import stage2_common as s2c
from stage2_model import MyCoordUnetDecoder, encode_with_resnet

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

DEVICE = "cuda:0"
VALID_VARIANTS = ['0001', '0003', '0004', '0005', '0006', '0009', '0010', '0011', '0012']


# =============================================================================
# V10.1 Model
# =============================================================================
class CoordDecoderBlock(nn.Module):
    def __init__(self, in_ch, skip_ch, out_ch, scale=2):
        super().__init__()
        self.scale = scale
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch + skip_ch + 2, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        )

    def forward(self, x, skip=None):
        x = F.interpolate(x, scale_factor=self.scale, mode='nearest')
        if skip is not None:
            x = torch.cat([x, skip], dim=1)
        b, c, h, w = x.shape
        cy, cx = torch.meshgrid(
            torch.linspace(-1, 1, h, device=x.device, dtype=x.dtype),
            torch.linspace(-1, 1, w, device=x.device, dtype=x.dtype),
            indexing='ij'
        )
        coord = torch.stack([cx, cy]).unsqueeze(0).expand(b, -1, -1, -1)
        x = torch.cat([x, coord], dim=1)
        return self.conv(x)


class ECGNetV10(nn.Module):
    def __init__(self, decoder_dims=[256, 128, 64, 32, 16]):
        super().__init__()
        self.encoder = timm.create_model(
            'efficientnet_b4.ra2_in1k', pretrained=False,
            features_only=True, out_indices=(0, 1, 2, 3, 4)
        )
        enc_dims = [24, 32, 56, 160, 448]
        
        self.dec_blocks = nn.ModuleList()
        in_ch = enc_dims[-1]
        skip_chs = enc_dims[:-1][::-1] + [0]
        while len(decoder_dims) < len(skip_chs):
            decoder_dims.append(decoder_dims[-1])
        decoder_dims = decoder_dims[:len(skip_chs)]
        
        for skip_ch, out_ch in zip(skip_chs, decoder_dims):
            self.dec_blocks.append(CoordDecoderBlock(in_ch, skip_ch, out_ch))
            in_ch = out_ch
        
        self.seg_head = nn.Conv2d(decoder_dims[-1], 4, 1)
        self.reg_head = nn.Sequential(
            nn.Conv2d(decoder_dims[-1], 64, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.AdaptiveAvgPool2d((1, None)),
        )
        self.reg_out = nn.Sequential(
            nn.Conv1d(64, 32, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv1d(32, 4, 1),
            nn.Sigmoid()
        )
    
    def forward(self, x):
        input_size = x.shape[2:]
        enc = self.encoder(x)
        d = enc[-1]
        skips = enc[:-1][::-1] + [None]
        for block, skip in zip(self.dec_blocks, skips):
            d = block(d, skip)
        if d.shape[2:] != input_size:
            d = F.interpolate(d, size=input_size, mode='bilinear', align_corners=False)
        seg_logits = self.seg_head(d)
        reg_feat = self.reg_head(d).squeeze(2)
        reg_coords = self.reg_out(reg_feat)
        return seg_logits, reg_coords


# =============================================================================
# Friend's Net3 Model
# =============================================================================
class Net3(nn.Module):
    def __init__(self, pretrained=True):
        super().__init__()
        encoder_dim = [64, 128, 256, 512]
        decoder_dim = [128, 64, 32, 16]
        self.encoder = timm.create_model(
            model_name='resnet34.a3_in1k',
            pretrained=pretrained,
            in_chans=3,
            num_classes=0,
            global_pool=''
        )
        self.decoder = MyCoordUnetDecoder(
            in_channel=encoder_dim[-1],
            skip_channel=encoder_dim[:-1][::-1] + [0],
            out_channel=decoder_dim,
            scale=[2, 2, 2, 2]
        )
        self.pixel = nn.Conv2d(decoder_dim[-1], 4, 1)

    def forward(self, image):
        encode = encode_with_resnet(self.encoder, image)
        last, _ = self.decoder(feature=encode[-1], skip=encode[:-1][::-1] + [None])
        return self.pixel(last)


# =============================================================================
# Helper Functions
# =============================================================================
def soft_argmax(heatmap, temperature=SOFT_ARGMAX_TEMP):
    B, C, H, W = heatmap.shape
    y_coords = torch.arange(H, device=heatmap.device, dtype=heatmap.dtype).view(1, 1, H, 1)
    weights = F.softmax(heatmap * temperature, dim=2)
    return (weights * y_coords).sum(dim=2)


def resample_signal(signal, target_length):
    if len(signal) == target_length:
        return signal
    x_old = np.linspace(0, 1, len(signal))
    x_new = np.linspace(0, 1, target_length)
    return np.interp(x_new, x_old, signal)


# =============================================================================
# Competition-style SNR (with alignment)
# =============================================================================
PERFECT_SCORE = 10**(384/10)  # Maximum SNR ratio
MAX_TIME_SHIFT = 0.2  # 200ms max shift (as fraction of second)

# Lead layout for extraction
ROW_LAYOUT = [
    ['I', 'aVR', 'V1', 'V4'],
    ['II_short', 'aVL', 'V2', 'V5'],
    ['III', 'aVF', 'V3', 'V6'],
]


def align_signals(label: np.ndarray, pred: np.ndarray, max_shift: float = float('inf')) -> np.ndarray:
    """Align prediction to label using cross-correlation (competition style)."""
    if np.any(~np.isfinite(label)):
        raise ValueError('values in label should all be finite')
    if np.sum(np.isfinite(pred)) == 0:
        raise ValueError('prediction can not all be infinite')

    label_arr = np.asarray(label, dtype=np.float64)
    pred_arr = np.asarray(pred, dtype=np.float64)
    label_mean = np.mean(label_arr)
    pred_mean = np.nanmean(pred_arr[np.isfinite(pred_arr)])
    label_arr_centered = label_arr - label_mean
    pred_arr_centered = np.where(np.isfinite(pred_arr), pred_arr - pred_mean, 0)
    
    correlation = scipy.signal.correlate(label_arr_centered, pred_arr_centered, mode='full')
    n_label = np.size(label_arr)
    n_pred = np.size(pred_arr)
    lags = scipy.signal.correlation_lags(n_label, n_pred, mode='full')
    valid_lags_mask = (lags >= -max_shift) & (lags <= max_shift)
    
    max_correlation = np.nanmax(correlation[valid_lags_mask])
    all_max_indices = np.flatnonzero(correlation == max_correlation)
    best_idx = min(all_max_indices, key=lambda i: abs(lags[i]))
    time_shift = lags[best_idx]
    
    start_padding_len = max(time_shift, 0)
    pred_slice_start = max(-time_shift, 0)
    pred_slice_end = min(n_label - time_shift, n_pred)
    end_padding_len = max(n_label - n_pred - time_shift, 0)
    
    aligned_pred = np.concatenate((
        np.full(start_padding_len, np.nan), 
        pred_arr[pred_slice_start:pred_slice_end], 
        np.full(end_padding_len, np.nan)
    ))

    def objective_func(v_shift):
        return np.nansum((label_arr - (aligned_pred - v_shift)) ** 2)

    if np.any(np.isfinite(label_arr) & np.isfinite(aligned_pred)):
        results = scipy.optimize.minimize_scalar(objective_func, method='Brent')
        vertical_shift = results.x
        aligned_pred -= vertical_shift
    
    return aligned_pred


def compute_power(label: np.ndarray, prediction: np.ndarray) -> Tuple[float, float]:
    """Compute signal and noise power (competition style)."""
    if label.ndim != 1 or prediction.ndim != 1:
        raise ValueError('Inputs must be 1-dimensional arrays.')
    finite_mask = np.isfinite(prediction)
    if not np.any(finite_mask):
        raise ValueError("The 'prediction' array contains no finite values.")
    prediction = prediction.copy()
    prediction[~np.isfinite(prediction)] = 0
    noise = label - prediction
    p_signal = np.sum(label**2)
    p_noise = np.sum(noise**2)
    return p_signal, p_noise


def series_to_leads(series_mv):
    """Convert 4-row mV series to 12-lead dictionary."""
    leads = {}
    segment_width = series_mv.shape[1] // 4
    
    for row_idx in range(3):
        for seg_idx, lead_name in enumerate(ROW_LAYOUT[row_idx]):
            actual_name = 'II_short' if lead_name == 'II_short' else lead_name
            start = seg_idx * segment_width
            end = (seg_idx + 1) * segment_width
            leads[actual_name] = series_mv[row_idx, start:end]
    
    leads['II'] = series_mv[3]  # Full rhythm strip
    return leads


def compute_competition_snr(pred_mv, gt_df, sig_len):
    """Compute SNR using competition-style scoring with alignment."""
    # Get sampling frequency (assume 10-second recording)
    fs = sig_len // 10 if sig_len >= 10 else 500
    max_shift = int(fs * MAX_TIME_SHIFT)
    
    # Convert pred to leads
    pred_leads = series_to_leads(pred_mv)
    
    # Get short lead length
    short_lead_len = len(gt_df['I'].dropna()) if 'I' in gt_df.columns else sig_len // 4
    
    sum_signal = 0.0
    sum_noise = 0.0
    
    # Standard leads (short duration)
    short_leads = ['I', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']
    
    for lead in short_leads:
        if lead not in gt_df.columns or lead not in pred_leads:
            continue
        
        gt_lead = gt_df[lead].dropna().values
        if len(gt_lead) == 0:
            continue
        
        pred_lead = pred_leads[lead]
        pred_resampled = resample_signal(pred_lead, len(gt_lead))
        
        try:
            aligned_pred = align_signals(gt_lead, pred_resampled, max_shift)
            p_signal, p_noise = compute_power(gt_lead, aligned_pred)
            sum_signal += p_signal
            sum_noise += p_noise
        except Exception:
            # Fallback - direct comparison
            p_signal = np.sum(gt_lead**2)
            p_noise = np.sum((gt_lead - pred_resampled)**2)
            sum_signal += p_signal
            sum_noise += p_noise
    
    # II_short (first 2.5s of II)
    if 'II' in gt_df.columns and 'II_short' in pred_leads:
        full_ii = gt_df['II'].dropna().values
        gt_ii_short = full_ii[:short_lead_len]
        pred_ii_short = pred_leads['II_short']
        pred_resampled = resample_signal(pred_ii_short, len(gt_ii_short))
        
        try:
            aligned_pred = align_signals(gt_ii_short, pred_resampled, max_shift)
            p_signal, p_noise = compute_power(gt_ii_short, aligned_pred)
            sum_signal += p_signal
            sum_noise += p_noise
        except Exception:
            p_signal = np.sum(gt_ii_short**2)
            p_noise = np.sum((gt_ii_short - pred_resampled)**2)
            sum_signal += p_signal
            sum_noise += p_noise
    
    # Full II (rhythm strip) - only count the portion NOT covered by II_short
    # Actually competition may score II separately. Let's include it.
    if 'II' in gt_df.columns and 'II' in pred_leads:
        gt_ii = gt_df['II'].dropna().values
        pred_ii = pred_leads['II']
        pred_resampled = resample_signal(pred_ii, len(gt_ii))
        
        try:
            aligned_pred = align_signals(gt_ii, pred_resampled, max_shift)
            p_signal, p_noise = compute_power(gt_ii, aligned_pred)
            sum_signal += p_signal
            sum_noise += p_noise
        except Exception:
            p_signal = np.sum(gt_ii**2)
            p_noise = np.sum((gt_ii - pred_resampled)**2)
            sum_signal += p_signal
            sum_noise += p_noise
    
    # Compute final SNR
    if sum_noise == 0:
        return PERFECT_SCORE
    elif sum_signal == 0:
        return 0
    else:
        snr_ratio = min((sum_signal / sum_noise), PERFECT_SCORE)
        return snr_ratio  # Return ratio, convert to dB later


def compute_snr_4row(pred_mv, gt_mv):
    """Compute average SNR across all 4 rows (legacy, for quick comparison)."""
    snrs = []
    for row in range(4):
        pred = resample_signal(pred_mv[row], len(gt_mv[row]))
        signal_power = np.mean(gt_mv[row] ** 2)
        noise_power = np.mean((pred - gt_mv[row]) ** 2)
        if noise_power > 1e-10 and signal_power > 1e-10:
            snr = 10 * np.log10(signal_power / noise_power)
            snrs.append(np.clip(snr, -10, 50))
    return np.mean(snrs) if snrs else 0.0


# =============================================================================
# Inference Functions
# =============================================================================
@torch.no_grad()
def process_v10(model, image_bgr, device):
    """V10 inference on stage1 image. Returns 4-row mV signal."""
    h, w = image_bgr.shape[:2]
    crop_h, crop_w = min(h, Y1), min(w, X1)
    image_cropped = image_bgr[:crop_h, :crop_w]
    image_resized = cv2.resize(image_cropped, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR)
    
    image_tensor = torch.from_numpy(image_resized.astype(np.float32) / 255.0)
    image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0).to(device)
    
    with torch.amp.autocast('cuda', dtype=torch.float32):
        seg_logits, _ = model(image_tensor)
    
    seg_logits = torch.nan_to_num(seg_logits, nan=0.0, posinf=0.0, neginf=0.0)
    seg_probs = torch.sigmoid(seg_logits.float())
    signal_full = soft_argmax(seg_probs).cpu().numpy()[0]
    
    signal_pixel = signal_full[:, T0:T1]
    signal_mv = np.zeros_like(signal_pixel)
    for row_idx in range(4):
        signal_mv[row_idx] = (ZERO_MV[row_idx] - signal_pixel[row_idx]) / MV_TO_PIXEL
    
    return np.clip(signal_mv, -7.0, 7.0)


@torch.no_grad()
def process_net3(model, image_bgr, sig_len, device):
    """Net3 (Friend's) inference on stage1 image. Returns 4-row mV signal."""
    resize = T.Resize((1696, 4352), interpolation=T.InterpolationMode.BILINEAR)
    
    # Use RGB for Net3 (friend's model expects RGB)
    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
    img = image_rgb[Y0:Y1, X0:X1] / 255.0
    batch = resize(torch.from_numpy(np.ascontiguousarray(img.transpose(2, 0, 1))).unsqueeze(0))
    batch = batch.float().to(device)
    
    with torch.amp.autocast('cuda', dtype=torch.float32):
        output = model(batch)
    
    pixel = torch.sigmoid(output).float().cpu().numpy()[0]
    series_in_pixel = s2c.pixel_to_series(pixel[..., T0:T1], ZERO_MV.tolist(), sig_len)
    series_mv = (ZERO_MV.reshape(4, 1) - series_in_pixel) / MV_TO_PIXEL
    
    return np.clip(series_mv, -7.0, 7.0)


def load_ground_truth(csv_path, sig_len):
    """Load ground truth from Kaggle CSV and convert to 4-row mV format."""
    df = pd.read_csv(csv_path)
    
    segment_width = sig_len // 4
    gt_4row = np.zeros((4, sig_len), dtype=np.float32)
    
    # Get short lead length (use I as reference)
    short_lead_len = len(df['I'].dropna()) if 'I' in df.columns else segment_width
    
    row_layout = [
        ['I', 'aVR', 'V1', 'V4'],
        ['II_short', 'aVL', 'V2', 'V5'],  # II_short = first part of II matching short lead length
        ['III', 'aVF', 'V3', 'V6'],
    ]
    
    for row_idx in range(3):
        for seg_idx, lead_name in enumerate(row_layout[row_idx]):
            # Handle II_short specially - use first portion of II
            if lead_name == 'II_short':
                if 'II' in df.columns:
                    full_ii = df['II'].dropna().values
                    # Take first short_lead_len samples (same duration as other short leads)
                    ii_short = full_ii[:short_lead_len]
                    resampled = resample_signal(ii_short, segment_width)
                    start = seg_idx * segment_width
                    end = (seg_idx + 1) * segment_width
                    gt_4row[row_idx, start:end] = resampled
            elif lead_name in df.columns:
                signal = df[lead_name].dropna().values
                resampled = resample_signal(signal, segment_width)
                start = seg_idx * segment_width
                end = (seg_idx + 1) * segment_width
                gt_4row[row_idx, start:end] = resampled
    
    # Row 3: full II
    if 'II' in df.columns:
        gt_4row[3] = resample_signal(df['II'].dropna().values, sig_len)
    
    return gt_4row


def ensemble_predictions(pred_v10, pred_friend, weight_v10):
    """Ensemble two 4-row predictions."""
    weight_friend = 1.0 - weight_v10
    target_len = max(pred_v10.shape[1], pred_friend.shape[1])
    
    result = np.zeros((4, target_len), dtype=np.float32)
    for row in range(4):
        v10_row = resample_signal(pred_v10[row], target_len)
        friend_row = resample_signal(pred_friend[row], target_len)
        result[row] = weight_v10 * v10_row + weight_friend * friend_row
    
    return result


def main():
    print("="*70)
    print("Finding Optimal Ensemble Ratio")
    print("="*70)
    
    # Paths - use pre-processed stage1 images
    data_dir = Path("/data/ecg-digitization/stage1_data/train")
    
    # Model paths - downloaded from Kaggle
    v10_weights = "/data/ecg-digitization/checkpoints/ecg_v10_best.pth"
    net3_weights = "/data/ecg-digitization/checkpoints/net3_kaggle/iter_0004200.pt"
    
    # Check if weights exist, otherwise try alternative locations
    if not os.path.exists(v10_weights):
        alt_path = "/data/ecg-digitization/checkpoints/v10_efficientnet_b4_best_reg.pth"
        if os.path.exists(alt_path):
            v10_weights = alt_path
        else:
            print(f"ERROR: V10 weights not found at {v10_weights}")
            print("Please download from: https://www.kaggle.com/models/udaybhatia/ecg-v10-best/PyTorch/pytorch/4")
            return
    
    if not os.path.exists(net3_weights):
        alt_path = f"{BASELINE_PATH}/weight/stage2-00005810.checkpoint.pth"
        if os.path.exists(alt_path):
            net3_weights = alt_path
        else:
            print(f"ERROR: Net3 weights not found at {net3_weights}")
            print("Please download from: https://www.kaggle.com/models/wasupandceacar/physio-seg-public/PyTorch/net3_009_4200/1")
            return
    
    # Load V10 model
    print(f"\nLoading V10 model from {v10_weights}...")
    v10_model = ECGNetV10()
    checkpoint = torch.load(v10_weights, map_location='cpu', weights_only=False)
    state_dict = checkpoint.get('model', checkpoint)
    if isinstance(state_dict, dict) and list(state_dict.keys())[0].startswith('module.'):
        state_dict = {k[7:]: v for k, v in state_dict.items()}
    v10_model.load_state_dict(state_dict)
    v10_model.to(DEVICE).eval()
    print(f"V10 loaded - epoch {checkpoint.get('epoch', '?')}")
    
    # Load Net3 model (friend's)
    print(f"\nLoading Net3 model from {net3_weights}...")
    net3_model = Net3(pretrained=False).to(DEVICE).eval()
    st = torch.load(net3_weights, map_location="cpu")
    if isinstance(st, dict) and "state_dict" in st:
        st = st["state_dict"]
    net3_model.load_state_dict(st, strict=True)
    print("Net3 loaded.")
    
    # Collect samples - use pre-processed stage1 images
    all_samples = []
    for sample_dir in sorted(data_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 variant in VALID_VARIANTS:
            img_path = sample_dir / f"{sample_dir.name}-{variant}.png"
            if img_path.exists():
                all_samples.append((img_path, csv_path, sample_dir.name, variant))
    
    print(f"\nTotal samples with all variants: {len(all_samples)}")
    
    # Random sample 1000
    random.seed(42)
    if len(all_samples) > 1000:
        samples = random.sample(all_samples, 1000)
    else:
        samples = all_samples
    
    print(f"Evaluating on {len(samples)} samples...")
    
    # Collect predictions
    results = []
    
    for img_path, csv_path, sample_id, variant in tqdm(samples, desc="Processing"):
        try:
            # Read pre-processed stage1 image (already rectified)
            img_bgr = cv2.imread(str(img_path))
            if img_bgr is None:
                continue
            
            # Get ground truth
            gt_df = pd.read_csv(csv_path)
            sig_len = len(gt_df['II'].dropna()) if 'II' in gt_df.columns else 5000
            gt_4row = load_ground_truth(csv_path, sig_len)
            
            # V10 prediction (expects BGR)
            pred_v10 = process_v10(v10_model, img_bgr, DEVICE)
            
            # Net3 prediction (friend's model)
            pred_net3 = process_net3(net3_model, img_bgr, sig_len, DEVICE)
            
            # Resample to same length
            target_len = max(pred_v10.shape[1], pred_net3.shape[1])
            pred_v10_resampled = np.zeros((4, target_len), dtype=np.float32)
            pred_net3_resampled = np.zeros((4, target_len), dtype=np.float32)
            for row in range(4):
                pred_v10_resampled[row] = resample_signal(pred_v10[row], target_len)
                pred_net3_resampled[row] = resample_signal(pred_net3[row], target_len)
            
            results.append({
                'sample_id': sample_id,
                'variant': variant,
                'pred_v10': pred_v10_resampled,
                'pred_friend': pred_net3_resampled,
                'gt': gt_4row,
                'gt_df': gt_df,
                'sig_len': sig_len
            })
            
        except Exception as e:
            print(f"Error {sample_id}-{variant}: {e}")
            continue
    
    print(f"\nSuccessfully processed {len(results)} samples")
    
    if len(results) == 0:
        print("ERROR: No samples processed successfully!")
        return
    
    # Evaluate different ratios
    print("\n" + "="*70)
    print("Evaluating ensemble ratios (competition-style with alignment)...")
    print("="*70)
    
    ratios = np.arange(0.0, 1.05, 0.05)
    ratio_results = []
    
    for ratio in ratios:
        snr_ratios = []  # SNR as ratio, not dB
        for r in results:
            ensemble = ensemble_predictions(r['pred_v10'], r['pred_friend'], ratio)
            snr_ratio = compute_competition_snr(ensemble, r['gt_df'], r['sig_len'])
            snr_ratios.append(snr_ratio)
        
        # Competition final score: mean of ratios, then to dB
        mean_ratio = np.mean(snr_ratios)
        avg_snr_db = max(10 * np.log10(mean_ratio), -384) if mean_ratio > 0 else -384
        ratio_results.append((ratio, avg_snr_db, mean_ratio))
        print(f"V10 weight: {ratio:.2f}, Friend weight: {1-ratio:.2f} -> SNR: {avg_snr_db:.3f} dB")
    
    # Find best ratio
    best_ratio, best_snr, _ = max(ratio_results, key=lambda x: x[1])
    
    print("\n" + "="*70)
    print(f"BEST RATIO: V10={best_ratio:.2f}, Friend={1-best_ratio:.2f}")
    print(f"Best SNR: {best_snr:.3f} dB")
    print("="*70)
    
    # Also show individual model performance with competition scoring
    v10_ratios = [compute_competition_snr(r['pred_v10'], r['gt_df'], r['sig_len']) for r in results]
    friend_ratios = [compute_competition_snr(r['pred_friend'], r['gt_df'], r['sig_len']) for r in results]
    
    v10_snr_db = max(10 * np.log10(np.mean(v10_ratios)), -384)
    friend_snr_db = max(10 * np.log10(np.mean(friend_ratios)), -384)
    
    print(f"\nV10 only:    {v10_snr_db:.3f} dB")
    print(f"Friend only: {friend_snr_db:.3f} dB")
    print(f"Ensemble:    {best_snr:.3f} dB (improvement: +{best_snr - max(v10_snr_db, friend_snr_db):.3f} dB)")
    
    # Breakdown by variant
    print("\n" + "="*70)
    print("SNR by variant (using best ensemble ratio):")
    print("="*70)
    
    for variant in VALID_VARIANTS:
        variant_results = [r for r in results if r['variant'] == variant]
        if variant_results:
            v10_r = [compute_competition_snr(r['pred_v10'], r['gt_df'], r['sig_len']) for r in variant_results]
            friend_r = [compute_competition_snr(r['pred_friend'], r['gt_df'], r['sig_len']) for r in variant_results]
            ensemble_r = []
            for r in variant_results:
                ensemble = ensemble_predictions(r['pred_v10'], r['pred_friend'], best_ratio)
                ensemble_r.append(compute_competition_snr(ensemble, r['gt_df'], r['sig_len']))
            
            v10_db = max(10 * np.log10(np.mean(v10_r)), -384)
            friend_db = max(10 * np.log10(np.mean(friend_r)), -384)
            ensemble_db = max(10 * np.log10(np.mean(ensemble_r)), -384)
            print(f"  {variant}: V10={v10_db:.2f}, Friend={friend_db:.2f}, Ensemble={ensemble_db:.2f} dB (n={len(variant_results)})")
    
    return best_ratio


if __name__ == "__main__":
    best_ratio = main()