code/scripts/dataset_v2.py

"""
ECG Dataset V2 - With Multi-Version Support and Style-Specific Augmentation

Based on expert recommendations:
- Use Stage 1 recovered images for training consistency
- Heavy augmentation based on different image types (0001-0012)
- Exclude bad data samples
- Train on all 9-12 image versions with augmentation

Usage:
    from dataset_v2 import create_dataloaders_v2
"""

import os
import torch
import numpy as np
import cv2
import pandas as pd
from pathlib import Path
from torch.utils.data import Dataset, DataLoader, ConcatDataset
import albumentations as A
from albumentations.pytorch import ToTensorV2
import random
import json


# Target specifications
TARGET_HEIGHT = 1696  
TARGET_WIDTH = 4352

# Calibration constants
ZERO_MV = [703.5, 987.5, 1271.5, 1531.5]
MV_TO_PIXEL = 78.5
T0, T1 = 235, 4161


# Style-specific augmentation strengths (0001-0012)
STYLE_AUGMENTATION = {
    '0001': {'noise': 0.3, 'blur': 0.2, 'color': 0.5, 'quality': 0.3},  # Clean print
    '0002': {'noise': 0.5, 'blur': 0.3, 'color': 0.6, 'quality': 0.4},  # Aged paper
    '0003': {'noise': 0.3, 'blur': 0.2, 'color': 0.5, 'quality': 0.3},  # Green grid
    '0004': {'noise': 0.5, 'blur': 0.3, 'color': 0.6, 'quality': 0.4},  # Green aged
    '0005': {'noise': 0.3, 'blur': 0.3, 'color': 0.7, 'quality': 0.3},  # Blue trace
    '0006': {'noise': 0.6, 'blur': 0.4, 'color': 0.7, 'quality': 0.5},  # Grey bg
    '0007': {'noise': 0.6, 'blur': 0.4, 'color': 0.7, 'quality': 0.5},  # Yellow aged
    '0008': {'noise': 0.4, 'blur': 0.3, 'color': 0.5, 'quality': 0.4},  # Faded green
    '0009': {'noise': 0.3, 'blur': 0.2, 'color': 0.6, 'quality': 0.3},  # Orange grid
    '0010': {'noise': 0.4, 'blur': 0.3, 'color': 0.6, 'quality': 0.4},  # High contrast
    '0011': {'noise': 0.5, 'blur': 0.3, 'color': 0.5, 'quality': 0.5},  # Low contrast
    '0012': {'noise': 0.3, 'blur': 0.2, 'color': 0.4, 'quality': 0.3},  # No grid
}


def get_style_transforms(style_key='0001', height=TARGET_HEIGHT, width=TARGET_WIDTH):
    """
    Get style-specific augmentation transforms.
    Different image styles need different augmentation intensities.
    """
    aug_params = STYLE_AUGMENTATION.get(style_key, STYLE_AUGMENTATION['0001'])
    
    transforms = [
        A.Resize(height, width),
    ]
    
    # Color augmentations (intensity varies by style)
    color_p = aug_params['color']
    transforms.append(
        A.OneOf([
            A.ColorJitter(
                brightness=0.15 * color_p,
                contrast=0.15 * color_p, 
                saturation=0.1 * color_p,
                hue=0.05 * color_p,
                p=0.8
            ),
            A.RandomBrightnessContrast(
                brightness_limit=0.15 * color_p,
                contrast_limit=0.15 * color_p,
                p=0.8
            ),
            A.HueSaturationValue(
                hue_shift_limit=int(10 * color_p),
                sat_shift_limit=int(20 * color_p),
                val_shift_limit=int(20 * color_p),
                p=0.8
            ),
        ], p=color_p)
    )
    
    # Noise augmentations
    noise_p = aug_params['noise']
    transforms.append(
        A.OneOf([
            A.GaussNoise(var_limit=(5, int(30 * noise_p)), p=0.5),
            A.ISONoise(
                color_shift=(0.01, 0.05 * noise_p),
                intensity=(0.1, 0.3 * noise_p),
                p=0.5
            ),
        ], p=noise_p)
    )
    
    # Blur augmentations
    blur_p = aug_params['blur']
    transforms.append(
        A.OneOf([
            A.GaussianBlur(blur_limit=(3, max(3, int(5 * blur_p))), p=0.5),
            A.MotionBlur(blur_limit=max(3, int(5 * blur_p)), p=0.3),
        ], p=blur_p)
    )
    
    # Quality degradation
    quality_p = aug_params['quality']
    transforms.append(
        A.OneOf([
            A.ImageCompression(
                quality_lower=max(50, int(100 - 30 * quality_p)),
                quality_upper=100,
                p=0.3
            ),
            A.Downscale(
                scale_min=max(0.7, 1.0 - 0.2 * quality_p),
                scale_max=0.95,
                p=0.3
            ),
        ], p=quality_p)
    )
    
    # Final transforms
    transforms.extend([
        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ToTensorV2(),
    ])
    
    return A.Compose(transforms)


def get_heavy_augmentation(height=TARGET_HEIGHT, width=TARGET_WIDTH):
    """
    Heavy augmentation for training on all image versions.
    Used when we want to be robust across all styles.
    """
    return A.Compose([
        A.Resize(height, width),
        
        # Strong color augmentations
        A.OneOf([
            A.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.1, p=1.0),
            A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=1.0),
            A.HueSaturationValue(hue_shift_limit=15, sat_shift_limit=30, val_shift_limit=30, p=1.0),
            A.RGBShift(r_shift_limit=20, g_shift_limit=20, b_shift_limit=20, p=1.0),
        ], p=0.8),
        
        # Channel manipulations
        A.OneOf([
            A.ChannelShuffle(p=0.3),
            A.ToGray(p=0.1),
            A.CLAHE(clip_limit=2.0, p=0.3),
        ], p=0.3),
        
        # Noise
        A.OneOf([
            A.GaussNoise(var_limit=(10, 50), p=0.5),
            A.ISONoise(color_shift=(0.01, 0.1), intensity=(0.1, 0.5), p=0.5),
            A.MultiplicativeNoise(multiplier=(0.9, 1.1), p=0.3),
        ], p=0.5),
        
        # Blur
        A.OneOf([
            A.GaussianBlur(blur_limit=(3, 7), p=0.4),
            A.MotionBlur(blur_limit=5, p=0.3),
            A.MedianBlur(blur_limit=5, p=0.2),
        ], p=0.3),
        
        # Quality
        A.OneOf([
            A.ImageCompression(quality_lower=60, quality_upper=100, p=0.4),
            A.Downscale(scale_min=0.7, scale_max=0.95, p=0.3),
        ], p=0.3),
        
        # Shadows and lighting
        A.RandomShadow(shadow_roi=(0, 0, 1, 1), num_shadows_lower=1, num_shadows_upper=3, p=0.2),
        
        # Normalize and convert
        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ToTensorV2(),
    ])


def get_val_transforms(height=TARGET_HEIGHT, width=TARGET_WIDTH):
    """Validation transforms - resize and normalize only."""
    return A.Compose([
        A.Resize(height, width),
        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ToTensorV2(),
    ])


class SyntheticECGDatasetV2(Dataset):
    """
    Synthetic ECG Dataset - loads Stage 1 processed images with ground truth.
    
    Directory structure (from generate_synthetic_v2.py):
    synthetic_dir/
        stage1/     <- Main training images (Stage 1 rectified)
            syn_00000000-0001.png
            ...
        gt/         <- Ground truth signals
            syn_00000000.npy
            ...
        manifest.json  <- Sample metadata
    """
    def __init__(self, synthetic_dir, transform=None, max_samples=None, use_styles=None):
        self.synthetic_dir = Path(synthetic_dir)
        self.transform = transform or get_heavy_augmentation()
        
        # Load manifest if exists
        manifest_path = self.synthetic_dir / 'manifest.json'
        if manifest_path.exists():
            with open(manifest_path) as f:
                self.manifest = json.load(f)
        else:
            self.manifest = None
        
        # Find stage1 images
        stage1_dir = self.synthetic_dir / 'stage1'
        if stage1_dir.exists():
            self.image_files = sorted(stage1_dir.glob('*.png'))
        else:
            # Fallback to images directory
            image_dir = self.synthetic_dir / 'images'
            if image_dir.exists():
                self.image_files = sorted(image_dir.glob('*.png'))
            else:
                self.image_files = []
        
        # Filter by style if specified
        if use_styles:
            self.image_files = [
                f for f in self.image_files
                if any(f'-{s}' in f.stem for s in use_styles)
            ]
        
        # Limit samples if specified
        if max_samples:
            self.image_files = self.image_files[:max_samples]
        
        print(f"SyntheticDatasetV2: {len(self.image_files)} images")
    
    def __len__(self):
        return len(self.image_files)
    
    def __getitem__(self, idx):
        image_path = self.image_files[idx]
        
        # Load image
        image = cv2.imread(str(image_path))
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # Determine style from filename (e.g., syn_00000000-0001.png)
        filename = image_path.stem
        style = '0001'  # default
        for s in STYLE_AUGMENTATION.keys():
            if f'-{s}' in filename:
                style = s
                break
        
        # Load ground truth
        # Extract sample_id (e.g., syn_00000000 from syn_00000000-0001)
        sample_id = filename.split('-')[0] if '-' in filename else filename
        gt_path = self.synthetic_dir / 'gt' / f'{sample_id}.npy'
        
        if gt_path.exists():
            target = np.load(gt_path)  # Shape: [4, width]
        else:
            # No GT, return zeros
            target = np.zeros((4, T1 - T0), dtype=np.float32)
        
        # Apply style-specific or heavy augmentation
        if self.transform:
            transformed = self.transform(image=image)
            image = transformed['image']
        
        # Convert mV to normalized pixel coordinates
        target_pixel = np.zeros_like(target)
        for row_idx in range(4):
            target_pixel[row_idx] = ZERO_MV[row_idx] - target[row_idx] * MV_TO_PIXEL
        
        target_normalized = target_pixel / TARGET_HEIGHT
        target_normalized = np.clip(target_normalized, 0, 1)
        
        # Ensure correct output width
        output_width = T1 - T0
        if target_normalized.shape[1] != output_width:
            target_resized = np.zeros((4, output_width), dtype=np.float32)
            for row_idx in range(4):
                x_old = np.linspace(0, 1, target_normalized.shape[1])
                x_new = np.linspace(0, 1, output_width)
                target_resized[row_idx] = np.interp(x_new, x_old, target_normalized[row_idx])
            target_normalized = target_resized
        
        return {
            'image': image,
            'target': torch.from_numpy(target_normalized.astype(np.float32)),
            'id': filename,
            'style': style,
        }


class KaggleECGDatasetV2(Dataset):
    """
    Kaggle ECG Dataset V2 - loads ALL image versions (0001-0012) with augmentation.
    
    Key features:
    - Uses Stage 1 rectified images for consistency
    - Loads all versions (0001-0012) of each image
    - Applies style-specific augmentation
    - Excludes bad data samples (configurable)
    
    Directory structure:
    kaggle_dir/
        train/
            {id}/
                {id}-0001.png, {id}-0002.png, ..., {id}-0012.png
                {id}.csv  (ground truth)
        stage1_output/
            {id}-0001.png, {id}-0002.png, ...  (rectified)
    """
    def __init__(
        self,
        kaggle_dir,
        stage1_dir=None,
        df=None,
        transform=None,
        use_all_versions=True,
        versions=None,
        exclude_ids=None,
    ):
        self.kaggle_dir = Path(kaggle_dir)
        self.stage1_dir = Path(stage1_dir) if stage1_dir else None
        self.transform = transform or get_heavy_augmentation()
        self.use_all_versions = use_all_versions
        self.versions = versions or [f'{i:04d}' for i in range(1, 13)]  # 0001-0012
        self.exclude_ids = set(exclude_ids or [])
        
        # Get image IDs from train directory or df
        if df is not None:
            self.image_ids = [str(x) for x in df['id'].unique() if str(x) not in self.exclude_ids]
        else:
            train_dir = self.kaggle_dir / 'train'
            if train_dir.exists():
                self.image_ids = [
                    d.name for d in train_dir.iterdir() 
                    if d.is_dir() and d.name not in self.exclude_ids
                ]
            else:
                self.image_ids = []
        
        # Build sample list: (image_id, version)
        self.samples = []
        for img_id in self.image_ids:
            if use_all_versions:
                for v in self.versions:
                    self.samples.append((img_id, v))
            else:
                # Just use first version
                self.samples.append((img_id, self.versions[0]))
        
        print(f"KaggleDatasetV2: {len(self.image_ids)} images, {len(self.samples)} samples (versions: {len(self.versions)})")
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        image_id, version = self.samples[idx]
        
        # Try Stage 1 rectified first
        image = None
        if self.stage1_dir:
            stage1_path = self.stage1_dir / f"{image_id}-{version}.png"
            if stage1_path.exists():
                image = cv2.imread(str(stage1_path))
            else:
                # Try without version suffix
                stage1_path = self.stage1_dir / f"{image_id}.png"
                if stage1_path.exists():
                    image = cv2.imread(str(stage1_path))
        
        # Fall back to original
        if image is None:
            original_path = self.kaggle_dir / 'train' / image_id / f"{image_id}-{version}.png"
            if original_path.exists():
                image = cv2.imread(str(original_path))
            else:
                # Last resort - try first version
                original_path = self.kaggle_dir / 'train' / image_id / f"{image_id}-0001.png"
                if original_path.exists():
                    image = cv2.imread(str(original_path))
                else:
                    # Return dummy image
                    image = np.zeros((TARGET_HEIGHT, TARGET_WIDTH, 3), dtype=np.uint8)
        
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # Load ground truth
        csv_path = self.kaggle_dir / 'train' / image_id / f"{image_id}.csv"
        if csv_path.exists():
            gt_df = pd.read_csv(csv_path)
            target = self._csv_to_target(gt_df)
        else:
            target = np.zeros((4, T1 - T0), dtype=np.float32)
        
        # Apply transforms (style-specific or heavy)
        if self.transform:
            transformed = self.transform(image=image)
            image = transformed['image']
        
        return {
            'image': image,
            'target': torch.from_numpy(target.astype(np.float32)),
            'id': image_id,
            'version': version,
        }
    
    def _csv_to_target(self, df):
        """Convert ground truth CSV to normalized target array."""
        lead_layout = [
            ['I', 'aVR', 'V1', 'V4'],
            ['II', 'aVL', 'V2', 'V5'],
            ['III', 'aVF', 'V3', 'V6'],
        ]
        
        output_width = T1 - T0
        target = np.zeros((4, output_width), dtype=np.float32)
        
        # Process rows 0-2
        for row_idx in range(3):
            row_signals = []
            quarter_width = output_width // 4
            
            for lead in lead_layout[row_idx]:
                if lead in df.columns:
                    signal = df[lead].dropna().values
                    if len(signal) > 0:
                        x_old = np.linspace(0, 1, len(signal))
                        x_new = np.linspace(0, 1, quarter_width)
                        signal_resampled = np.interp(x_new, x_old, signal)
                    else:
                        signal_resampled = np.zeros(quarter_width)
                else:
                    signal_resampled = np.zeros(quarter_width)
                row_signals.append(signal_resampled)
            
            target[row_idx] = np.concatenate(row_signals)
        
        # Row 3: Full II rhythm strip
        if 'II' in df.columns:
            signal_ii = df['II'].dropna().values
            if len(signal_ii) > 0:
                x_old = np.linspace(0, 1, len(signal_ii))
                x_new = np.linspace(0, 1, output_width)
                target[3] = np.interp(x_new, x_old, signal_ii)
        
        # Convert to normalized pixel coordinates
        target_pixel = np.zeros_like(target)
        for row_idx in range(4):
            target_pixel[row_idx] = ZERO_MV[row_idx] - target[row_idx] * MV_TO_PIXEL
        
        target_normalized = target_pixel / TARGET_HEIGHT
        return np.clip(target_normalized, 0, 1)


class MixedDatasetV2(Dataset):
    """
    Combined synthetic + real dataset for training.
    """
    def __init__(self, synthetic_dataset, kaggle_dataset, synthetic_ratio=0.3):
        self.synthetic_dataset = synthetic_dataset
        self.kaggle_dataset = kaggle_dataset
        
        n_kaggle = len(kaggle_dataset)
        n_synthetic_to_use = int(n_kaggle * synthetic_ratio / (1 - synthetic_ratio))
        
        # Create index mapping
        self.indices = []
        
        # Add all kaggle samples
        for i in range(n_kaggle):
            self.indices.append(('kaggle', i))
        
        # Add synthetic samples (with wrap-around)
        for i in range(n_synthetic_to_use):
            syn_idx = i % len(synthetic_dataset)
            self.indices.append(('synthetic', syn_idx))
        
        random.shuffle(self.indices)
        print(f"MixedDatasetV2: {n_kaggle} kaggle + {n_synthetic_to_use} synthetic = {len(self.indices)} total")
    
    def __len__(self):
        return len(self.indices)
    
    def __getitem__(self, idx):
        source, real_idx = self.indices[idx]
        if source == 'kaggle':
            return self.kaggle_dataset[real_idx]
        else:
            return self.synthetic_dataset[real_idx]


def create_dataloaders_v2(
    synthetic_dir=None,
    kaggle_dir=None,
    stage1_dir=None,
    train_df=None,
    val_df=None,
    batch_size=8,
    num_workers=4,
    synthetic_ratio=0.3,
    use_all_versions=True,
    exclude_ids=None,
    distributed=False,
):
    """
    Create DataLoaders V2 with multi-version support.
    
    Args:
        synthetic_dir: Path to synthetic data (with stage1/ subdirectory)
        kaggle_dir: Path to Kaggle competition data
        stage1_dir: Path to Stage 1 processed Kaggle images
        train_df: Training DataFrame (image IDs)
        val_df: Validation DataFrame
        batch_size: Batch size per GPU
        num_workers: Number of data loading workers
        synthetic_ratio: Ratio of synthetic data in mixed training
        use_all_versions: Whether to use all 12 image versions
        exclude_ids: List of image IDs to exclude (bad data)
        distributed: Whether using DDP
    
    Returns:
        train_loader, val_loader
    """
    train_transform = get_heavy_augmentation()
    val_transform = get_val_transforms()
    
    # Create training dataset
    if synthetic_dir and kaggle_dir:
        # Mixed training
        synthetic_ds = SyntheticECGDatasetV2(synthetic_dir, train_transform)
        kaggle_ds = KaggleECGDatasetV2(
            kaggle_dir,
            stage1_dir,
            train_df,
            train_transform,
            use_all_versions=use_all_versions,
            exclude_ids=exclude_ids,
        )
        train_dataset = MixedDatasetV2(synthetic_ds, kaggle_ds, synthetic_ratio)
        
    elif synthetic_dir:
        # Synthetic pre-training
        train_dataset = SyntheticECGDatasetV2(synthetic_dir, train_transform)
        
    elif kaggle_dir:
        # Kaggle fine-tuning
        train_dataset = KaggleECGDatasetV2(
            kaggle_dir,
            stage1_dir,
            train_df,
            train_transform,
            use_all_versions=use_all_versions,
            exclude_ids=exclude_ids,
        )
    else:
        raise ValueError("Must provide synthetic_dir or kaggle_dir")
    
    # Validation dataset (use -0001 version only, no heavy augmentation)
    val_dataset = None
    if val_df is not None and kaggle_dir:
        val_dataset = KaggleECGDatasetV2(
            kaggle_dir,
            stage1_dir,
            val_df,
            val_transform,
            use_all_versions=False,  # Only use 0001 for validation
            exclude_ids=exclude_ids,
        )
    
    # Samplers for distributed training
    train_sampler = None
    val_sampler = None
    if distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
        if val_dataset:
            val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False)
    
    # Create loaders
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=(train_sampler is None),
        sampler=train_sampler,
        num_workers=num_workers,
        pin_memory=True,
        drop_last=True,
    )
    
    val_loader = None
    if val_dataset:
        val_loader = DataLoader(
            val_dataset,
            batch_size=batch_size,
            shuffle=False,
            sampler=val_sampler,
            num_workers=num_workers,
            pin_memory=True,
        )
    
    return train_loader, val_loader


if __name__ == '__main__':
    # Test the datasets
    print("Testing Dataset V2 classes...")
    
    # Test transforms
    train_tf = get_heavy_augmentation(424, 1088)
    dummy_image = np.random.randint(0, 255, (424, 1088, 3), dtype=np.uint8)
    result = train_tf(image=dummy_image)
    print(f"Transform output shape: {result['image'].shape}")
    
    # Test style-specific transforms
    for style in ['0001', '0006', '0012']:
        tf = get_style_transforms(style, 424, 1088)
        result = tf(image=dummy_image)
        print(f"Style {style} transform: {result['image'].shape}")
    
    print("\nDataset V2 tests completed!")