#!/usr/bin/env python3
"""
Compute Per-Lead Baseline Offsets for V23

This script runs V23 on training images and computes the median offset
(prediction - ground_truth) for each lead. These offsets can then be
subtracted during inference to improve accuracy.

The offsets are typically very small (~0.005-0.007 mV) but help with
systematic bias in the model.

Usage:
  python compute_baseline_offsets_v23.py \
    --checkpoint /data/ecg-digitization/checkpoints/v23_latest.pth \
    --num_samples 1000

Output:
  - Prints BASELINE_OFFSETS dict to copy into inference script
  - Saves raw data to CSV for analysis
"""

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 import stats
from scipy.signal import savgol_filter

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.cuda.amp import autocast

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

# Don't use validation samples for baseline computation
VAL_SAMPLE_IDS = set([
    '1006427285', '1006867983', '1012423188', '10140238', '1015663939',
    '102150619', '1026034238', '1041099777', '104573050', '1048962695',
    '1052007218', '1053922973', '1059602762', '1063816858',
    '1067371646', '1067975047', '1068062585', '1072767337',
    '1084993373', '108599929'
])

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


# =============================================================================
# Model Architecture (V23)
# =============================================================================
class MultiScaleFeatureFusion(nn.Module):
    def __init__(self, l3_channels, l4_channels, out_channels=256):
        super().__init__()
        self.l3_proj = nn.Sequential(
            nn.Conv2d(l3_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.GELU(),
        )
        self.l4_proj = nn.Sequential(
            nn.Conv2d(l4_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.GELU(),
        )
        self.fusion = nn.Sequential(
            nn.Conv2d(out_channels * 2, out_channels, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.GELU(),
            nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.GELU(),
        )
        self.out_channels = out_channels
    
    def forward(self, l3_feat, l4_feat):
        l3_proj = self.l3_proj(l3_feat)
        l4_proj = self.l4_proj(l4_feat)
        l4_up = F.interpolate(l4_proj, size=l3_proj.shape[2:], mode='bilinear', align_corners=True)
        fused = torch.cat([l3_proj, l4_up], dim=1)
        fused = self.fusion(fused)
        return fused


class HeightAttention(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.attention = nn.Sequential(
            nn.Conv2d(in_channels + 2, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.GELU(),
            nn.Conv2d(64, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.GELU(),
            nn.Conv2d(32, 1, 1),
        )
    
    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_coord = torch.cat([x, yy, xx], dim=1)
        attn = self.attention(x_coord)
        attn = F.softmax(attn, dim=2)
        pooled = (x * attn).sum(dim=2)
        return pooled


class BiLSTMTemporal(nn.Module):
    def __init__(self, input_dim, hidden_dim=128, num_layers=2, dropout=0.1):
        super().__init__()
        self.lstm = nn.LSTM(
            input_size=input_dim, hidden_size=hidden_dim,
            num_layers=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)
        out = self.output_proj(lstm_out)
        return out


class FinalHead(nn.Module):
    def __init__(self, input_dim, hidden_dim=64):
        super().__init__()
        self.conv1d = nn.Sequential(
            nn.Conv1d(input_dim, hidden_dim, kernel_size=5, padding=2),
            nn.BatchNorm1d(hidden_dim),
            nn.GELU(),
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1),
            nn.BatchNorm1d(hidden_dim),
            nn.GELU(),
        )
        self.linear = nn.Sequential(
            nn.Linear(hidden_dim, 32),
            nn.GELU(),
            nn.Linear(32, 1),
            nn.Sigmoid(),
        )
    
    def forward(self, x):
        x = x.permute(0, 2, 1)
        x = self.conv1d(x)
        x = x.permute(0, 2, 1)
        out = self.linear(x).squeeze(-1)
        return out


class ColumnDSNT(nn.Module):
    def __init__(self, in_channels, hidden_channels=32):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, hidden_channels, 3, padding=1),
            nn.BatchNorm2d(hidden_channels),
            nn.GELU(),
            nn.Conv2d(hidden_channels, hidden_channels, 3, padding=1),
            nn.BatchNorm2d(hidden_channels),
            nn.GELU(),
            nn.Conv2d(hidden_channels, 1, 1),
        )
        
    def forward(self, x):
        B, C, H, W = x.shape
        logits = self.conv(x)
        probs = F.softmax(logits, dim=2)
        y_coords_normalized = torch.linspace(0, 1, H, device=x.device)
        y_coords_normalized = y_coords_normalized.view(1, 1, H, 1)
        y_expected = (probs * y_coords_normalized).sum(dim=2)
        y_expected = y_expected.squeeze(1)
        return y_expected, probs


class PerLeadNetV23(nn.Module):
    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()
        
        self.ms_fusion = MultiScaleFeatureFusion(
            l3_channels=enc_channels[2], l4_channels=enc_channels[3], out_channels=256,
        )
        
        self.dec_blocks = nn.ModuleList()
        self.dec_blocks.append(self._make_dec_block(256, enc_channels[1], 128))
        self.dec_blocks.append(self._make_dec_block(128, enc_channels[0], 64))
        self.dec_blocks.append(self._make_dec_block(64, 0, 32))
        self.dec_blocks.append(self._make_dec_block(32, 0, 32))
        
        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 = HeightAttention(in_channels=32)
        self.bilstm = BiLSTMTemporal(input_dim=32, hidden_dim=128, num_layers=2, dropout=0.1)
        self.head = FinalHead(input_dim=128, hidden_dim=64)
        self.dsnt_head = ColumnDSNT(in_channels=32, hidden_channels=32)
        self.dsnt_weight = nn.Parameter(torch.tensor(0.3))
    
    def _make_dec_block(self, in_ch, skip_ch, out_ch):
        return nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            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(),
        )
    
    def forward(self, x):
        B, C, H, W = x.shape
        features = self.encoder(x)
        l1, l2, l3, l4 = features
        
        fused = self.ms_fusion(l3, l4)
        
        d = self.dec_blocks[0][0](fused)
        if d.shape[2:] != l2.shape[2:]:
            d = F.interpolate(d, size=l2.shape[2:], mode='bilinear', align_corners=True)
        d = torch.cat([d, l2], dim=1)
        d = self.dec_blocks[0][1:](d)
        
        d = self.dec_blocks[1][0](d)
        if d.shape[2:] != l1.shape[2:]:
            d = F.interpolate(d, size=l1.shape[2:], mode='bilinear', align_corners=True)
        d = torch.cat([d, l1], dim=1)
        d = self.dec_blocks[1][1:](d)
        
        d = self.dec_blocks[2](d)
        d = self.dec_blocks[3](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)
        
        y_dsnt, _ = self.dsnt_head(d)
        pooled = self.height_attention(d)
        temporal = self.bilstm(pooled)
        y_v22 = self.head(temporal)
        
        w = torch.sigmoid(self.dsnt_weight)
        y_pred = (1 - w) * y_v22 + w * y_dsnt
        
        return y_pred


# =============================================================================
# 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 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 load_model(checkpoint_path, device):
    model = PerLeadNetV23(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()
    
    dsnt_w = torch.sigmoid(model.dsnt_weight).item()
    print(f"Loaded V23 from epoch {checkpoint.get('epoch', 'N/A')}")
    print(f"  SNR: {checkpoint.get('snr', 'N/A'):.2f} dB")
    print(f"  DSNT weight: {dsnt_w:.3f} (V22: {1-dsnt_w:.3f})")
    return model


def compute_baseline_offset(pred_mv_segment, gt_mv):
    """Compute median offset between prediction and ground truth."""
    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(model, image, df, device):
    """Process image and compute per-lead offsets."""
    segment_width = OUTPUT_WIDTH // 4
    pred_mv_rows = {}
    
    # Get predictions for all rows
    for row_idx in range(3):  # Only rows 0-2 (not rhythm strip)
        row_crop = crop_row(image, row_idx)
        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 autocast():
                pred = model(image_tensor)
        
        pred_y = pred[0].cpu().numpy() * ROW_HEIGHT
        baseline_y = ZERO_MV[row_idx]
        pred_mv = (CROP_HALF_HEIGHT - pred_y) / MV_TO_PIXEL
        pred_mv = apply_savgol_smoothing(pred_mv, window=7, polyorder=2)
        pred_mv = clamp_ecg_amplitude(pred_mv)
        pred_mv_rows[row_idx] = pred_mv
    
    # Apply Einthoven correction
    pred_mv_rows = apply_einthoven_correction(pred_mv_rows, alpha=0.33)
    
    # Compute offsets 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 (10s in GT, 2.5s in 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]
            
            offsets[lead_name] = compute_baseline_offset(pred_segment, gt_mv)
    
    return offsets


def main():
    parser = argparse.ArgumentParser(description='Compute V23 baseline offsets')
    parser.add_argument('--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=1000)
    parser.add_argument('--output_dir', type=str, default='/home/azureuser/tmp/baseline_analysis_v23')
    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"{'='*70}")
    print(f"Baseline Offset Computation for V23")
    print(f"{'='*70}")
    print(f"Checkpoint: {args.checkpoint}")
    print(f"Num samples: {args.num_samples}")
    print(f"Device: {device}")
    print(f"{'='*70}")
    
    model = load_model(args.checkpoint, device)
    
    # Collect sample paths (excluding validation samples)
    kaggle_dir = Path(args.kaggle_data)
    sample_paths = []
    sample_dirs = list(kaggle_dir.iterdir())
    random.shuffle(sample_dirs)
    
    for sample_dir in sample_dirs:
        if not sample_dir.is_dir():
            continue
        if sample_dir.name in VAL_SAMPLE_IDS:
            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 (excluding validation)")
    
    # Compute offsets for 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(model, 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_df.to_csv(output_dir / 'baseline_offsets_raw.csv', index=False)
    print(f"\nSaved raw offsets to {output_dir / 'baseline_offsets_raw.csv'}")
    
    # Compute summary statistics
    print(f"\n{'='*70}")
    print("Per-Lead Baseline Offset Statistics (mV)")
    print(f"{'='*70}")
    print(f"{'Lead':<6} {'Count':>6} {'Mean':>10} {'Median':>10} {'Std':>10}")
    print(f"{'-'*50}")
    
    median_offsets = {}
    for lead in ALL_LEADS:
        values = np.array(all_offsets[lead])
        if len(values) == 0:
            continue
        median_offsets[lead] = np.median(values)
        print(f"{lead:<6} {len(values):>6} {np.mean(values):>10.4f} {np.median(values):>10.4f} {np.std(values):>10.4f}")
    
    # Print Python dict format for inference script
    print(f"\n{'='*70}")
    print("Copy this into inference_v23.py:")
    print(f"{'='*70}")
    print("BASELINE_OFFSETS = {")
    for i, lead in enumerate(ALL_LEADS):
        offset = median_offsets.get(lead, 0.0)
        comma = "," if i < len(ALL_LEADS) - 1 else ""
        print(f"    '{lead}': {offset:.4f}{comma}")
    print("}")
    print(f"{'='*70}")
    
    # Save summary
    summary = {lead: {'median': median_offsets.get(lead, 0.0), 
                      'count': len(all_offsets[lead]),
                      'mean': np.mean(all_offsets[lead]) if all_offsets[lead] else 0.0,
                      'std': np.std(all_offsets[lead]) if all_offsets[lead] else 0.0}
               for lead in ALL_LEADS}
    
    with open(output_dir / 'baseline_offsets_summary.json', 'w') as f:
        import json
        json.dump(summary, f, indent=2)
    
    print(f"\nSaved summary to {output_dir / 'baseline_offsets_summary.json'}")
    print("Done!")


if __name__ == '__main__':
    main()