#!/usr/bin/env python3
"""
Compute per-segment baseline offsets for V18 model.

This analyzes what the typical baseline offset is between V18 predictions
and ground truth for each lead.
"""

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

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

import cv2
import timm


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

CROP_HALF_HEIGHT = 250
ROW_HEIGHT = 500

ECG_MV_MIN, ECG_MV_MAX = -10.0, 10.0

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

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

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


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


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


class PerLeadNetV16(nn.Module):
    """V16 Per-Lead ECG Network (frozen, used as prior)."""
    
    def __init__(self, encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=True):
        super().__init__()
        
        self.encoder = timm.create_model(
            encoder_name, pretrained=pretrained, features_only=True, out_indices=(0, 1, 2, 3),
        )
        enc_channels = self.encoder.feature_info.channels()
        
        decoder_dims = [256, 128, 64, 32]
        self.dec_blocks = nn.ModuleList()
        in_ch = enc_channels[-1]
        skip_channels = enc_channels[:-1][::-1] + [0]
        
        for skip_ch, out_ch in zip(skip_channels, decoder_dims):
            self.dec_blocks.append(UNetDecoderBlock(in_ch, skip_ch, out_ch))
            in_ch = out_ch
        
        self.final_up = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            nn.Conv2d(decoder_dims[-1], decoder_dims[-1], 3, padding=1, bias=False),
            nn.BatchNorm2d(decoder_dims[-1]), nn.GELU(),
        )
        
        self.height_attention = nn.Sequential(
            CoordConv2d(decoder_dims[-1], 64, 3, padding=1),
            nn.BatchNorm2d(64), nn.GELU(),
            nn.Conv2d(64, 1, 1),
        )
        
        self.regression_head = nn.Sequential(
            nn.Conv1d(decoder_dims[-1], 128, 7, padding=3),
            nn.BatchNorm1d(128), nn.GELU(),
            nn.Conv1d(128, 64, 5, padding=2),
            nn.BatchNorm1d(64), nn.GELU(),
            nn.Conv1d(64, 1, 1),
        )
    
    def forward(self, x):
        B, C, H, W = x.shape
        features = self.encoder(x)
        d = features[-1]
        skips = features[:-1][::-1] + [None]
        for block, skip in zip(self.dec_blocks, skips):
            d = block(d, skip)
        d = self.final_up(d)
        if d.shape[3] != W:
            d = F.interpolate(d, size=(d.shape[2], W), mode='bilinear', align_corners=True)
        attn = self.height_attention(d)
        attn = F.softmax(attn, dim=2)
        d = (d * attn).sum(dim=2)
        out = self.regression_head(d)
        out = torch.sigmoid(out)
        return out.squeeze(1)


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


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)
        x = x.reshape(B, W, num_rows, C).permute(0, 2, 3, 1)
        return x


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(embed_dim=cross_row_dim, num_heads=num_heads, mlp_ratio=2.0, dropout=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)
        y_grid = y_grid.view(1, height, 1)
        y_pred = y_pred.unsqueeze(1)
        guide = torch.exp(-0.5 * ((y_grid - y_pred) / sigma) ** 2)
        guide = guide.unsqueeze(1)
        return guide
    
    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


# =============================================================================
# Helper Functions
# =============================================================================
def apply_savgol_smoothing(signal_mv, window=7, polyorder=2):
    if len(signal_mv) >= window:
        return savgol_filter(signal_mv, window_length=window, polyorder=polyorder)
    return signal_mv


def apply_einthoven_correction(pred_mv_rows, alpha=0.33):
    segment_width = len(pred_mv_rows[0]) // 4
    lead_I = pred_mv_rows[0][:segment_width].copy()
    lead_II_short = pred_mv_rows[1][:segment_width].copy()
    lead_III = pred_mv_rows[2][:segment_width].copy()
    derived_II = lead_I + lead_III
    error = lead_II_short - derived_II
    pred_mv_rows[0][:segment_width] = lead_I + alpha * error
    pred_mv_rows[2][:segment_width] = lead_III + alpha * error
    return pred_mv_rows


def clamp_ecg_amplitude(signal_mv):
    return np.clip(signal_mv, ECG_MV_MIN, ECG_MV_MAX)


def interpolate_nan(signal_1d):
    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 crop_row(image, row_idx):
    baseline_y = int(ZERO_MV[row_idx])
    y_start = max(0, baseline_y - CROP_HALF_HEIGHT)
    y_end = min(TARGET_HEIGHT, baseline_y + CROP_HALF_HEIGHT)
    row_crop = image[y_start:y_end, T0:T1, :].copy()
    if row_crop.shape[0] < ROW_HEIGHT:
        pad_top = max(0, CROP_HALF_HEIGHT - baseline_y)
        pad_bottom = max(0, (baseline_y + CROP_HALF_HEIGHT) - TARGET_HEIGHT)
        row_crop = np.pad(row_crop, ((pad_top, pad_bottom), (0, 0), (0, 0)), mode='edge')
    return row_crop


def convert_crop_to_full(pred_y_crop, row_idx):
    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 pred_y_full


def load_v16_model(checkpoint_path, device):
    model = PerLeadNetV16(encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=False)
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    state_dict = checkpoint['model']
    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)
    model.eval()
    print(f"Loaded V16 from epoch {checkpoint.get('epoch', 'N/A')}, SNR: {checkpoint.get('snr', 'N/A'):.2f} dB")
    return model


def load_v18_refiner(checkpoint_path, device):
    refiner = V18RefinerNet(encoder_name='efficientnet_b0', pretrained=False,
                            cross_row_layers=3, cross_row_dim=128, num_heads=4)
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    
    # Try different checkpoint formats
    if 'refiner' in checkpoint:
        state_dict = checkpoint['refiner']
    elif 'model' in checkpoint:
        state_dict = checkpoint['model']
    elif 'state_dict' in checkpoint:
        state_dict = checkpoint['state_dict']
    else:
        state_dict = checkpoint
    
    # Extract refiner weights only (they may be prefixed with 'refiner.')
    refiner_state_dict = {}
    for k, v in state_dict.items():
        if k.startswith('refiner.'):
            refiner_state_dict[k.replace('refiner.', '')] = v
        elif k.startswith('module.refiner.'):
            refiner_state_dict[k.replace('module.refiner.', '')] = v
        elif k.startswith('module.'):
            refiner_state_dict[k.replace('module.', '')] = v
        elif not k.startswith('v16.') and not k.startswith('module.v16.'):
            refiner_state_dict[k] = v
    
    refiner.load_state_dict(refiner_state_dict)
    refiner = refiner.to(device)
    refiner.eval()
    print(f"Loaded V18 refiner from epoch {checkpoint.get('epoch', 'N/A')}, SNR: {checkpoint.get('snr', 'N/A'):.2f} dB")
    return refiner


def compute_baseline_offset(pred_mv_segment, gt_mv):
    """Compute baseline offset between prediction and GT."""
    x_pred = np.linspace(0, 1, len(pred_mv_segment))
    x_gt = np.linspace(0, 1, len(gt_mv))
    pred_resampled = np.interp(x_gt, x_pred, pred_mv_segment)
    return np.median(pred_resampled - gt_mv)


def process_image(v16_model, v18_refiner, image, df, device):
    """
    Process a single image with V18 and return baseline offsets for all 12 leads.
    """
    segment_width = OUTPUT_WIDTH // 4
    
    # Crop all 4 rows
    row_crops = [crop_row(image, row_idx) for row_idx in range(4)]
    
    # Run V16 on all rows
    v16_preds = []
    for row_crop in row_crops:
        image_tensor = torch.from_numpy(row_crop.astype(np.float32) / 255.0)
        image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0).to(device)
        with torch.no_grad():
            with torch.amp.autocast('cuda'):
                pred = v16_model(image_tensor)
        v16_preds.append(pred)
    
    # Stack for V18 refiner
    v16_preds_tensor = torch.stack(v16_preds, dim=1)  # (1, 4, W)
    
    # Prepare images for refiner
    images_tensor = []
    for row_crop in row_crops:
        image_tensor = torch.from_numpy(row_crop.astype(np.float32) / 255.0)
        image_tensor = image_tensor.permute(2, 0, 1)
        images_tensor.append(image_tensor)
    images_tensor = torch.stack(images_tensor, dim=0).unsqueeze(0).to(device)  # (1, 4, 3, H, W)
    
    # Run V18 refiner
    with torch.no_grad():
        with torch.amp.autocast('cuda'):
            v18_preds, residuals = v18_refiner(images_tensor, v16_preds_tensor)
    
    # Convert V18 predictions to mV
    pred_mv_rows = {}
    for row_idx in range(3):  # Only rows 0-2 (not rhythm strip for leads)
        pred_normalized = v18_preds[0, row_idx].cpu().numpy()
        pred_y_crop = pred_normalized * ROW_HEIGHT
        pred_y_full = convert_crop_to_full(pred_y_crop, row_idx)
        
        baseline_y = ZERO_MV[row_idx]
        pred_mv = (baseline_y - pred_y_full) / MV_TO_PIXEL
        pred_mv = apply_savgol_smoothing(pred_mv, window=7, polyorder=2)
        pred_mv = clamp_ecg_amplitude(pred_mv)
        pred_mv = interpolate_nan(pred_mv.copy())
        
        pred_mv_rows[row_idx] = pred_mv
    
    # Apply Einthoven correction
    pred_mv_rows = apply_einthoven_correction(pred_mv_rows, alpha=0.33)
    
    # Compute baseline offset for each lead
    offsets = {}
    
    for row_idx in range(3):
        lead_names = LEAD_LAYOUT[row_idx]
        for seg_idx, lead_name in enumerate(lead_names):
            if lead_name not in df.columns:
                offsets[lead_name] = np.nan
                continue
            
            gt_mv = df[lead_name].dropna().values
            if len(gt_mv) == 0:
                offsets[lead_name] = np.nan
                continue
            
            # Handle Lead II short strip
            if lead_name == 'II':
                ref_lead = 'I' if 'I' in df.columns else 'III'
                if ref_lead in df.columns:
                    ref_len = len(df[ref_lead].dropna().values)
                    if len(gt_mv) > ref_len:
                        gt_mv = gt_mv[:ref_len]
            
            seg_start = seg_idx * segment_width
            seg_end = (seg_idx + 1) * segment_width
            pred_segment = pred_mv_rows[row_idx][seg_start:seg_end]
            
            offset = compute_baseline_offset(pred_segment, gt_mv)
            offsets[lead_name] = offset
    
    return offsets


def describe_distribution(values, lead_name):
    values = np.array([v for v in values if not np.isnan(v)])
    if len(values) == 0:
        return None
    
    binned = np.round(values * 100) / 100
    mode_result = stats.mode(binned, keepdims=True)
    mode_val = mode_result.mode[0]
    mode_count = mode_result.count[0]
    
    return {
        'lead': lead_name,
        'count': len(values),
        'mean': np.mean(values),
        'median': np.median(values),
        'mode': mode_val,
        'mode_count': mode_count,
        'std': np.std(values),
        'min': np.min(values),
        'max': np.max(values),
        'range': np.max(values) - np.min(values),
        'p5': np.percentile(values, 5),
        'p25': np.percentile(values, 25),
        'p75': np.percentile(values, 75),
        'p95': np.percentile(values, 95),
        'iqr': np.percentile(values, 75) - np.percentile(values, 25),
    }


def main():
    parser = argparse.ArgumentParser(description='Compute per-lead baseline offsets for V18')
    parser.add_argument('--v16_checkpoint', type=str, required=True)
    parser.add_argument('--v18_checkpoint', type=str, required=True)
    parser.add_argument('--kaggle_data', type=str, default='/data/ecg-digitization/stage1_data/train')
    parser.add_argument('--num_samples', type=int, default=4000)
    parser.add_argument('--output_dir', type=str, default='/home/azureuser/tmp/baseline_analysis_v18')
    parser.add_argument('--seed', type=int, default=42)
    args = parser.parse_args()
    
    random.seed(args.seed)
    np.random.seed(args.seed)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    print(f"{'='*80}")
    print(f"Baseline Offset Analysis for V18")
    print(f"{'='*80}")
    print(f"V16 Checkpoint: {args.v16_checkpoint}")
    print(f"V18 Checkpoint: {args.v18_checkpoint}")
    print(f"Num samples: {args.num_samples}")
    print(f"Output dir: {output_dir}")
    print(f"Device: {device}")
    print(f"{'='*80}")
    
    # Load models
    v16_model = load_v16_model(args.v16_checkpoint, device)
    v18_refiner = load_v18_refiner(args.v18_checkpoint, device)
    
    # Find samples
    kaggle_dir = Path(args.kaggle_data)
    sample_paths = []
    
    print("Finding samples...")
    sample_dirs = list(kaggle_dir.iterdir())
    random.shuffle(sample_dirs)
    
    for sample_dir in sample_dirs:
        if not sample_dir.is_dir():
            continue
        csv_files = list(sample_dir.glob('*.csv'))
        if len(csv_files) != 1:
            continue
        csv_path = csv_files[0]
        
        for img_path in sample_dir.glob('*.png'):
            variant = img_path.stem.split('-')[-1] if '-' in img_path.stem else '0000'
            if variant in VALID_VARIANTS:
                sample_paths.append((img_path, csv_path))
                break
        
        if len(sample_paths) >= args.num_samples:
            break
    
    print(f"Found {len(sample_paths)} samples")
    
    # Process all samples
    all_offsets = {lead: [] for lead in ALL_LEADS}
    raw_data = []
    
    for img_path, csv_path in tqdm(sample_paths, desc="Computing offsets"):
        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)
            
            offsets = process_image(v16_model, v18_refiner, image, df, device)
            
            row = {'image': img_path.stem}
            for lead in ALL_LEADS:
                offset = offsets.get(lead, np.nan)
                row[lead] = offset
                if not np.isnan(offset):
                    all_offsets[lead].append(offset)
            raw_data.append(row)
            
        except Exception as e:
            print(f"Error processing {img_path}: {e}")
            continue
    
    # Save raw data
    raw_df = pd.DataFrame(raw_data)
    raw_csv_path = output_dir / 'baseline_offsets_raw.csv'
    raw_df.to_csv(raw_csv_path, index=False)
    print(f"\nSaved raw data to: {raw_csv_path}")
    
    # Print statistics
    print(f"\n{'='*80}")
    print(f"BASELINE OFFSET STATISTICS (mV) - V18")
    print(f"{'='*80}")
    print(f"Offset = median(prediction - ground_truth)")
    print(f"Positive offset means model predicts HIGHER than GT")
    print(f"{'='*80}\n")
    
    stats_data = []
    
    print(f"{'Lead':<8} {'Count':>6} {'Mean':>8} {'Median':>8} {'Mode':>8} {'Std':>8} {'Min':>8} {'Max':>8} {'Range':>8} {'IQR':>8}")
    print("-" * 90)
    
    for lead in ALL_LEADS:
        stats_info = describe_distribution(all_offsets[lead], lead)
        if stats_info:
            stats_data.append(stats_info)
            print(f"{lead:<8} {stats_info['count']:>6} {stats_info['mean']:>8.4f} {stats_info['median']:>8.4f} "
                  f"{stats_info['mode']:>8.4f} {stats_info['std']:>8.4f} {stats_info['min']:>8.4f} "
                  f"{stats_info['max']:>8.4f} {stats_info['range']:>8.4f} {stats_info['iqr']:>8.4f}")
    
    # Print recommended offsets
    print(f"\n{'='*80}")
    print(f"RECOMMENDED HARDCODED BASELINE OFFSETS (mV) - V18")
    print(f"{'='*80}")
    print()
    
    recommended = {}
    print("BASELINE_OFFSETS_V18 = {")
    for lead in ALL_LEADS:
        values = all_offsets[lead]
        if values:
            median_offset = np.median(values)
            recommended[lead] = float(median_offset)
            print(f"    '{lead}': {median_offset:.4f},")
    print("}")
    
    # Save stats
    stats_json_path = output_dir / 'baseline_stats.json'
    with open(stats_json_path, 'w') as f:
        json.dump({
            'stats': stats_data,
            'recommended_offsets': recommended,
            'num_samples': len(sample_paths),
            'v16_checkpoint': args.v16_checkpoint,
            'v18_checkpoint': args.v18_checkpoint,
        }, f, indent=2)
    print(f"\nSaved stats to: {stats_json_path}")
    
    print(f"\n{'='*80}")
    print("Done!")


if __name__ == '__main__':
    main()