#!/usr/bin/env python3
"""
Evaluate individual models and ensembles on validation set.

Models:
- NET3 (baseline public solution)
- V16 (ConvNeXt-Base per-lead)
- V18 (V16 + Cross-Row Refiner)
- V19_base (ConvNeXt-Base + BiLSTM + DeformableConv)

Metrics:
- MAE (Mean Absolute Error)
- Normalized MSE (MSE / var(gt))
- Scale Invariant SNR (SI-SNR)
- SNR (Signal-to-Noise Ratio)
- Log-Cosh Loss
"""

import os
import sys
import argparse
import numpy as np
import pandas as pd
from pathlib import Path
from tqdm import tqdm
from scipy.signal import savgol_filter
from itertools import combinations
import json

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

import cv2
import timm

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

try:
    import stage0_common as s0c
    import stage1_common as s1c
    import stage2_common as s2c
    from stage0_model import Net as Stage0Net
    from stage1_model import Net as Stage1Net
    from stage2_model import MyCoordUnetDecoder, encode_with_resnet
    HAS_BASELINE = True
except ImportError:
    HAS_BASELINE = False
    print("Warning: Baseline not available, Net3 will be skipped")

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

VAL_SAMPLE_IDS = [
    '1006427285', '1006867983', '1012423188', '10140238', '1015663939',
    '102150619', '1026034238', '1041099777', '104573050', '1048962695',
    '1052007218', '1053922973', '1059602762', '1063816858',
    '1067371646', '1067975047', '1068062585', '1072767337',
    '1084993373', '108599929'
]

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

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


# =============================================================================
# Metrics
# =============================================================================
def compute_mae(pred, gt):
    """Mean Absolute Error."""
    return np.mean(np.abs(pred - gt))


def compute_normalized_mse(pred, gt):
    """Normalized MSE = MSE / var(gt)."""
    mse = np.mean((pred - gt) ** 2)
    var_gt = np.var(gt)
    if var_gt < 1e-10:
        return 0.0
    return mse / var_gt


def compute_snr(pred, gt, epsilon=1e-10):
    """Signal-to-Noise Ratio in dB."""
    signal_power = np.mean(gt ** 2)
    noise_power = np.mean((pred - gt) ** 2)
    if noise_power < epsilon:
        return 50.0
    return 10 * np.log10(signal_power / noise_power)


def compute_si_snr(pred, gt, epsilon=1e-10):
    """Scale-Invariant SNR in dB."""
    # SI-SNR: project pred onto gt, then compute SNR
    # s_target = <pred, gt> / ||gt||^2 * gt
    gt_norm_sq = np.sum(gt ** 2)
    if gt_norm_sq < epsilon:
        return 0.0
    alpha = np.sum(pred * gt) / gt_norm_sq
    s_target = alpha * gt
    e_noise = pred - s_target
    
    signal_power = np.mean(s_target ** 2)
    noise_power = np.mean(e_noise ** 2)
    
    if noise_power < epsilon:
        return 50.0
    return 10 * np.log10(signal_power / noise_power)


def compute_log_cosh(pred, gt):
    """Log-Cosh loss."""
    diff = pred - gt
    return np.mean(np.log(np.cosh(diff + 1e-12)))


def compute_all_metrics(pred, gt):
    """Compute all metrics for a single signal pair."""
    return {
        'mae': compute_mae(pred, gt),
        'nmse': compute_normalized_mse(pred, gt),
        'si_snr': compute_si_snr(pred, gt),
        'snr': compute_snr(pred, gt),
        'log_cosh': compute_log_cosh(pred, gt),
    }


# =============================================================================
# V16 Model Architecture
# =============================================================================
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)
        return self.conv(torch.cat([x, yy, xx], dim=1))


class UNetDecoderBlock(nn.Module):
    def __init__(self, in_ch, skip_ch, out_ch):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch + skip_ch, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch), nn.GELU(),
            nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch), nn.GELU(),
        )
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
    
    def forward(self, x, skip=None):
        x = self.upsample(x)
        if skip is not None:
            if x.shape[2:] != skip.shape[2:]:
                x = F.interpolate(x, size=skip.shape[2:], mode='bilinear', align_corners=True)
            x = torch.cat([x, skip], dim=1)
        return self.conv(x)


class PerLeadNet(nn.Module):
    """V16 Per-Lead ECG Extraction Network."""
    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(32, 32, 3, padding=1, bias=False), nn.BatchNorm2d(32), nn.GELU())
        
        self.height_attention = nn.Sequential(
            CoordConv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU(), nn.Conv2d(64, 1, 1))
        
        self.regression_head = nn.Sequential(
            nn.Conv1d(32, 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 = F.softmax(self.height_attention(d), dim=2)
        d = (d * attn).sum(dim=2)
        out = torch.sigmoid(self.regression_head(d))
        return out.squeeze(1)


# =============================================================================
# V18 Refiner Model Architecture
# =============================================================================
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)
        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
        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.norm(x)
        x = self.ffn(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
        device = v16_pred.device
        y_pred = v16_pred * height
        y_grid = torch.arange(height, device=device, dtype=torch.float32).view(1, height, 1)
        y_pred = y_pred.unsqueeze(1)
        guide = torch.exp(-0.5 * ((y_grid - y_pred) / sigma) ** 2)
        return guide.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):
            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


# =============================================================================
# V19 Model Architecture
# =============================================================================
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:
            offset = self.offset_conv(x)
            out = self.deform_conv(x, offset)
        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):
        x = x.permute(0, 2, 1)
        lstm_out, _ = self.lstm(x)
        return self.output_proj(lstm_out)


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):
    """V19: ConvNeXt-Base + BiLSTM + Deformable Conv."""
    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)):
            use_deform = (i >= 2)
            self.dec_blocks.append(UNetDecoderBlockV19(in_ch, skip_ch, out_ch, use_deform))
            in_ch = out_ch
        
        self.final_up = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            nn.Conv2d(32, 32, 3, padding=1, bias=False), nn.BatchNorm2d(32), nn.GELU())
        
        self.height_attention = nn.Sequential(
            CoordConv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU(), nn.Conv2d(64, 1, 1))
        
        self.bilstm = BiLSTMHead(32, 128, 2, 0.1)
        self.regression_head = nn.Sequential(nn.Linear(128, 64), nn.GELU(), nn.Linear(64, 1), nn.Sigmoid())
        self.aux_heads = AuxiliaryHeads(128)
    
    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, self.aux_heads(features_2d, temporal_features)
        return y_pred


# =============================================================================
# Net3 Model Architecture (use baseline's stage2 model)
# =============================================================================
if HAS_BASELINE:
    from stage2_model import Net as Stage2Net


# =============================================================================
# Inference Helpers
# =============================================================================
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 convert_crop_to_mv(pred_normalized, row_idx):
    """Convert normalized prediction to mV."""
    pred_y_crop = pred_normalized * 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
    return (ZERO_MV[row_idx] - pred_y_full) / MV_TO_PIXEL


def resample_signal(signal, target_length):
    """Resample signal to 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)


# =============================================================================
# Model Loaders
# =============================================================================
def load_v16(checkpoint_path, device):
    model = PerLeadNet(encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=False)
    ckpt = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
    state_dict = ckpt['model'] if 'model' in ckpt else ckpt
    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).eval()
    print(f"V16 loaded: epoch {ckpt.get('epoch', '?')}, SNR: {ckpt.get('snr', 0):.2f} dB")
    return model


def load_v18(checkpoint_path, device):
    model = V18RefinerNet(encoder_name='efficientnet_b0', pretrained=False)
    ckpt = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
    state_dict = ckpt['refiner'] if 'refiner' in ckpt else ckpt.get('model', ckpt)
    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).eval()
    print(f"V18 loaded: epoch {ckpt.get('epoch', '?')}, SNR: {ckpt.get('snr', 0):.2f} dB")
    return model


def load_v19(checkpoint_path, device):
    model = PerLeadNetV19(encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=False)
    ckpt = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
    state_dict = ckpt['model'] if 'model' in ckpt else ckpt
    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).eval()
    print(f"V19 loaded: epoch {ckpt.get('epoch', '?')}, SNR: {ckpt.get('snr', 0):.2f} dB")
    return model


def load_net3(checkpoint_path, device):
    if not HAS_BASELINE:
        return None
    model = Stage2Net(pretrained=False)
    ckpt = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
    state_dict = ckpt['state_dict'] if 'state_dict' in ckpt else ckpt
    model.load_state_dict(state_dict)
    model = model.to(device).eval()
    print(f"Net3 (Stage2) loaded")
    return model


# =============================================================================
# Inference Functions
# =============================================================================
@torch.no_grad()
def predict_v16(model, row_crops, device):
    """Run V16 on row crops, return 4-row mV signal."""
    series_mv = []
    for row_idx, crop in enumerate(row_crops):
        tensor = torch.from_numpy(crop.astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0).to(device)
        with torch.amp.autocast('cuda'):
            pred = model(tensor)
        pred_np = pred[0].cpu().numpy()
        series_mv.append(convert_crop_to_mv(pred_np, row_idx))
    return np.array(series_mv)


@torch.no_grad()
def predict_v16v18(v16_model, v18_model, row_crops, device):
    """Run V16 + V18 refiner on row crops, return 4-row mV signal."""
    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)  # [1, 4, 3, H, W]
    
    with torch.amp.autocast('cuda'):
        v16_preds = []
        for row_idx in range(4):
            pred = v16_model(images[:, row_idx])
            v16_preds.append(pred)
        v16_preds = torch.stack(v16_preds, dim=1)
        refined, _ = v18_model(images, v16_preds)
    
    refined_np = refined[0].cpu().numpy()
    series_mv = np.array([convert_crop_to_mv(refined_np[i], i) for i in range(4)])
    return series_mv


@torch.no_grad()
def predict_v19(model, row_crops, device):
    """Run V19 on row crops, return 4-row mV signal."""
    series_mv = []
    for row_idx, crop in enumerate(row_crops):
        tensor = torch.from_numpy(crop.astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0).to(device)
        with torch.amp.autocast('cuda'):
            pred = model(tensor, return_aux=False)
        pred_np = pred[0].cpu().numpy()
        series_mv.append(convert_crop_to_mv(pred_np, row_idx))
    return np.array(series_mv)


@torch.no_grad()
def predict_net3(model, image_rgb, device):
    """Run Net3 (Stage2) on image, return 4-row mV signal."""
    if model is None:
        return None
    
    # Crop and resize
    x0, x1 = 0, 2176
    y0, y1 = 0, 1696
    img = image_rgb[y0:y1, x0:x1] / 255.0
    batch = torch.from_numpy(np.ascontiguousarray(img.transpose(2, 0, 1))).unsqueeze(0).float()
    batch = F.interpolate(batch, size=(TARGET_HEIGHT, TARGET_WIDTH), mode='bilinear', align_corners=True)
    batch = batch.to(device)
    
    with torch.amp.autocast('cuda'):
        output = model(batch)
    
    pixel = torch.sigmoid(output).float().cpu().numpy()[0]
    # pixel shape: [4, H, W]
    
    # Use pixel_to_series from baseline
    series_in_pixel = s2c.pixel_to_series(pixel[..., T0:T1], ZERO_MV, length=OUTPUT_WIDTH)
    series_mv = (ZERO_MV.reshape(4, 1) - series_in_pixel) / MV_TO_PIXEL
    return series_mv


def series_to_leads(series_mv):
    """Convert 4-row 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(LEAD_LAYOUT[row_idx]):
            start = seg_idx * segment_width
            end = (seg_idx + 1) * segment_width
            leads[lead_name] = series_mv[row_idx, start:end]
    
    # Lead II short strip (first segment of row 1)
    leads['II'] = series_mv[1, :segment_width]  # Use short strip for fair comparison
    return leads


def get_gt_leads(df):
    """Get ground truth leads from CSV."""
    leads = {}
    for lead in ALL_LEADS:
        if lead in df.columns:
            signal = df[lead].dropna().values
            if lead == 'II' and len(signal) > 5000:
                # Use first quarter (short strip equivalent)
                signal = signal[:len(signal)//4]
            leads[lead] = signal
    return leads


# =============================================================================
# Ensemble Function
# =============================================================================
def ensemble_signals(signals_dict, weights):
    """
    Ensemble multiple model predictions.
    signals_dict: {'model_name': leads_dict}
    weights: {'model_name': weight}
    """
    model_names = list(signals_dict.keys())
    all_leads = list(signals_dict[model_names[0]].keys())
    
    ensemble = {}
    for lead in all_leads:
        signals = []
        ws = []
        target_len = 0
        for model_name in model_names:
            if lead in signals_dict[model_name]:
                sig = signals_dict[model_name][lead]
                target_len = max(target_len, len(sig))
                signals.append(sig)
                ws.append(weights[model_name])
        
        if not signals:
            continue
        
        # Resample to same length
        resampled = [resample_signal(s, target_len) for s in signals]
        
        # Weighted average
        total_weight = sum(ws)
        ensemble[lead] = sum(w * s for w, s in zip(ws, resampled)) / total_weight
    
    return ensemble


# =============================================================================
# Main Evaluation
# =============================================================================
def evaluate_on_validation(models, kaggle_data, device):
    """
    Evaluate all models on validation set.
    Returns: dict of {model_name: {metric: value}}
    """
    kaggle_dir = Path(kaggle_data)
    val_sample_set = set(VAL_SAMPLE_IDS)
    
    # Collect all samples
    all_samples = []
    for sample_dir in kaggle_dir.iterdir():
        if not sample_dir.is_dir() or 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))
                break  # One image per sample
    
    print(f"Found {len(all_samples)} validation samples")
    
    # Initialize metrics storage
    model_names = list(models.keys())
    metrics_per_model = {name: {'mae': [], 'nmse': [], 'si_snr': [], 'snr': [], 'log_cosh': []}
                         for name in model_names}
    
    for img_path, csv_path in tqdm(all_samples, desc="Evaluating"):
        try:
            # Load image
            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)
            
            # Load GT
            df = pd.read_csv(csv_path)
            gt_leads = get_gt_leads(df)
            
            # Crop rows
            row_crops = [crop_row(image, i) for i in range(4)]
            
            # Get predictions from each model
            preds = {}
            
            if 'v16' in models and models['v16'] is not None:
                series = predict_v16(models['v16'], row_crops, device)
                preds['v16'] = series_to_leads(np.clip(series, ECG_MV_MIN, ECG_MV_MAX))
            
            if 'v16v18' in models and models['v16v18'] is not None:
                v16_m, v18_m = models['v16v18']
                series = predict_v16v18(v16_m, v18_m, row_crops, device)
                preds['v16v18'] = series_to_leads(np.clip(series, ECG_MV_MIN, ECG_MV_MAX))
            
            if 'v19' in models and models['v19'] is not None:
                series = predict_v19(models['v19'], row_crops, device)
                preds['v19'] = series_to_leads(np.clip(series, ECG_MV_MIN, ECG_MV_MAX))
            
            if 'net3' in models and models['net3'] is not None:
                series = predict_net3(models['net3'], image, device)
                if series is not None:
                    preds['net3'] = series_to_leads(np.clip(series, ECG_MV_MIN, ECG_MV_MAX))
            
            # Compute metrics for each model
            for model_name, pred_leads in preds.items():
                for lead in ALL_LEADS:
                    if lead not in pred_leads or lead not in gt_leads:
                        continue
                    
                    pred = pred_leads[lead]
                    gt = gt_leads[lead]
                    
                    # Resample to same length
                    target_len = len(gt)
                    pred_resampled = resample_signal(pred, target_len)
                    
                    # Compute metrics
                    m = compute_all_metrics(pred_resampled, gt)
                    for metric_name, value in m.items():
                        metrics_per_model[model_name][metric_name].append(value)
        
        except Exception as e:
            print(f"Error: {img_path}: {e}")
            continue
    
    # Average metrics
    results = {}
    for model_name in model_names:
        results[model_name] = {}
        for metric_name in ['mae', 'nmse', 'si_snr', 'snr', 'log_cosh']:
            values = metrics_per_model[model_name][metric_name]
            if values:
                results[model_name][metric_name] = np.mean(values)
            else:
                results[model_name][metric_name] = None
    
    return results


def evaluate_ensembles(models, kaggle_data, device, weight_steps=11):
    """
    Evaluate all 2-model and 3-model ensembles.
    Returns: dict of {ensemble_name: {metric: value, weights: {...}}}
    """
    kaggle_dir = Path(kaggle_data)
    val_sample_set = set(VAL_SAMPLE_IDS)
    
    # Collect samples
    all_samples = []
    for sample_dir in kaggle_dir.iterdir():
        if not sample_dir.is_dir() or 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))
                break
    
    print(f"Found {len(all_samples)} validation samples for ensemble evaluation")
    
    # Get all predictions first
    all_preds = {name: [] for name in models.keys()}
    all_gts = []
    
    print("Collecting predictions...")
    for img_path, csv_path in tqdm(all_samples, desc="Collecting"):
        try:
            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)
            gt_leads = get_gt_leads(df)
            all_gts.append(gt_leads)
            
            row_crops = [crop_row(image, i) for i in range(4)]
            
            if 'v16' in models and models['v16'] is not None:
                series = predict_v16(models['v16'], row_crops, device)
                all_preds['v16'].append(series_to_leads(np.clip(series, ECG_MV_MIN, ECG_MV_MAX)))
            
            if 'v16v18' in models and models['v16v18'] is not None:
                v16_m, v18_m = models['v16v18']
                series = predict_v16v18(v16_m, v18_m, row_crops, device)
                all_preds['v16v18'].append(series_to_leads(np.clip(series, ECG_MV_MIN, ECG_MV_MAX)))
            
            if 'v19' in models and models['v19'] is not None:
                series = predict_v19(models['v19'], row_crops, device)
                all_preds['v19'].append(series_to_leads(np.clip(series, ECG_MV_MIN, ECG_MV_MAX)))
            
            if 'net3' in models and models['net3'] is not None:
                series = predict_net3(models['net3'], image, device)
                if series is not None:
                    all_preds['net3'].append(series_to_leads(np.clip(series, ECG_MV_MIN, ECG_MV_MAX)))
        
        except Exception as e:
            print(f"Error: {e}")
            continue
    
    # Available models
    available_models = [name for name in models.keys() if len(all_preds[name]) > 0]
    print(f"Available models: {available_models}")
    
    # Test 2-model ensembles
    print("\nTesting 2-model ensembles...")
    best_ensemble_per_metric = {metric: {'score': None, 'models': None, 'weights': None}
                                 for metric in ['mae', 'nmse', 'si_snr', 'snr', 'log_cosh']}
    
    all_ensemble_results = []
    
    for m1, m2 in combinations(available_models, 2):
        for w1 in np.linspace(0.1, 0.9, weight_steps):
            w2 = 1.0 - w1
            weights = {m1: w1, m2: w2}
            
            # Compute ensemble metrics
            metrics = {'mae': [], 'nmse': [], 'si_snr': [], 'snr': [], 'log_cosh': []}
            
            for i, gt_leads in enumerate(all_gts):
                if i >= len(all_preds[m1]) or i >= len(all_preds[m2]):
                    continue
                
                signals = {m1: all_preds[m1][i], m2: all_preds[m2][i]}
                ensemble = ensemble_signals(signals, weights)
                
                for lead in ALL_LEADS:
                    if lead not in ensemble or lead not in gt_leads:
                        continue
                    
                    pred = ensemble[lead]
                    gt = gt_leads[lead]
                    pred_resampled = resample_signal(pred, len(gt))
                    
                    m = compute_all_metrics(pred_resampled, gt)
                    for metric_name, value in m.items():
                        metrics[metric_name].append(value)
            
            # Average
            avg_metrics = {k: np.mean(v) if v else None for k, v in metrics.items()}
            
            result = {
                'models': f"{m1}+{m2}",
                'weights': f"{w1:.1f}/{w2:.1f}",
                **avg_metrics
            }
            all_ensemble_results.append(result)
            
            # Update best
            for metric in ['mae', 'nmse', 'log_cosh']:  # Lower is better
                if avg_metrics[metric] is not None:
                    if best_ensemble_per_metric[metric]['score'] is None or avg_metrics[metric] < best_ensemble_per_metric[metric]['score']:
                        best_ensemble_per_metric[metric] = {'score': avg_metrics[metric], 'models': f"{m1}+{m2}", 'weights': weights}
            
            for metric in ['si_snr', 'snr']:  # Higher is better
                if avg_metrics[metric] is not None:
                    if best_ensemble_per_metric[metric]['score'] is None or avg_metrics[metric] > best_ensemble_per_metric[metric]['score']:
                        best_ensemble_per_metric[metric] = {'score': avg_metrics[metric], 'models': f"{m1}+{m2}", 'weights': weights}
    
    # Test 3-model ensembles
    print("Testing 3-model ensembles...")
    for m1, m2, m3 in combinations(available_models, 3):
        for w1 in np.linspace(0.2, 0.6, 5):
            for w2 in np.linspace(0.2, 0.6, 5):
                w3 = 1.0 - w1 - w2
                if w3 < 0.1:
                    continue
                
                weights = {m1: w1, m2: w2, m3: w3}
                
                metrics = {'mae': [], 'nmse': [], 'si_snr': [], 'snr': [], 'log_cosh': []}
                
                for i, gt_leads in enumerate(all_gts):
                    if i >= len(all_preds[m1]) or i >= len(all_preds[m2]) or i >= len(all_preds[m3]):
                        continue
                    
                    signals = {m1: all_preds[m1][i], m2: all_preds[m2][i], m3: all_preds[m3][i]}
                    ensemble = ensemble_signals(signals, weights)
                    
                    for lead in ALL_LEADS:
                        if lead not in ensemble or lead not in gt_leads:
                            continue
                        
                        pred = ensemble[lead]
                        gt = gt_leads[lead]
                        pred_resampled = resample_signal(pred, len(gt))
                        
                        m = compute_all_metrics(pred_resampled, gt)
                        for metric_name, value in m.items():
                            metrics[metric_name].append(value)
                
                avg_metrics = {k: np.mean(v) if v else None for k, v in metrics.items()}
                
                result = {
                    'models': f"{m1}+{m2}+{m3}",
                    'weights': f"{w1:.1f}/{w2:.1f}/{w3:.1f}",
                    **avg_metrics
                }
                all_ensemble_results.append(result)
                
                # Update best
                for metric in ['mae', 'nmse', 'log_cosh']:
                    if avg_metrics[metric] is not None:
                        if best_ensemble_per_metric[metric]['score'] is None or avg_metrics[metric] < best_ensemble_per_metric[metric]['score']:
                            best_ensemble_per_metric[metric] = {'score': avg_metrics[metric], 'models': f"{m1}+{m2}+{m3}", 'weights': weights}
                
                for metric in ['si_snr', 'snr']:
                    if avg_metrics[metric] is not None:
                        if best_ensemble_per_metric[metric]['score'] is None or avg_metrics[metric] > best_ensemble_per_metric[metric]['score']:
                            best_ensemble_per_metric[metric] = {'score': avg_metrics[metric], 'models': f"{m1}+{m2}+{m3}", 'weights': weights}
    
    return best_ensemble_per_metric, all_ensemble_results


def main():
    parser = argparse.ArgumentParser(description='Evaluate models and ensembles')
    parser.add_argument('--kaggle_data', type=str, default='/data/ecg-digitization/stage1_data/train')
    parser.add_argument('--v16_ckpt', type=str, default='/data/ecg-digitization/checkpoints/v16_perlead_latest.pth')
    parser.add_argument('--v18_ckpt', type=str, default='/data/ecg-digitization/checkpoints/v18_refiner_best.pth')
    parser.add_argument('--v19_ckpt', type=str, default='/data/ecg-digitization/checkpoints/v19_enhanced_epoch025.pth')
    parser.add_argument('--net3_ckpt', type=str, default='/data/ecg-digitization/hengck23-submit-physionet/hengck23-submit-physionet/weight/stage2-00005810.checkpoint.pth')
    parser.add_argument('--output', type=str, default='/home/azureuser/tmp/ensemble_results.json')
    args = parser.parse_args()
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}")
    
    # Load models
    print("\n" + "="*70)
    print("Loading models...")
    print("="*70)
    
    models = {}
    
    # V16
    if os.path.exists(args.v16_ckpt):
        models['v16'] = load_v16(args.v16_ckpt, device)
    else:
        print(f"V16 checkpoint not found: {args.v16_ckpt}")
        models['v16'] = None
    
    # V18 (requires V16)
    if os.path.exists(args.v18_ckpt) and models.get('v16') is not None:
        v18 = load_v18(args.v18_ckpt, device)
        models['v16v18'] = (models['v16'], v18)
    else:
        print(f"V18 checkpoint not found or V16 missing")
        models['v16v18'] = None
    
    # V19
    if os.path.exists(args.v19_ckpt):
        models['v19'] = load_v19(args.v19_ckpt, device)
    else:
        print(f"V19 checkpoint not found: {args.v19_ckpt}")
        models['v19'] = None
    
    # Net3
    if HAS_BASELINE and os.path.exists(args.net3_ckpt):
        models['net3'] = load_net3(args.net3_ckpt, device)
    else:
        print(f"Net3 not available or checkpoint not found")
        models['net3'] = None
    
    # Evaluate individual models
    print("\n" + "="*70)
    print("Evaluating individual models...")
    print("="*70)
    
    individual_results = evaluate_on_validation(models, args.kaggle_data, device)
    
    print("\n" + "="*70)
    print("INDIVIDUAL MODEL RESULTS")
    print("="*70)
    print(f"{'Model':<12} {'MAE':>10} {'NMSE':>10} {'SI-SNR':>10} {'SNR':>10} {'LogCosh':>10}")
    print("-"*70)
    
    for model_name, metrics in individual_results.items():
        if metrics['mae'] is not None:
            print(f"{model_name:<12} {metrics['mae']:>10.4f} {metrics['nmse']:>10.4f} "
                  f"{metrics['si_snr']:>10.2f} {metrics['snr']:>10.2f} {metrics['log_cosh']:>10.4f}")
    
    # Evaluate ensembles
    print("\n" + "="*70)
    print("Evaluating ensembles...")
    print("="*70)
    
    best_ensembles, all_results = evaluate_ensembles(models, args.kaggle_data, device)
    
    print("\n" + "="*70)
    print("BEST ENSEMBLE FOR EACH METRIC")
    print("="*70)
    
    for metric, info in best_ensembles.items():
        if info['score'] is not None:
            weights_str = '/'.join([f"{v:.1f}" for v in info['weights'].values()])
            print(f"{metric.upper():<10}: {info['models']:<20} ({weights_str}) = {info['score']:.4f}")
    
    # Save results
    output = {
        'individual': individual_results,
        'best_ensembles': {k: {'score': v['score'], 'models': v['models'], 
                               'weights': {m: float(w) for m, w in v['weights'].items()}}
                           for k, v in best_ensembles.items() if v['score'] is not None},
    }
    
    with open(args.output, 'w') as f:
        json.dump(output, f, indent=2)
    
    print(f"\nResults saved to: {args.output}")


if __name__ == '__main__':
    main()