#!/usr/bin/env python3
"""
Evaluate individual models and ensembles on validation set using preprocessed images.
Uses preprocessed stage1 images directly (no stage0/1 processing).

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 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 for Net3
BASELINE_PATH = '/data/ecg-digitization/hengck23-submit-physionet/hengck23-submit-physionet'
sys.path.insert(0, BASELINE_PATH)

try:
    import stage2_common as s2c
    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  # For V16/V18/V19
MV_TO_PIXEL_NET3 = 78.8  # Net3 uses slightly different value
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

# Validation sample IDs (holdout set)
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."""
    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 - Matches Kaggle checkpoint
# =============================================================================
if HAS_BASELINE:
    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)


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


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]
    return leads


def get_gt_leads(df):
    """
    Get ground truth leads from CSV.
    
    IMPORTANT: The GT CSV structure is:
    - Most leads (I, III, aVR, aVL, aVF, V1-V6): Already 2.5s segments (~2562 samples)
    - Lead II: Full 10 seconds (~10250 samples)
    
    For the SHORT Lead II strip (row 1, segment 0), use first 2.5s of GT.
    For rhythm strip comparison, we would need full II, but we're only comparing short leads.
    """
    leads = {}
    
    # Get reference length for short leads
    ref_lead = 'I' if 'I' in df.columns else 'III'
    ref_len = len(df[ref_lead].dropna().values) if ref_lead in df.columns else 2500
    
    for lead in ALL_LEADS:
        if lead in df.columns:
            signal = df[lead].dropna().values.astype(np.float32)
            
            # Special handling for Lead II (10s data in GT)
            # The SHORT strip shows first 2.5s (0-2.5s)
            if lead == 'II' and len(signal) > ref_len * 2:
                # Lead II is at least 2x longer = it's 10s data
                # Use first quarter for SHORT strip comparison
                quarter_len = len(signal) // 4
                leads[lead] = signal[:quarter_len]
            else:
                leads[lead] = signal
    
    return leads


# =============================================================================
# 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 = Net3(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
    # Filter out unexpected keys like D, mean, std
    model_keys = set(model.state_dict().keys())
    state_dict = {k: v for k, v in state_dict.items() if k in model_keys}
    model.load_state_dict(state_dict, strict=False)
    model = model.to(device).eval()
    print(f"Net3 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 ALREADY preprocessed image, return 4-row mV signal.
    
    IMPORTANT: image_rgb should already be:
    - Cropped to [0:1696, 0:2176] from stage1 output
    - Resized to (1696, 4352)
    
    So we just normalize and run through the model.
    """
    if model is None:
        return None
    
    # Image is already preprocessed to (1696, 4352, 3)
    # Just normalize to [0, 1] and run
    img = image_rgb.astype(np.float32) / 255.0
    batch = torch.from_numpy(np.ascontiguousarray(img.transpose(2, 0, 1))).unsqueeze(0).float()
    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] = [4, 1696, 4352]
    
    # Use pixel_to_series from baseline to convert heatmap to pixel coords
    # Slice to signal region T0:T1 (same as reference)
    series_in_pixel = s2c.pixel_to_series(pixel[..., T0:T1], ZERO_MV, length=OUTPUT_WIDTH)
    # Convert to mV using Net3's MV_TO_PIXEL value (78.8, not 78.5!)
    series_mv = (ZERO_MV.reshape(4, 1) - series_in_pixel) / MV_TO_PIXEL_NET3
    return series_mv


# =============================================================================
# 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())
    if not model_names:
        return {}
    
    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 collect_samples(stage1_dir):
    """
    Collect all validation samples.
    Uses preprocessed images and GT from stage1_dir.
    Structure: stage1_dir/{sample_id}/{sample_id}-{variant}.png and {sample_id}.csv
    """
    stage1_path = Path(stage1_dir)
    val_sample_set = set(VAL_SAMPLE_IDS)
    
    samples = []
    
    # Iterate through sample folders
    for sample_folder in stage1_path.iterdir():
        if not sample_folder.is_dir():
            continue
        
        sample_id = sample_folder.name
        
        if sample_id not in val_sample_set:
            continue
        
        # Find GT CSV in the same folder
        csv_path = sample_folder / f"{sample_id}.csv"
        if not csv_path.exists():
            continue
        
        # Find all variant images
        for img_path in sample_folder.glob(f"{sample_id}-*.png"):
            # Format: 1006427285-0001.png
            name = img_path.stem
            parts = name.split('-')
            if len(parts) != 2:
                continue
            
            variant = parts[1]
            
            if variant not in VALID_VARIANTS:
                continue
            
            samples.append({
                'image_path': img_path,
                'csv_path': csv_path,
                'sample_id': sample_id,
                'variant': variant,
            })
    
    return samples


def evaluate_models(models, samples, device):
    """
    Evaluate all models on samples.
    Returns predictions and ground truth for each sample.
    """
    all_preds = {name: [] for name in models.keys()}
    all_gts = []
    sample_ids = []
    
    for sample in tqdm(samples, desc="Evaluating"):
        try:
            # Load preprocessed image (already stage0/1 processed)
            image = cv2.imread(str(sample['image_path']), cv2.IMREAD_COLOR)
            if image is None:
                continue
            
            # Preprocess: crop then resize (same as infer_v16)
            # stage1_data images are ~(1700, 2200), need to crop to (Y1, X1) then resize
            h, w = image.shape[:2]
            crop_h = min(h, Y1)  # min(1700, 1696) = 1696
            crop_w = min(w, X1)  # min(2200, 2176) = 2176
            image = image[:crop_h, :crop_w]
            image = cv2.resize(image, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR)
            
            # Load GT
            df = pd.read_csv(sample['csv_path'])
            gt_leads = get_gt_leads(df)
            if not gt_leads:
                continue
            
            all_gts.append(gt_leads)
            sample_ids.append(f"{sample['sample_id']}-{sample['variant']}")
            
            # Crop rows for V16/V18/V19
            row_crops = [crop_row(image, i) for i in range(4)]
            
            # Get predictions from each model
            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 'v18' in models and models['v18'] is not None:
                v16_m, v18_m = models['v18']
                series = predict_v16v18(v16_m, v18_m, row_crops, device)
                all_preds['v18'].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:
                image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                series = predict_net3(models['net3'], image_rgb, 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: {sample['image_path']}: {e}")
            continue
    
    return all_preds, all_gts, sample_ids


def compute_metrics_for_predictions(all_preds, all_gts, model_name):
    """Compute metrics for a single model's predictions."""
    metrics = {'mae': [], 'nmse': [], 'si_snr': [], 'snr': [], 'log_cosh': []}
    
    preds_list = all_preds.get(model_name, [])
    if not preds_list:
        return None
    
    for i, gt_leads in enumerate(all_gts):
        if i >= len(preds_list):
            break
        
        pred_leads = preds_list[i]
        
        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
            pred_resampled = resample_signal(pred, len(gt))
            
            # Compute metrics
            m = compute_all_metrics(pred_resampled, gt)
            for metric_name, value in m.items():
                metrics[metric_name].append(value)
    
    # Average
    return {k: np.mean(v) if v else None for k, v in metrics.items()}


def compute_metrics_for_ensemble(all_preds, all_gts, model_names, weights):
    """Compute metrics for an ensemble."""
    metrics = {'mae': [], 'nmse': [], 'si_snr': [], 'snr': [], 'log_cosh': []}
    
    for i, gt_leads in enumerate(all_gts):
        # Check all models have predictions for this sample
        valid = True
        for name in model_names:
            if i >= len(all_preds.get(name, [])):
                valid = False
                break
        
        if not valid:
            continue
        
        # Build signals dict for this sample
        signals = {name: all_preds[name][i] for name in model_names}
        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)
    
    return {k: np.mean(v) if v else None for k, v in metrics.items()}


def main():
    parser = argparse.ArgumentParser(description='Evaluate models and ensembles on preprocessed images')
    parser.add_argument('--stage1_dir', 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_augraphy_latest.pth')
    parser.add_argument('--net3_ckpt', type=str, default='/data/ecg-digitization/checkpoints/net3_kaggle/iter_0004200.pt')
    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}")
    
    # Collect samples
    print("\n" + "="*70)
    print("Collecting validation samples...")
    print("="*70)
    
    samples = collect_samples(args.stage1_dir)
    print(f"Found {len(samples)} validation samples")
    
    if len(samples) == 0:
        print("ERROR: No samples found!")
        return
    
    # Load models
    print("\n" + "="*70)
    print("Loading models...")
    print("="*70)
    
    models = {}
    
    # V16
    if os.path.exists(args.v16_ckpt):
        v16_model = load_v16(args.v16_ckpt, device)
        models['v16'] = v16_model
    else:
        print(f"V16 checkpoint not found: {args.v16_ckpt}")
    
    # V18 (requires V16)
    if os.path.exists(args.v18_ckpt) and 'v16' in models:
        v18_model = load_v18(args.v18_ckpt, device)
        models['v18'] = (models['v16'], v18_model)
    else:
        print(f"V18 checkpoint not found or V16 missing")
    
    # 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}")
    
    # # 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")
    
    print(f"\nLoaded models: {list(models.keys())}")
    
    # Evaluate
    print("\n" + "="*70)
    print("Running inference...")
    print("="*70)
    
    all_preds, all_gts, sample_ids = evaluate_models(models, samples, device)
    
    print(f"\nProcessed {len(all_gts)} samples")
    
    # Compute individual metrics
    print("\n" + "="*70)
    print("INDIVIDUAL MODEL RESULTS")
    print("="*70)
    
    individual_results = {}
    print(f"\n{'Model':<12} {'MAE':>10} {'NMSE':>10} {'SI-SNR':>10} {'SNR':>10} {'LogCosh':>10}")
    print("-"*70)
    
    for model_name in ['v16', 'v18', 'v19', 'net3']:
        if model_name in models:
            metrics = compute_metrics_for_predictions(all_preds, all_gts, model_name)
            if metrics and metrics['mae'] is not None:
                individual_results[model_name] = metrics
                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}")
    
    # Find best ensemble for each metric
    print("\n" + "="*70)
    print("SEARCHING FOR BEST ENSEMBLES...")
    print("="*70)
    
    available_models = [name for name in models.keys() if len(all_preds.get(name, [])) > 0]
    print(f"Available models for ensemble: {available_models}")
    
    best_ensembles = {
        'mae': {'score': None, 'models': None, 'weights': None},
        'nmse': {'score': None, 'models': None, 'weights': None},
        'si_snr': {'score': None, 'models': None, 'weights': None},
        'snr': {'score': None, 'models': None, 'weights': None},
        'log_cosh': {'score': None, 'models': None, 'weights': None},
    }
    
    all_ensemble_results = []
    
    # Test 2-model ensembles
    print("\nTesting 2-model ensembles...")
    for m1, m2 in combinations(available_models, 2):
        for w1 in np.linspace(0.1, 0.9, 9):
            w2 = 1.0 - w1
            weights = {m1: w1, m2: w2}
            
            metrics = compute_metrics_for_ensemble(all_preds, all_gts, [m1, m2], weights)
            
            if metrics['mae'] is None:
                continue
            
            result = {
                'models': f"{m1}+{m2}",
                'weights': {m1: float(w1), m2: float(w2)},
                **metrics
            }
            all_ensemble_results.append(result)
            
            # Update best (lower is better for mae, nmse, log_cosh)
            for metric in ['mae', 'nmse', 'log_cosh']:
                if best_ensembles[metric]['score'] is None or metrics[metric] < best_ensembles[metric]['score']:
                    best_ensembles[metric] = {'score': metrics[metric], 'models': f"{m1}+{m2}", 'weights': weights.copy()}
            
            # Higher is better for si_snr, snr
            for metric in ['si_snr', 'snr']:
                if best_ensembles[metric]['score'] is None or metrics[metric] > best_ensembles[metric]['score']:
                    best_ensembles[metric] = {'score': metrics[metric], 'models': f"{m1}+{m2}", 'weights': weights.copy()}
    
    # Test 3-model ensembles
    if len(available_models) >= 3:
        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 = compute_metrics_for_ensemble(all_preds, all_gts, [m1, m2, m3], weights)
                    
                    if metrics['mae'] is None:
                        continue
                    
                    result = {
                        'models': f"{m1}+{m2}+{m3}",
                        'weights': {m1: float(w1), m2: float(w2), m3: float(w3)},
                        **metrics
                    }
                    all_ensemble_results.append(result)
                    
                    for metric in ['mae', 'nmse', 'log_cosh']:
                        if best_ensembles[metric]['score'] is None or metrics[metric] < best_ensembles[metric]['score']:
                            best_ensembles[metric] = {'score': metrics[metric], 'models': f"{m1}+{m2}+{m3}", 'weights': weights.copy()}
                    
                    for metric in ['si_snr', 'snr']:
                        if best_ensembles[metric]['score'] is None or metrics[metric] > best_ensembles[metric]['score']:
                            best_ensembles[metric] = {'score': metrics[metric], 'models': f"{m1}+{m2}+{m3}", 'weights': weights.copy()}
    
    # Test 4-model ensembles
    if len(available_models) >= 4:
        print("Testing 4-model ensembles...")
        for w1 in np.linspace(0.2, 0.4, 3):
            for w2 in np.linspace(0.2, 0.4, 3):
                for w3 in np.linspace(0.1, 0.3, 3):
                    w4 = 1.0 - w1 - w2 - w3
                    if w4 < 0.1:
                        continue
                    
                    weights = {available_models[0]: w1, available_models[1]: w2, 
                              available_models[2]: w3, available_models[3]: w4}
                    metrics = compute_metrics_for_ensemble(all_preds, all_gts, available_models, weights)
                    
                    if metrics['mae'] is None:
                        continue
                    
                    result = {
                        'models': '+'.join(available_models),
                        'weights': {k: float(v) for k, v in weights.items()},
                        **metrics
                    }
                    all_ensemble_results.append(result)
                    
                    for metric in ['mae', 'nmse', 'log_cosh']:
                        if best_ensembles[metric]['score'] is None or metrics[metric] < best_ensembles[metric]['score']:
                            best_ensembles[metric] = {'score': metrics[metric], 'models': '+'.join(available_models), 'weights': weights.copy()}
                    
                    for metric in ['si_snr', 'snr']:
                        if best_ensembles[metric]['score'] is None or metrics[metric] > best_ensembles[metric]['score']:
                            best_ensembles[metric] = {'score': metrics[metric], 'models': '+'.join(available_models), 'weights': weights.copy()}
    
    # Print results
    print("\n" + "="*70)
    print("BEST ENSEMBLE FOR EACH METRIC")
    print("="*70)
    
    for metric in ['mae', 'nmse', 'si_snr', 'snr', 'log_cosh']:
        info = best_ensembles[metric]
        if info['score'] is not None:
            weights_str = ', '.join([f"{k}:{v:.2f}" for k, v in info['weights'].items()])
            print(f"{metric.upper():<10}: {info['models']:<25} ({weights_str}) = {info['score']:.4f}")
    
    # Print ALL ensemble results sorted by SNR (most useful metric)
    print("\n" + "="*70)
    print("ALL ENSEMBLE RESULTS (sorted by SNR, descending)")
    print("="*70)
    
    # Sort by SNR descending
    sorted_by_snr = sorted(all_ensemble_results, key=lambda x: x['snr'] if x['snr'] is not None else -999, reverse=True)
    
    print(f"{'Models':<30} {'Weights':<40} {'MAE':>8} {'NMSE':>8} {'SNR':>8} {'SI-SNR':>8} {'LogCosh':>8}")
    print("-" * 110)
    
    for result in sorted_by_snr[:50]:  # Top 50 by SNR
        weights_str = '/'.join([f"{result['weights'].get(m, 0):.2f}" for m in sorted(result['weights'].keys())])
        models_short = result['models'].replace('+', '/')
        print(f"{models_short:<30} {weights_str:<40} {result['mae']:>8.4f} {result['nmse']:>8.2f} {result['snr']:>8.2f} {result['si_snr']:>8.2f} {result['log_cosh']:>8.4f}")
    
    # Also show top 10 by MAE
    print("\n" + "="*70)
    print("TOP 10 BY MAE (ascending)")
    print("="*70)
    sorted_by_mae = sorted(all_ensemble_results, key=lambda x: x['mae'] if x['mae'] is not None else 999)
    print(f"{'Models':<30} {'Weights':<40} {'MAE':>8} {'SNR':>8}")
    print("-" * 90)
    for result in sorted_by_mae[:10]:
        weights_str = '/'.join([f"{result['weights'].get(m, 0):.2f}" for m in sorted(result['weights'].keys())])
        models_short = result['models'].replace('+', '/')
        print(f"{models_short:<30} {weights_str:<40} {result['mae']:>8.4f} {result['snr']:>8.2f}")
    
    # 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},
        'all_ensemble_results': all_ensemble_results,
    }
    
    os.makedirs(os.path.dirname(args.output), exist_ok=True)
    with open(args.output, 'w') as f:
        json.dump(output, f, indent=2)
    
    print(f"\nResults saved to: {args.output}")


if __name__ == '__main__':
    main()