#!/usr/bin/env python3
"""
V19 Fine-tuning with Soft-DTW Loss

Soft-DTW allows temporal alignment tolerance during training,
which can improve performance on signals where exact pixel alignment varies.

Strategy:
- Start from best V19 checkpoint (epoch 19)
- Fine-tune with Soft-DTW loss + MSE loss
- Use very low learning rate to preserve existing features
- Train for few epochs only (2-5)
"""

import os
import sys
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import timm
from torch.utils.data import DataLoader, Dataset
from pathlib import Path
import cv2
from tqdm import tqdm
import random
from datetime import datetime

# Soft-DTW implementation
try:
    from tslearn.metrics import soft_dtw
    from tslearn.backend import Backend
    HAS_TSLEARN = True
except ImportError:
    HAS_TSLEARN = False
    print("Warning: tslearn not found, using pure PyTorch Soft-DTW")

try:
    from torchvision.ops import DeformConv2d
    HAS_DEFORM_CONV = True
except ImportError:
    HAS_DEFORM_CONV = False

# Paths
BASELINE_PATH = '/home/azureuser/tmp/hengck23/hengck23-submit-physionet'
V19_CHECKPOINT = '/data/ecg-digitization/checkpoints/v19_enhanced_epoch019.pth'
TRAIN_DIR = Path('/data/ecg-digitization/data/physionet.org/files/ecg-image-database/1.0.0/synthetic_images')
GT_DIR = Path('/data/ecg-digitization/data/physionet.org/files/ecg-image-database/1.0.0/ground_truth')
OUTPUT_DIR = Path('/data/ecg-digitization/checkpoints')

# Config
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
LEARNING_RATE = 1e-5  # Very low for fine-tuning
EPOCHS = 3
BATCH_SIZE = 4
DTW_GAMMA = 1.0  # Soft-DTW smoothness parameter
LAMBDA_DTW = 0.3  # Weight for DTW loss (0.3 DTW + 0.7 MSE)

# ECG 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
CROP_HALF_HEIGHT = 250
ROW_HEIGHT = 500
OUTPUT_WIDTH = T1 - T0
ECG_MV_MIN, ECG_MV_MAX = -7.0, 7.0

# ======================= Soft-DTW Loss =======================
class SoftDTWLoss(nn.Module):
    """
    Differentiable Soft-DTW loss for time series alignment.
    Based on Cuturi & Blondel 2017.
    """
    def __init__(self, gamma=1.0):
        super().__init__()
        self.gamma = gamma
    
    def pairwise_distances(self, x, y):
        """Compute pairwise squared Euclidean distances."""
        # x: [B, T1], y: [B, T2]
        x = x.unsqueeze(2)  # [B, T1, 1]
        y = y.unsqueeze(1)  # [B, 1, T2]
        return (x - y) ** 2  # [B, T1, T2]
    
    def forward(self, pred, target):
        """
        Compute Soft-DTW loss.
        pred: [B, T] predicted signal
        target: [B, T] ground truth signal
        """
        B, T = pred.shape
        D = self.pairwise_distances(pred, target)  # [B, T, T]
        
        # Initialize R matrix
        R = torch.full((B, T + 1, T + 1), float('inf'), device=pred.device)
        R[:, 0, 0] = 0
        
        # Forward pass with soft-min
        for i in range(1, T + 1):
            for j in range(1, T + 1):
                r_options = torch.stack([
                    R[:, i-1, j-1],
                    R[:, i-1, j],
                    R[:, i, j-1]
                ], dim=1)  # [B, 3]
                
                # Soft-min
                R[:, i, j] = -self.gamma * torch.logsumexp(-r_options / self.gamma, dim=1) + D[:, i-1, j-1]
        
        return R[:, T, T].mean()


class CombinedLoss(nn.Module):
    """Combined MSE + Soft-DTW loss."""
    def __init__(self, lambda_dtw=0.3, gamma=1.0):
        super().__init__()
        self.lambda_dtw = lambda_dtw
        self.mse = nn.MSELoss()
        self.soft_dtw = SoftDTWLoss(gamma=gamma)
    
    def forward(self, pred, target):
        mse_loss = self.mse(pred, target)
        
        # DTW on downsampled signals for efficiency
        stride = 8
        pred_ds = pred[:, ::stride]
        target_ds = target[:, ::stride]
        dtw_loss = self.soft_dtw(pred_ds, target_ds)
        
        total = (1 - self.lambda_dtw) * mse_loss + self.lambda_dtw * dtw_loss
        return total, mse_loss, dtw_loss


# ======================= V19 Model =======================
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))
            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:
            out = self.deform_conv(x, self.offset_conv(x))
        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 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 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):
    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]
        for i, (skip_ch, out_ch) in enumerate(zip(enc_channels[:-1][::-1] + [0], decoder_dims)):
            self.dec_blocks.append(UNetDecoderBlockV19(in_ch, skip_ch, out_ch, use_deform=(i >= 2)))
            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, hidden_dim=128, num_layers=2, dropout=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]
        for block, skip in zip(self.dec_blocks, features[:-1][::-1] + [None]):
            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


# ======================= Dataset =======================
class ECGDataset(Dataset):
    def __init__(self, train_dir, gt_dir, max_samples=None):
        self.samples = []
        
        # Find training images
        for variant_dir in sorted(train_dir.iterdir()):
            if not variant_dir.is_dir():
                continue
            for img_path in sorted(variant_dir.glob('*.png')):
                sample_id = img_path.stem.replace('-0', '-').replace('-', '/')
                gt_path = gt_dir / f"{img_path.stem.split('-')[0]}.csv"
                if gt_path.exists():
                    self.samples.append({
                        'image_path': img_path,
                        'gt_path': gt_path,
                        'row_idx': int(img_path.stem.split('-')[1])
                    })
        
        if max_samples:
            self.samples = self.samples[:max_samples]
        
        print(f"Found {len(self.samples)} training samples")
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        sample = self.samples[idx]
        
        # Load and preprocess image
        img = cv2.imread(str(sample['image_path']))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        # Load ground truth
        gt_df = np.loadtxt(sample['gt_path'], delimiter=',', skiprows=1)
        row_idx = sample['row_idx']
        
        # Get y values and convert to normalized coordinates
        y_values = gt_df[:, row_idx]  # Assuming CSV has rows as columns
        y_normalized = (y_values - y_values.min()) / (y_values.max() - y_values.min() + 1e-6)
        
        # Resize to match output width
        y_normalized = np.interp(
            np.linspace(0, 1, OUTPUT_WIDTH),
            np.linspace(0, 1, len(y_normalized)),
            y_normalized
        )
        
        # Convert to tensors
        img_tensor = torch.from_numpy(img.astype(np.float32) / 255.0).permute(2, 0, 1)
        y_tensor = torch.from_numpy(y_normalized.astype(np.float32))
        
        return img_tensor, y_tensor


class SimpleRowDataset(Dataset):
    """Simple dataset loading pre-cropped rows."""
    def __init__(self, data_dir, max_samples=None):
        self.samples = []
        data_dir = Path(data_dir)
        
        for row_dir in sorted(data_dir.glob('row_*')):
            for img_path in sorted(row_dir.glob('*.npz')):
                self.samples.append(str(img_path))
        
        if max_samples:
            random.shuffle(self.samples)
            self.samples = self.samples[:max_samples]
        
        print(f"Found {len(self.samples)} pre-cropped samples")
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        data = np.load(self.samples[idx])
        img = torch.from_numpy(data['image'].astype(np.float32) / 255.0)
        if img.dim() == 3 and img.shape[2] == 3:
            img = img.permute(2, 0, 1)
        y = torch.from_numpy(data['y_normalized'].astype(np.float32))
        return img, y


# ======================= Training =======================
def train_epoch(model, dataloader, optimizer, criterion, device, epoch):
    model.train()
    total_loss = 0
    total_mse = 0
    total_dtw = 0
    
    pbar = tqdm(dataloader, desc=f'Epoch {epoch}')
    for batch_idx, (images, targets) in enumerate(pbar):
        images = images.to(device)
        targets = targets.to(device)
        
        optimizer.zero_grad()
        
        with torch.amp.autocast('cuda', dtype=torch.float16):
            pred = model(images, return_aux=False)
            loss, mse_loss, dtw_loss = criterion(pred, targets)
        
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        
        total_loss += loss.item()
        total_mse += mse_loss.item()
        total_dtw += dtw_loss.item()
        
        pbar.set_postfix({
            'loss': f'{loss.item():.4f}',
            'mse': f'{mse_loss.item():.4f}',
            'dtw': f'{dtw_loss.item():.4f}'
        })
    
    n = len(dataloader)
    return total_loss / n, total_mse / n, total_dtw / n


def validate(model, dataloader, device):
    model.eval()
    total_snr = 0
    n_samples = 0
    
    with torch.no_grad():
        for images, targets in tqdm(dataloader, desc='Validation'):
            images = images.to(device)
            targets = targets.to(device)
            
            with torch.amp.autocast('cuda', dtype=torch.float16):
                pred = model(images, return_aux=False)
            
            # Compute SNR
            mse = ((pred - targets) ** 2).mean(dim=1)
            signal_power = (targets ** 2).mean(dim=1)
            snr = 10 * torch.log10(signal_power / (mse + 1e-10))
            
            total_snr += snr.sum().item()
            n_samples += images.shape[0]
    
    return total_snr / n_samples if n_samples > 0 else 0


def main():
    print(f"Device: {DEVICE}")
    print(f"Deformable Conv: {HAS_DEFORM_CONV}")
    
    # Load model
    model = PerLeadNetV19(pretrained=False)
    checkpoint = torch.load(V19_CHECKPOINT, map_location='cpu', weights_only=False)
    model.load_state_dict(checkpoint['model'], strict=True)
    model.to(DEVICE)
    print(f"Loaded V19 epoch {checkpoint['epoch']}, SNR={checkpoint['snr']:.2f} dB")
    
    # Create combined loss
    criterion = CombinedLoss(lambda_dtw=LAMBDA_DTW, gamma=DTW_GAMMA)
    
    # Optimizer - lower LR for fine-tuning
    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=0.01)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS, eta_min=LEARNING_RATE / 10)
    
    # Check for pre-cropped data
    precropped_dir = Path('/data/ecg-digitization/precropped_rows')
    if precropped_dir.exists():
        dataset = SimpleRowDataset(precropped_dir, max_samples=1000)
    else:
        print("Warning: Pre-cropped data not found. Creating minimal dataset...")
        # Use existing validation data as proxy
        dataset = SimpleRowDataset('/data/ecg-digitization/validation_rows', max_samples=500) if Path('/data/ecg-digitization/validation_rows').exists() else None
    
    if dataset is None or len(dataset) == 0:
        print("Error: No training data found!")
        return
    
    # Split into train/val
    n_val = min(100, len(dataset) // 5)
    n_train = len(dataset) - n_val
    train_dataset, val_dataset = torch.utils.data.random_split(dataset, [n_train, n_val])
    
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True)
    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4, pin_memory=True)
    
    print(f"Training samples: {len(train_dataset)}, Validation samples: {len(val_dataset)}")
    
    best_snr = checkpoint['snr']
    
    for epoch in range(1, EPOCHS + 1):
        loss, mse, dtw = train_epoch(model, train_loader, optimizer, criterion, DEVICE, epoch)
        val_snr = validate(model, val_loader, DEVICE)
        scheduler.step()
        
        print(f"Epoch {epoch}: Loss={loss:.4f}, MSE={mse:.4f}, DTW={dtw:.4f}, Val SNR={val_snr:.2f} dB")
        
        # Save checkpoint
        ckpt_path = OUTPUT_DIR / f'v19_softdtw_epoch{epoch:03d}.pth'
        torch.save({
            'epoch': checkpoint['epoch'] + epoch,
            'model': model.state_dict(),
            'optimizer': optimizer.state_dict(),
            'snr': val_snr,
            'loss': loss
        }, ckpt_path)
        print(f"Saved: {ckpt_path}")
        
        if val_snr > best_snr:
            best_snr = val_snr
            best_path = OUTPUT_DIR / 'v19_softdtw_best.pth'
            torch.save({
                'epoch': checkpoint['epoch'] + epoch,
                'model': model.state_dict(),
                'snr': val_snr
            }, best_path)
            print(f"New best! SNR={val_snr:.2f} dB")
    
    print(f"\nFine-tuning complete. Best SNR: {best_snr:.2f} dB")


if __name__ == '__main__':
    main()