#!/usr/bin/env python3
"""
Find optimal ensemble ratio between V10.1 model and Friend's Net3 model.
Uses the EXACT same SNR calculation as validate_v10_variants.ipynb:
- Extract 12 individual leads from 4-row prediction
- Resample GT to match prediction length
- Compute SNR per lead: 10 * log10(sum(gt^2) / sum((gt-pred)^2))
- Average all 12 lead SNRs
"""

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
import random

# 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
ECG_MV_MIN, ECG_MV_MAX = -7.0, 7.0

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

# Lead layout - EXACTLY matches validate_v10_variants.ipynb
ROW_LAYOUT = [
    ['I', 'aVR', 'V1', 'V4'],
    ['II_short', 'aVL', 'V2', 'V5'],  # II_short not used, we use rhythm strip for II
    ['III', 'aVF', 'V3', 'V6'],
]


# =============================================================================
# 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 - EXACTLY as in validate_v10_variants.ipynb
# =============================================================================
def soft_argmax(heatmap, temperature=SOFT_ARGMAX_TEMP):
    """Extract sub-pixel coordinates using soft-argmax."""
    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 interpolate_nan(signal_1d):
    """Interpolate NaN values from valid neighbors. Falls back to 0 if all NaN."""
    valid_mask = np.isfinite(signal_1d)
    if valid_mask.all():
        return signal_1d
    if not valid_mask.any():
        return np.zeros_like(signal_1d)
    x = np.arange(len(signal_1d))
    signal_1d[~valid_mask] = np.interp(x[~valid_mask], x[valid_mask], signal_1d[valid_mask])
    return signal_1d


def 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 compute_snr(pred_mv, gt_mv):
    """Compute SNR in dB between prediction and ground truth.
    
    SNR = 10 * log10(sum(gt^2) / sum((gt - pred)^2))
    
    EXACTLY as in validate_v10_variants.ipynb
    """
    # Ensure same length
    min_len = min(len(pred_mv), len(gt_mv))
    pred = pred_mv[:min_len]
    gt = gt_mv[:min_len]
    
    sig_power = np.sum(gt ** 2)
    noise_power = np.sum((gt - pred) ** 2)
    
    if noise_power < 1e-10:
        return 50.0  # Perfect match
    if sig_power < 1e-10:
        return 0.0  # No signal
    
    snr = 10 * np.log10(sig_power / noise_power)
    return np.clip(snr, -20, 50)


def extract_leads_from_prediction(pred_mv):
    """Extract 12 leads from 4-row prediction - matches Kaggle inference exactly.
    
    EXACTLY as in validate_v10_variants.ipynb
    """
    segment_width = pred_mv.shape[1] // 4
    leads = {}
    
    # Rows 0-2: short leads (4 segments each)
    for row_idx in range(3):
        for seg_idx, lead_name in enumerate(ROW_LAYOUT[row_idx]):
            if lead_name == 'II_short':
                continue  # Skip, we use rhythm strip for II
            start = seg_idx * segment_width
            end = (seg_idx + 1) * segment_width
            leads[lead_name] = pred_mv[row_idx, start:end]
    
    # Row 3: Full rhythm strip (lead II) - THIS IS WHAT KAGGLE USES FOR II
    leads['II'] = pred_mv[3]
    
    return leads


def extract_leads_from_gt(csv_path):
    """Extract 12 leads from ground truth CSV - matches Kaggle format.
    
    EXACTLY as in validate_v10_variants.ipynb
    """
    df = pd.read_csv(csv_path)
    leads = {}
    
    # All 12 standard leads
    for lead in ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']:
        if lead in df.columns:
            leads[lead] = df[lead].dropna().values
    
    return leads


def compute_sample_snr(pred_mv, csv_path):
    """Compute per-lead SNR and return average - matches Kaggle evaluation.
    
    Kaggle computes SNR for each lead separately, then averages.
    Lead II uses the RHYTHM STRIP (row 3), not the short segment.
    
    EXACTLY as in validate_v10_variants.ipynb
    """
    pred_leads = extract_leads_from_prediction(pred_mv)
    gt_leads = extract_leads_from_gt(csv_path)
    
    lead_snrs = []
    for lead in ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']:
        if lead not in pred_leads or lead not in gt_leads:
            continue
        
        pred = pred_leads[lead]
        gt = gt_leads[lead]
        
        # Resample GT to match prediction length
        gt_resampled = resample_signal(gt, len(pred))
        
        snr = compute_snr(pred, gt_resampled)
        if np.isfinite(snr):
            lead_snrs.append(snr)
    
    if lead_snrs:
        return np.mean(lead_snrs)
    return np.nan


# =============================================================================
# Inference Functions
# =============================================================================
@torch.no_grad()
def process_v10(model, image_bgr, device):
    """V10 inference on stage1 image. Returns 4-row mV signal.
    
    EXACTLY as in validate_v10_variants.ipynb: process_stage1_image
    """
    h, w = image_bgr.shape[:2]
    
    # Crop to left region
    crop_h = min(h, Y1)
    crop_w = min(w, X1)
    image_cropped = image_bgr[:crop_h, :crop_w]
    
    # Resize to model input size
    image_resized = cv2.resize(image_cropped, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR)
    
    # Normalize to [0, 1]
    image_tensor = torch.from_numpy(image_resized.astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0)
    image_tensor = image_tensor.to(device)
    
    with torch.amp.autocast('cuda', dtype=torch.float32):
        seg_logits, reg_coords = model(image_tensor)
    
    # Handle NaN/Inf in model output
    seg_has_nan = torch.isnan(seg_logits).any() or torch.isinf(seg_logits).any()
    if seg_has_nan:
        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]  # [4, 4352]
    
    # Handle NaN in soft_argmax output
    for row_idx in range(4):
        if not np.isfinite(signal_full[row_idx]).all():
            # Try regression head as fallback
            reg_signal = reg_coords.float().cpu().numpy()[0] * (TARGET_HEIGHT - 1)
            nan_mask = ~np.isfinite(signal_full[row_idx])
            if np.isfinite(reg_signal[row_idx]).all():
                signal_full[row_idx, nan_mask] = reg_signal[row_idx, nan_mask]
            else:
                # Interpolate from neighbors
                signal_full[row_idx] = interpolate_nan(signal_full[row_idx].copy())
                remaining_nan = ~np.isfinite(signal_full[row_idx])
                if remaining_nan.any():
                    signal_full[row_idx, remaining_nan] = ZERO_MV[row_idx]
    
    # Extract signal region T0:T1
    signal_pixel = signal_full[:, T0:T1]  # [4, OUTPUT_WIDTH]
    
    # Convert to mV
    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
    signal_mv = np.clip(signal_mv, ECG_MV_MIN, ECG_MV_MAX)
    
    return signal_mv


@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 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("Using EXACT same SNR calculation as validate_v10_variants.ipynb")
    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}")
            return None
    
    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}")
            return None
    
    # 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 and compute SNRs
    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 sig_len from GT
            gt_df = pd.read_csv(csv_path)
            sig_len = len(gt_df['II'].dropna()) if 'II' in gt_df.columns else 5000
            
            # 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)
            
            results.append({
                'sample_id': sample_id,
                'variant': variant,
                'csv_path': csv_path,
                'pred_v10': pred_v10,
                'pred_friend': pred_net3,
            })
            
        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 None
    
    # Evaluate different ratios
    print("\n" + "="*70)
    print("Evaluating ensemble ratios...")
    print("SNR = average of per-lead SNRs (12 leads)")
    print("="*70)
    
    ratios = np.arange(0.0, 1.05, 0.05)
    ratio_results = []
    
    for ratio in ratios:
        snrs = []
        for r in results:
            ensemble = ensemble_predictions(r['pred_v10'], r['pred_friend'], ratio)
            snr = compute_sample_snr(ensemble, r['csv_path'])
            if np.isfinite(snr):
                snrs.append(snr)
        
        avg_snr = np.mean(snrs) if snrs else 0.0
        ratio_results.append((ratio, avg_snr))
        print(f"V10 weight: {ratio:.2f}, Friend weight: {1-ratio:.2f} -> SNR: {avg_snr:.2f} dB")
    
    # Find best ratio
    best_ratio, best_snr = max(ratio_results, key=lambda x: x[1])
    
    # Compute individual model SNRs
    v10_snrs = []
    friend_snrs = []
    for r in results:
        v10_snr = compute_sample_snr(r['pred_v10'], r['csv_path'])
        friend_snr = compute_sample_snr(r['pred_friend'], r['csv_path'])
        if np.isfinite(v10_snr):
            v10_snrs.append(v10_snr)
        if np.isfinite(friend_snr):
            friend_snrs.append(friend_snr)
    
    v10_avg = np.mean(v10_snrs) if v10_snrs else 0.0
    friend_avg = np.mean(friend_snrs) if friend_snrs else 0.0
    
    print("\n" + "="*70)
    print(f"BEST RATIO: V10={best_ratio:.2f}, Friend={1-best_ratio:.2f}")
    print(f"Best SNR: {best_snr:.2f} dB")
    print("="*70)
    
    print(f"\nV10 only:    {v10_avg:.2f} dB")
    print(f"Friend only: {friend_avg:.2f} dB")
    print(f"Ensemble:    {best_snr:.2f} dB (improvement: +{best_snr - max(v10_avg, friend_avg):.2f} 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_var_snrs = []
            friend_var_snrs = []
            ensemble_var_snrs = []
            for r in variant_results:
                v10_snr = compute_sample_snr(r['pred_v10'], r['csv_path'])
                friend_snr = compute_sample_snr(r['pred_friend'], r['csv_path'])
                ensemble = ensemble_predictions(r['pred_v10'], r['pred_friend'], best_ratio)
                ensemble_snr = compute_sample_snr(ensemble, r['csv_path'])
                
                if np.isfinite(v10_snr):
                    v10_var_snrs.append(v10_snr)
                if np.isfinite(friend_snr):
                    friend_var_snrs.append(friend_snr)
                if np.isfinite(ensemble_snr):
                    ensemble_var_snrs.append(ensemble_snr)
            
            v10_var = np.mean(v10_var_snrs) if v10_var_snrs else 0.0
            friend_var = np.mean(friend_var_snrs) if friend_var_snrs else 0.0
            ensemble_var = np.mean(ensemble_var_snrs) if ensemble_var_snrs else 0.0
            print(f"  {variant}: V10={v10_var:.2f}, Friend={friend_var:.2f}, Ensemble={ensemble_var:.2f} dB (n={len(variant_results)})")
    
    return best_ratio


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