%%bash
pip install torch albumentationsx opencv-python matplotlib tqdm segmentation_models_pytorch kaggle

import torch
from torch import nn
from torch import optim
from torch.utils.data import Dataset, DataLoader
import albumentations as A
import numpy as np
import cv2
from pathlib import Path
from tqdm import tqdm
import matplotlib.pyplot as plt
import segmentation_models_pytorch as smp

# Reproducibility
SEED = 137
torch.manual_seed(SEED)
np.random.seed(SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(SEED)

# Device setup
def get_device():
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print(f"Using CUDA: {torch.cuda.get_device_name(0)}")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
        print("Using MPS (Apple Silicon)")
    else:
        device = torch.device("cpu")
        print("Using CPU")
    return device

device = get_device()

def download_kaggle_helen(data_dir="data"):
    """Download HELEN dataset with 68 facial landmarks from Kaggle."""
    try:
        import kaggle
    except ImportError:
        print("Please install kaggle: pip install kaggle")
        return False
    
    data_path = Path(data_dir)
    download_path = data_path / "helen_raw"
    
    try:
        # Download from Kaggle
        print("Downloading HELEN dataset...")
        kaggle.api.dataset_download_files(
            'vietpro/helen-with-68-facial-landmarks',
            path=str(download_path),
            unzip=True
        )
        
        # Find image and landmark files
        img_files = list(download_path.glob("*.jpg"))
        if not img_files:
            print("No images found in downloaded dataset")
            return False
        
        print(f"Found {len(img_files)} images")
        
        # Process files
        samples = []
        for img_file in tqdm(img_files, desc="Processing"):
            pts_file = img_file.with_suffix('.pts')
            if not pts_file.exists():
                continue
            
            # Read image
            image = cv2.imread(str(img_file))
            if image is None:
                continue
            
            # Read landmarks from .pts file
            with open(pts_file, 'r') as f:
                lines = f.readlines()
            
            # Parse coordinates (skip header lines)
            coords = []
            for line in lines:
                line = line.strip()
                if not line or line.startswith(('#', '//', 'version', 'n_points', '{', '}')):
                    continue
                parts = line.split()
                if len(parts) >= 2:
                    coords.append([float(parts[0]), float(parts[1])])
            
            if len(coords) != 68:
                continue
            
            landmarks = np.array(coords, dtype=np.float32)  # [68, 2] as [x, y]
            
            # Sanity check: landmarks should be within image bounds
            h, w = image.shape[:2]
            if landmarks[:, 0].max() > w or landmarks[:, 1].max() > h:
                # Try swapping if needed
                if landmarks[:, 1].max() <= w and landmarks[:, 0].max() <= h:
                    landmarks = landmarks[:, [1, 0]]
                else:
                    continue
            
            samples.append((img_file.name, image, landmarks))
        
        if len(samples) == 0:
            print("No valid samples found")
            return False
        
        print(f"\n✓ Processed {len(samples)} samples")
        
        # Shuffle and split 90/10
        rng = np.random.RandomState(SEED)
        rng.shuffle(samples)
        n_val = len(samples) // 10
        n_train = len(samples) - n_val
        
        train_samples = samples[:n_train]
        val_samples = samples[n_train:]
        
        # Save training set
        train_img_dir = data_path / "train" / "images"
        train_lm_dir = data_path / "train" / "landmarks"
        train_img_dir.mkdir(parents=True, exist_ok=True)
        train_lm_dir.mkdir(parents=True, exist_ok=True)
        
        for i, (name, img, lm) in enumerate(tqdm(train_samples, desc="Saving train")):
            cv2.imwrite(str(train_img_dir / f"face_{i:04d}.jpg"), img)
            np.save(train_lm_dir / f"face_{i:04d}.npy", lm)
        
        # Save validation set
        val_img_dir = data_path / "val" / "images"
        val_lm_dir = data_path / "val" / "landmarks"
        val_img_dir.mkdir(parents=True, exist_ok=True)
        val_lm_dir.mkdir(parents=True, exist_ok=True)
        
        for i, (name, img, lm) in enumerate(tqdm(val_samples, desc="Saving val")):
            cv2.imwrite(str(val_img_dir / f"face_{i:04d}.jpg"), img)
            np.save(val_lm_dir / f"face_{i:04d}.npy", lm)
        
        print(f"\n✓ Dataset ready: {n_train} train, {n_val} val")
        return True
        
    except Exception as e:
        print(f"Error downloading dataset: {e}")
        return False

# Download dataset
DATA_DIR = "data"
TRAIN_IMAGES_DIR = f"{DATA_DIR}/train/images"

if not Path(TRAIN_IMAGES_DIR).exists() or len(list(Path(TRAIN_IMAGES_DIR).glob("*"))) == 0:
    print("Downloading HELEN dataset from Kaggle...")
    success = download_kaggle_helen(DATA_DIR)
    if not success:
        print("❌ Failed to download. Please check Kaggle API setup.")
else:
    print("✓ Dataset already exists")

# Load one sample image for demonstration
DATA_DIR = "data"
np.random.seed(SEED)
all_images = sorted(list(Path(f"{DATA_DIR}/train/images").glob("*.jpg")))
sample_idx = np.random.randint(0, len(all_images))
sample_image_path = all_images[sample_idx]
sample_landmark_path = Path(f"{DATA_DIR}/train/landmarks") / f"{sample_image_path.stem}.npy"

# Load image and landmarks
image = cv2.imread(str(sample_image_path))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
landmarks = np.load(sample_landmark_path)

# Select a few key points to visualize labels
# Note: "left" and "right" are ANATOMICAL (person's left/right, mirrored for viewer)
# When person looks at themselves: their left eye appears on RIGHT side of image to viewer
label_points = {
    36: 'right_eye',   # Person's RIGHT eye → appears on LEFT side of image for viewer
    45: 'left_eye',    # Person's LEFT eye → appears on RIGHT side of image for viewer
    48: 'right_mouth', # Person's RIGHT mouth → appears on LEFT side for viewer
    54: 'left_mouth'   # Person's LEFT mouth → appears on RIGHT side for viewer
}

# Create transform WITHOUT label mapping (standard Albumentations)
transform_standard = A.Compose([
    A.Resize(256, 256, interpolation=cv2.INTER_CUBIC, area_for_downscale="image"),
    A.HorizontalFlip(p=1.0),  # Always flip for demo
], keypoint_params=A.KeypointParams(format='xy', remove_invisible=False), strict=True, seed=SEED)

# Create transform WITH label mapping (AlbumentationsX feature)
# First, let's define a simple mapping just for the key points
FACE_68_HFLIP_MAPPING = {
    # Jawline (0-16): mirror across center
    0: 16, 1: 15, 2: 14, 3: 13, 4: 12, 5: 11, 6: 10, 7: 9, 8: 8,
    16: 0, 15: 1, 14: 2, 13: 3, 12: 4, 11: 5, 10: 6, 9: 7,
    # Eyebrows: left (17-21) ↔ right (22-26)
    17: 26, 18: 25, 19: 24, 20: 23, 21: 22,
    26: 17, 25: 18, 24: 19, 23: 20, 22: 21,
    # Nose bridge (27-30): no swap (central)
    27: 27, 28: 28, 29: 29, 30: 30,
    # Nose bottom: left (31-32) ↔ right (34-35), center (33)
    31: 35, 32: 34, 33: 33, 34: 32, 35: 31,
    # Eyes: left (36-41) ↔ right (42-47)
    36: 45, 37: 44, 38: 43, 39: 42, 40: 47, 41: 46,
    45: 36, 44: 37, 43: 38, 42: 39, 47: 40, 46: 41,
    # Outer lips: left ↔ right, keep centers
    48: 54, 49: 53, 50: 52, 51: 51,
    54: 48, 53: 49, 52: 50,
    55: 59, 56: 58, 57: 57,
    59: 55, 58: 56,
    # Inner lips: left ↔ right, keep centers
    60: 64, 61: 63, 62: 62,
    64: 60, 63: 61,
    65: 67, 66: 66,
    67: 65,
}

transform_with_mapping = A.Compose([
    A.Resize(256, 256, interpolation=cv2.INTER_CUBIC, area_for_downscale="image"),
    A.HorizontalFlip(p=1.0),  # Always flip for demo
], keypoint_params=A.KeypointParams(
    format='xy', 
    label_fields=['keypoint_labels'],
    remove_invisible=False,
    label_mapping={'HorizontalFlip': {'keypoint_labels': FACE_68_HFLIP_MAPPING}}
), strict=True, seed=SEED)

# Apply standard transform
keypoint_labels = np.arange(68)  # 0, 1, 2, ..., 67
result_standard = transform_standard(image=image.copy(), keypoints=landmarks)
image_standard = result_standard['image']
kpts_standard = np.array(result_standard['keypoints'])

# Apply transform with label mapping
result_mapped = transform_with_mapping(
    image=image.copy(), 
    keypoints=landmarks,
    keypoint_labels=keypoint_labels
)
image_mapped = result_mapped['image']
kpts_mapped = np.array(result_mapped['keypoints'])
labels_mapped = np.array(result_mapped['keypoint_labels'])

# Visualize comparison
fig, axes = plt.subplots(1, 3, figsize=(18, 6))

# Original
ax = axes[0]
ax.imshow(cv2.resize(image, (256, 256)))
resized_landmarks = landmarks * (256 / image.shape[1])  # Simple resize
for idx, label_name in label_points.items():
    x, y = resized_landmarks[idx]
    ax.scatter(x, y, c='lime', s=100, edgecolors='black', linewidths=2)
    ax.text(x, y-10, label_name, fontsize=10, color='lime', weight='bold',
            ha='center', bbox=dict(boxstyle='round,pad=0.3', facecolor='black', alpha=0.7))
ax.set_title('Original Image', fontsize=12)
ax.axis('off')

# Standard HFlip (coordinates flipped, but labels stay the same!)
ax = axes[1]
ax.imshow(image_standard)
for idx, label_name in label_points.items():
    x, y = kpts_standard[idx]  # WRONG: "left_eye" is now physically on the right!
    color = 'red'
    ax.scatter(x, y, c=color, s=100, edgecolors='black', linewidths=2)
    ax.text(x, y-10, label_name, fontsize=10, color=color, weight='bold',
            ha='center', bbox=dict(boxstyle='round,pad=0.3', facecolor='black', alpha=0.7))
ax.set_title('❌ Standard HFlip\n("left_eye" is anatomically the RIGHT eye now - WRONG!)', fontsize=12, color='red')
ax.axis('off')

# With label mapping (coordinates AND labels swapped!)
ax = axes[2]
ax.imshow(image_mapped)
# Create reverse mapping to find which label is at each position
idx_to_label = {}
for idx, label_name in label_points.items():
    # Find where this label ended up after mapping
    mapped_idx = labels_mapped[idx]
    idx_to_label[int(mapped_idx)] = label_name

for orig_idx, label_name in label_points.items():
    # After mapping, find where this label ended up
    new_idx = int(labels_mapped[orig_idx])
    x, y = kpts_mapped[new_idx]
    color = 'lime'
    ax.scatter(x, y, c=color, s=100, edgecolors='black', linewidths=2)
    # Show the semantic label (it's now in the correct position!)
    ax.text(x, y-10, label_name, fontsize=10, color=color, weight='bold',
            ha='center', bbox=dict(boxstyle='round,pad=0.3', facecolor='black', alpha=0.7))
ax.set_title('✅ HFlip + Label Mapping\n(Labels match anatomy - CORRECT!)', fontsize=12, color='green')
ax.axis('off')

plt.suptitle('AlbumentationsX Feature: Keypoint Label Remapping During HorizontalFlip', 
             fontsize=14, weight='bold', y=0.98)
plt.tight_layout()
plt.show()

# Print what happened
print("\n" + "="*70)
print("WHAT HAPPENED:")
print("="*70)
print("\nOriginal image:")
print(f"  - 'left_eye' (idx 36):  at pixel ({resized_landmarks[36][0]:.0f}, {resized_landmarks[36][1]:.0f})")
print(f"  - 'right_eye' (idx 45): at pixel ({resized_landmarks[45][0]:.0f}, {resized_landmarks[45][1]:.0f})")

print("\n❌ Standard HFlip (PROBLEM):")
print(f"  - 'left_eye' label: at pixel ({kpts_standard[36][0]:.0f}, {kpts_standard[36][1]:.0f})")
print(f"  - 'right_eye' label: at pixel ({kpts_standard[45][0]:.0f}, {kpts_standard[45][1]:.0f})")
print("  ⚠️ After flip, 'left_eye' is anatomically the RIGHT eye of flipped person - labels are WRONG!")

print("\n✅ HFlip + Label Mapping (FIXED):")
# Find where left_eye and right_eye labels ended up
left_eye_new_idx = int(labels_mapped[36])
right_eye_new_idx = int(labels_mapped[45])
print(f"  - 'left_eye' label: at pixel ({kpts_mapped[left_eye_new_idx][0]:.0f}, {kpts_mapped[left_eye_new_idx][1]:.0f})")
print(f"  - 'right_eye' label: at pixel ({kpts_mapped[right_eye_new_idx][0]:.0f}, {kpts_mapped[right_eye_new_idx][1]:.0f})")
print(f"\n  Explanation:")
print(f"    - Original array[36] ('left_eye') → swapped to array[{FACE_68_HFLIP_MAPPING[36]}] ('right_eye')")
print(f"    - Original array[45] ('right_eye') → swapped to array[{FACE_68_HFLIP_MAPPING[45]}] ('left_eye')")
print(f"    - Now 'left_eye' label points to the person's anatomical left eye in flipped image!")
print("  ✓ Labels are anatomically correct for the flipped person!")
print("="*70)

def get_train_transforms(image_size=256, aug_type="medium"):
    """Get training augmentation pipeline."""
    # area_for_downscale="image": uses cv2.INTER_AREA for downscaling, cv2.INTER_CUBIC otherwise (reduces artifacts)
    base = [A.Resize(image_size, image_size, interpolation=cv2.INTER_CUBIC, area_for_downscale="image")]
    
    if aug_type == "minimal":
        # Minimal: HorizontalFlip with label swapping (demonstrates the key feature!)
        augs = [
            A.HorizontalFlip(p=0.5),
        ]
        label_mapping = {
            'HorizontalFlip': {
                'keypoint_labels': FACE_68_HFLIP_MAPPING
            }
        }
    
    elif aug_type == "medium":
        # Balanced augmentations with keypoint label swapping
        augs = [
            # Basic geometric (with label swapping for symmetric keypoints!)
            A.HorizontalFlip(p=0.5),
            
            # Dropout/Occlusion (HIGH IMPACT - teaches robustness to occlusions)
            A.CoarseDropout(
                num_holes_range=(1, 5),
                hole_height_range=(0.05, 0.15),  # 5-15% of image size
                hole_width_range=(0.05, 0.15),
                fill=0,
                p=0.3
            ),
            
            # Affine transformations
            A.Affine(
                scale=(0.8, 1.2),
                rotate=(-20, 20),
                keep_ratio=True,
                p=0.7
            ),
            
            # Color augmentations (realistic lighting conditions)
            A.OneOf([
                A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1, p=1.0),
                A.PlanckianJitter(mode='blackbody', p=1.0),  # Realistic temperature shifts
            ], p=0.5),
            
            # Image quality degradation (JPEG compression)
            A.ImageCompression(quality_range=(50, 95), compression_type='jpeg', p=0.3),
            
            # Blur and noise (camera effects)
            A.GaussNoise(std_range=(0.02, 0.1), p=0.3),
            A.OneOf([
                A.Blur(blur_limit=3, p=1.0),
                A.MedianBlur(blur_limit=3, p=1.0),
            ], p=0.2),
        ]
        label_mapping = {
            'HorizontalFlip': {
                'keypoint_labels': FACE_68_HFLIP_MAPPING
            }
        }
    else:
        raise ValueError(f"Unknown aug_type: {aug_type}")
    
    # Normalization and tensor conversion (always applied)
    post = [
        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        A.ToTensorV2(),
    ]
    
    return A.Compose(
        base + augs + post,
        keypoint_params=A.KeypointParams(
            format='xy',
            label_fields=['keypoint_labels'] if label_mapping else None,
            remove_invisible=False,
            label_mapping=label_mapping
        ),
        strict=True,
        seed=SEED
    )


def get_val_transforms(image_size=256):
    """Validation transforms (no augmentation)."""
    return A.Compose([
        A.Resize(image_size, image_size, interpolation=cv2.INTER_CUBIC, area_for_downscale="image"),
        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        A.ToTensorV2(),
    ], keypoint_params=A.KeypointParams(
        format='xy',
        label_fields=['keypoint_labels'],
        remove_invisible=False,
        label_mapping={}  # Empty mapping to suppress warning (no label swapping in validation)
    ), strict=True, seed=SEED)

class FaceLandmarkDataset(Dataset):
    """Dataset for face landmarks."""
    def __init__(self, images_dir, landmarks_dir, transform=None, image_size=256):
        self.images_dir = Path(images_dir)
        self.landmarks_dir = Path(landmarks_dir)
        self.transform = transform
        self.image_size = image_size
        self.image_files = sorted(list(self.images_dir.glob("*.jpg")) + 
                                  list(self.images_dir.glob("*.png")))
    
    def __len__(self):
        return len(self.image_files)
    
    def __getitem__(self, idx):
        # Load image and landmarks
        img_path = self.image_files[idx]
        image = cv2.imread(str(img_path))
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        landmark_path = self.landmarks_dir / f"{img_path.stem}.npy"
        landmarks = np.load(landmark_path).astype(np.float32)  # [68, 2]
        
        # Create keypoint labels for label swapping
        keypoint_labels = np.arange(len(landmarks))
        
        # Apply augmentations
        transformed = self.transform(
            image=image,
            keypoints=landmarks,
            keypoint_labels=keypoint_labels
        )
        
        image = transformed["image"]
        landmarks = transformed["keypoints"]
        landmarks_tensor = torch.from_numpy(np.array(landmarks, dtype=np.float32))
        
        return image, landmarks_tensor


def visualize_augmentations(aug_type, num_samples=4):
    """Visualize how augmentations affect the images."""
    dataset = FaceLandmarkDataset(
        f"{DATA_DIR}/train/images",
        f"{DATA_DIR}/train/landmarks",
        transform=get_train_transforms(256, aug_type),
        image_size=256
    )
    
    fig, axes = plt.subplots(2, num_samples // 2, figsize=(15, 8))
    axes = axes.flatten()
    
    np.random.seed(SEED)
    indices = np.random.choice(len(dataset), num_samples, replace=False)
    
    for idx, ax in zip(indices, axes):
        image, landmarks = dataset[idx]
        
        # Denormalize image for display
        image_np = image.permute(1, 2, 0).cpu().numpy()
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        image_np = np.clip(image_np * std + mean, 0, 1)
        
        ax.imshow(image_np)
        ax.scatter(landmarks[:, 0], landmarks[:, 1], c='red', s=5, alpha=0.7)
        ax.axis('off')
        ax.set_title(f'{aug_type.capitalize()} Augmentation')
    
    plt.tight_layout()
    plt.show()

# Visualize minimal augmentation (HFlip with label swapping)
print("=" * 60)
print("MINIMAL Augmentation")
print("=" * 60)
print("Just HorizontalFlip with label_mapping (keypoint reordering)")
print("Notice how symmetric keypoints swap positions correctly!")
visualize_augmentations("minimal", num_samples=4)

# Visualize medium augmentation
print("\n" + "=" * 60)
print("MEDIUM Augmentation")
print("=" * 60)
print("Full pipeline: HFlip + CoarseDropout + Affine + Color + Compression + Blur + Noise")
visualize_augmentations("medium", num_samples=4)

class HeatmapModel(nn.Module):
    """U-Net model for heatmap-based landmark detection."""
    def __init__(self, num_keypoints=68, heatmap_size=256, encoder_name="resnet50"):
        super().__init__()
        self.heatmap_size = heatmap_size
        
        # U-Net from segmentation_models_pytorch
        self.unet = smp.Unet(
            encoder_name=encoder_name,
            encoder_weights="imagenet",
            in_channels=3,
            classes=num_keypoints,
            decoder_channels=(256, 128, 64, 32, 16),
            activation=None,
        )
    
    def forward(self, x):
        heatmaps = self.unet(x)
        
        # Resize if needed
        if self.heatmap_size != x.shape[-1]:
            heatmaps = nn.functional.interpolate(
                heatmaps,
                size=(self.heatmap_size, self.heatmap_size),
                mode='bilinear',
                align_corners=False
            )
        
        return heatmaps

def generate_heatmaps(landmarks, image_size=256, heatmap_size=256, sigma=8.0):
    """Generate Gaussian heatmaps from landmark coordinates."""
    B, K, _ = landmarks.shape
    device = landmarks.device
    H = W = heatmap_size
    
    # Scale landmarks to heatmap size
    scale = heatmap_size / image_size
    landmarks_scaled = landmarks * scale
    
    # Create coordinate grids
    x = torch.arange(W, device=device, dtype=torch.float32)
    y = torch.arange(H, device=device, dtype=torch.float32)
    yy, xx = torch.meshgrid(y, x, indexing='ij')
    
    # Generate Gaussians
    cx = landmarks_scaled[:, :, 0].view(B, K, 1, 1)
    cy = landmarks_scaled[:, :, 1].view(B, K, 1, 1)
    
    return torch.exp(-((xx - cx)**2 + (yy - cy)**2) / (2 * sigma**2))

def soft_argmax_2d(heatmaps, temperature=1.0):
    """Extract sub-pixel coordinates from heatmaps using differentiable soft-argmax."""
    B, K, H, W = heatmaps.shape
    device = heatmaps.device
    
    # Flatten and apply softmax
    heatmaps_flat = heatmaps.view(B, K, -1) / temperature
    probs = torch.softmax(heatmaps_flat, dim=-1)
    
    # Create coordinate grids
    x_coords = torch.arange(W, device=device, dtype=torch.float32)
    y_coords = torch.arange(H, device=device, dtype=torch.float32)
    yy, xx = torch.meshgrid(y_coords, x_coords, indexing='ij')
    
    xx = xx.flatten()
    yy = yy.flatten()
    
    # Weighted sum: E[x] = Σ(x * p(x))
    x = (probs * xx.view(1, 1, -1)).sum(dim=-1)
    y = (probs * yy.view(1, 1, -1)).sum(dim=-1)
    
    return torch.stack([x, y], dim=-1)

class WingLoss(nn.Module):
    """Wing Loss for robust coordinate regression."""
    def __init__(self, omega=10, epsilon=2):
        super().__init__()
        self.omega = omega
        self.epsilon = epsilon
        self.C = self.omega - self.omega * np.log(1 + self.omega / self.epsilon)
    
    def forward(self, pred, target):
        delta = (target - pred).abs()
        loss = torch.where(
            delta < self.omega,
            self.omega * torch.log(1 + delta / self.epsilon),
            delta - self.C
        )
        return loss.mean()

def train_epoch(model, dataloader, heatmap_criterion, coord_criterion, optimizer, device, epoch, 
                image_size=256, heatmap_size=256, coord_loss_weight=0.1):
    """Train for one epoch."""
    model.train()
    running_total_loss = 0.0
    running_heatmap_loss = 0.0
    running_coord_loss = 0.0
    
    pbar = tqdm(dataloader, desc=f"Epoch {epoch}", leave=False)
    for images, landmarks in pbar:
        images = images.to(device)
        landmarks = landmarks.to(device)
        
        optimizer.zero_grad()
        
        # Generate target heatmaps
        target_heatmaps = generate_heatmaps(landmarks, image_size, heatmap_size, sigma=8.0)
        
        # Forward pass
        pred_heatmaps = model(images)
        heatmap_loss = heatmap_criterion(pred_heatmaps, target_heatmaps)
        
        # Auxiliary coordinate loss
        pred_coords = soft_argmax_2d(pred_heatmaps)
        scale = image_size / heatmap_size
        pred_coords_pixels = pred_coords * scale
        coord_loss = coord_criterion(pred_coords_pixels, landmarks)
        
        # Combined loss
        total_loss = heatmap_loss + coord_loss_weight * coord_loss
        total_loss.backward()
        optimizer.step()
        
        running_total_loss += total_loss.item()
        running_heatmap_loss += heatmap_loss.item()
        running_coord_loss += coord_loss.item()
        
        pbar.set_postfix({
            'total': f'{total_loss.item():.4f}',
            'heatmap': f'{heatmap_loss.item():.4f}',
            'coord': f'{coord_loss.item():.2f}'
        })
    
    return (running_total_loss / len(dataloader), 
            running_heatmap_loss / len(dataloader),
            running_coord_loss / len(dataloader))

def validate(model, dataloader, criterion, device, image_size=256, heatmap_size=256):
    """Validate model: returns (heatmap_loss, mean_pixel_error, nme_percentage)."""
    model.eval()
    running_loss = 0.0
    total_pixel_error = 0.0
    total_nme = 0.0
    num_samples = 0
    
    with torch.no_grad():
        for images, landmarks in tqdm(dataloader, desc="Validating", leave=False):
            images = images.to(device)
            landmarks = landmarks.to(device)
            
            # Generate target heatmaps and compute loss
            target_heatmaps = generate_heatmaps(landmarks, image_size, heatmap_size, sigma=8.0)
            pred_heatmaps = model(images)
            loss = criterion(pred_heatmaps, target_heatmaps)
            running_loss += loss.item()
            
            # Extract coordinates and compute metrics
            pred_coords_heatmap = soft_argmax_2d(pred_heatmaps)
            pred_coords_pixels = pred_coords_heatmap * (image_size / heatmap_size)
            
            # Pixel error
            pixel_errors = torch.norm(pred_coords_pixels - landmarks, dim=-1)
            total_pixel_error += pixel_errors.mean().item() * images.shape[0]
            
            # NME (normalized by inter-ocular distance: landmarks 36 and 45)
            inter_ocular = torch.norm(landmarks[:, 36] - landmarks[:, 45], dim=-1)
            nme = (pixel_errors.mean(dim=1) / inter_ocular * 100).mean().item()
            total_nme += nme * images.shape[0]
            
            num_samples += images.shape[0]
    
    return running_loss / len(dataloader), total_pixel_error / num_samples, total_nme / num_samples

def visualize_predictions(model, dataset, device, num_samples=4, image_size=256, heatmap_size=256):
    """Visualize model predictions."""
    model.eval()
    fig, axes = plt.subplots(2, num_samples // 2, figsize=(15, 8))
    axes = axes.flatten()
    
    np.random.seed(SEED)
    indices = np.random.choice(len(dataset), num_samples, replace=False)
    
    with torch.no_grad():
        for idx, ax in zip(indices, axes):
            image, landmarks_gt = dataset[idx]
            
            # Get prediction
            image_input = image.unsqueeze(0).to(device)
            pred_heatmaps = model(image_input)
            
            # Extract coordinates
            pred_coords_heatmap = soft_argmax_2d(pred_heatmaps)
            scale = image_size / heatmap_size
            landmarks_pred = pred_coords_heatmap.squeeze(0).cpu() * scale
            landmarks_gt = landmarks_gt.cpu()
            
            # Denormalize image
            image_np = image.permute(1, 2, 0).cpu().numpy()
            mean = np.array([0.485, 0.456, 0.406])
            std = np.array([0.229, 0.224, 0.225])
            image_np = np.clip(image_np * std + mean, 0, 1)
            
            ax.imshow(image_np)
            ax.scatter(landmarks_gt[:, 0], landmarks_gt[:, 1], c='green', s=10, label='GT')
            ax.scatter(landmarks_pred[:, 0], landmarks_pred[:, 1], c='red', s=10, marker='x', label='Pred')
            ax.axis('off')
            if idx == indices[0]:
                ax.legend()
    
    plt.tight_layout()
    plt.show()

def train_model(aug_type="minimal", num_epochs=150):
    """Train model with specified augmentation type."""
    print(f"\n{'='*60}")
    print(f"Training with {aug_type.upper()} augmentation")
    print(f"{'='*60}\n")
    
    # Hyperparameters
    IMAGE_SIZE = 256
    HEATMAP_SIZE = 256
    NUM_KEYPOINTS = 68
    BATCH_SIZE = 64
    LEARNING_RATE = 1e-3
    
    # Create datasets
    train_dataset = FaceLandmarkDataset(
        f"{DATA_DIR}/train/images",
        f"{DATA_DIR}/train/landmarks",
        transform=get_train_transforms(IMAGE_SIZE, aug_type),
        image_size=IMAGE_SIZE
    )
    val_dataset = FaceLandmarkDataset(
        f"{DATA_DIR}/val/images",
        f"{DATA_DIR}/val/landmarks",
        transform=get_val_transforms(IMAGE_SIZE),
        image_size=IMAGE_SIZE
    )
    print(f"Dataset: {len(train_dataset)} train, {len(val_dataset)} val")
    
    # DataLoaders (num_workers=0 for Jupyter compatibility)
    # In a regular Python script, you can use num_workers=4 for faster loading
    g = torch.Generator().manual_seed(SEED)
    
    train_loader = DataLoader(
        train_dataset, batch_size=BATCH_SIZE, shuffle=True,
        num_workers=0, generator=g
    )
    val_loader = DataLoader(
        val_dataset, batch_size=BATCH_SIZE, shuffle=False,
        num_workers=0
    )
    
    # Create model
    model = HeatmapModel(
        num_keypoints=NUM_KEYPOINTS,
        heatmap_size=HEATMAP_SIZE,
        encoder_name="resnet50"
    ).to(device)
    
    num_params = sum(p.numel() for p in model.parameters())
    print(f"Model: U-Net (ResNet50) → {NUM_KEYPOINTS} heatmaps @ {HEATMAP_SIZE}x{HEATMAP_SIZE}")
    print(f"Parameters: {num_params:,}")
    
    # Setup training
    heatmap_criterion = nn.MSELoss()
    coord_criterion = WingLoss(omega=10, epsilon=2)
    coord_loss_weight = 0.1
    
    optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.5, patience=5)
    
    print(f"Loss: MSE (heatmap) + WingLoss (coord, weight={coord_loss_weight})")
    print(f"LR: {LEARNING_RATE}, Batch: {BATCH_SIZE}, Epochs: {num_epochs}\n")
    
    # Training loop
    best_val_loss = float('inf')
    best_nme = float('inf')
    patience_counter = 0
    early_stop_patience = 20  # Stop if no improvement for 20 epochs
    train_losses, val_losses, val_pixel_errors, val_nmes = [], [], [], []
    
    for epoch in range(1, num_epochs + 1):
        train_loss_total, train_loss_heatmap, train_loss_coord = train_epoch(
            model, train_loader, heatmap_criterion, coord_criterion, 
            optimizer, device, epoch,
            IMAGE_SIZE, HEATMAP_SIZE, coord_loss_weight
        )
        val_loss, val_pixel_error, val_nme = validate(
            model, val_loader, heatmap_criterion, device,
            IMAGE_SIZE, HEATMAP_SIZE
        )
        
        train_losses.append(train_loss_total)
        val_losses.append(val_loss)
        val_pixel_errors.append(val_pixel_error)
        val_nmes.append(val_nme)
        
        # Print every 10 epochs (less verbose for notebook)
        if epoch % 10 == 0 or epoch == 1 or epoch == num_epochs:
            print(f"Epoch {epoch}/{num_epochs} - "
                  f"Val Loss: {val_loss:.4f} | "
                  f"Pixel Error: {val_pixel_error:.2f}px | "
                  f"NME: {val_nme:.2f}%")
        
        scheduler.step(val_loss)
        
        # Track best metrics and early stopping
        if val_nme < best_nme:
            best_nme = val_nme
            patience_counter = 0
        else:
            patience_counter += 1
        
        if val_loss < best_val_loss:
            best_val_loss = val_loss
        
        # Early stopping
        if patience_counter >= early_stop_patience:
            print(f"\n⚠️ Early stopping at epoch {epoch} (no improvement for {early_stop_patience} epochs)")
            break
    
    # Plot training history
    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
    
    axes[0].plot(train_losses, label='Train'); axes[0].plot(val_losses, label='Val')
    axes[0].set(xlabel='Epoch', ylabel='Heatmap Loss (MSE)', ylim=(0, None), title='Training Loss')
    axes[0].legend(); axes[0].grid(True)
    
    axes[1].plot(val_pixel_errors, color='orange')
    axes[1].set(xlabel='Epoch', ylabel='Mean Pixel Error', ylim=(0, None), title='Pixel-Space Error')
    axes[1].grid(True)
    
    axes[2].plot(val_nmes, color='green', label='Val NME')
    axes[2].axhline(y=5.0, color='r', linestyle='--', label='SOTA (~5%)')
    axes[2].set(xlabel='Epoch', ylabel='NME (%)', ylim=(0, None), title='Normalized Mean Error')
    axes[2].legend(); axes[2].grid(True)
    
    plt.suptitle(f'{aug_type.upper()} Augmentation', fontsize=16, y=1.02)
    plt.tight_layout()
    plt.show()
    
    print(f"\n✓ Training complete! Best NME: {best_nme:.2f}% (SOTA: ~3-5%)")
    
    # Visualize predictions
    print(f"\nFinal predictions ({aug_type}):")
    visualize_predictions(model, val_dataset, device, num_samples=4, 
                         image_size=IMAGE_SIZE, heatmap_size=HEATMAP_SIZE)
    
    return model, {
        'train_losses': train_losses,
        'val_losses': val_losses,
        'val_pixel_errors': val_pixel_errors,
        'val_nmes': val_nmes,
        'best_nme': best_nme
    }

# Train for up to 300 epochs (will stop early if no improvement for 20 epochs)
# For minimal augmentation, typically stops around epoch 110-120
model_minimal, results_minimal = train_model(aug_type="minimal", num_epochs=300)

# Train for up to 300 epochs (will stop early if no improvement for 20 epochs)
# For medium augmentation, typically needs 250+ epochs
model_medium, results_medium = train_model(aug_type="medium", num_epochs=300)

# Compare final NME
print("\n" + "="*60)
print("FINAL COMPARISON")
print("="*60)
print(f"Minimal Augmentation:  Best NME = {results_minimal['best_nme']:.2f}%")
print(f"Medium Augmentation:   Best NME = {results_medium['best_nme']:.2f}%")
print(f"Improvement:           {results_minimal['best_nme'] - results_medium['best_nme']:.2f}% reduction")
print("\nSOTA (State-of-the-Art) on HELEN: ~3-5% NME")
print("="*60)

# Plot side-by-side comparison of training curves (log scale for better visibility)
fig, axes = plt.subplots(1, 3, figsize=(18, 5))

# Loss comparison (log scale)
axes[0].plot(results_minimal['val_losses'], label='Minimal', alpha=0.7, linewidth=2)
axes[0].plot(results_medium['val_losses'], label='Medium', alpha=0.7, linewidth=2)
axes[0].set(xlabel='Epoch', ylabel='Validation Loss (log scale)', title='Validation Loss')
axes[0].set_yscale('log')
axes[0].legend(); axes[0].grid(True, which='both', alpha=0.3)

# Pixel error comparison (log scale)
axes[1].plot(results_minimal['val_pixel_errors'], label='Minimal', alpha=0.7, linewidth=2)
axes[1].plot(results_medium['val_pixel_errors'], label='Medium', alpha=0.7, linewidth=2)
axes[1].set(xlabel='Epoch', ylabel='Mean Pixel Error (log scale)', title='Pixel-Space Error')
axes[1].set_yscale('log')
axes[1].legend(); axes[1].grid(True, which='both', alpha=0.3)

# NME comparison (log scale)
axes[2].plot(results_minimal['val_nmes'], label='Minimal', alpha=0.7, linewidth=2)
axes[2].plot(results_medium['val_nmes'], label='Medium', alpha=0.7, linewidth=2)
axes[2].axhline(y=5.0, color='r', linestyle='--', label='SOTA (~5%)', alpha=0.5)
axes[2].set(xlabel='Epoch', ylabel='NME (%) - log scale', title='Normalized Mean Error')
axes[2].set_yscale('log')
axes[2].legend(); axes[2].grid(True, which='both', alpha=0.3)

plt.suptitle('Minimal vs Medium Augmentation (Log Scale)', fontsize=16, y=1.02)
plt.tight_layout()
plt.show()

def compare_predictions(model_minimal, model_medium, dataset, device, num_samples=4, 
                       image_size=256, heatmap_size=256):
    """Compare predictions from two models on the same images."""
    fig, axes = plt.subplots(2, num_samples, figsize=(20, 10))
    
    # Use fixed indices for reproducible comparison
    np.random.seed(SEED)
    indices = np.random.choice(len(dataset), num_samples, replace=False)
    
    model_minimal.eval()
    model_medium.eval()
    
    with torch.no_grad():
        for col_idx, img_idx in enumerate(indices):
            image, landmarks_gt = dataset[img_idx]
            
            # Get predictions from both models
            image_input = image.unsqueeze(0).to(device)
            
            # Minimal model
            pred_heatmaps_min = model_minimal(image_input)
            pred_coords_min = soft_argmax_2d(pred_heatmaps_min)
            scale = image_size / heatmap_size
            landmarks_pred_min = pred_coords_min.squeeze(0).cpu() * scale
            
            # Medium model
            pred_heatmaps_med = model_medium(image_input)
            pred_coords_med = soft_argmax_2d(pred_heatmaps_med)
            landmarks_pred_med = pred_coords_med.squeeze(0).cpu() * scale
            
            landmarks_gt = landmarks_gt.cpu()
            
            # Denormalize image
            image_np = image.permute(1, 2, 0).cpu().numpy()
            mean = np.array([0.485, 0.456, 0.406])
            std = np.array([0.229, 0.224, 0.225])
            image_np = np.clip(image_np * std + mean, 0, 1)
            
            # Compute errors for this image
            error_min = torch.norm(landmarks_pred_min - landmarks_gt, dim=-1).mean().item()
            error_med = torch.norm(landmarks_pred_med - landmarks_gt, dim=-1).mean().item()
            
            # Plot minimal model
            ax = axes[0, col_idx]
            ax.imshow(image_np)
            ax.scatter(landmarks_gt[:, 0], landmarks_gt[:, 1], c='green', s=15, label='GT', alpha=0.6)
            ax.scatter(landmarks_pred_min[:, 0], landmarks_pred_min[:, 1], c='red', s=15, 
                      marker='x', label='Pred', linewidths=2)
            ax.set_title(f'Minimal\nError: {error_min:.2f}px', fontsize=10)
            ax.axis('off')
            if col_idx == 0:
                ax.legend(loc='upper left', fontsize=8)
            
            # Plot medium model
            ax = axes[1, col_idx]
            ax.imshow(image_np)
            ax.scatter(landmarks_gt[:, 0], landmarks_gt[:, 1], c='green', s=15, label='GT', alpha=0.6)
            ax.scatter(landmarks_pred_med[:, 0], landmarks_pred_med[:, 1], c='blue', s=15, 
                      marker='x', label='Pred', linewidths=2)
            ax.set_title(f'Medium\nError: {error_med:.2f}px', fontsize=10)
            ax.axis('off')
            if col_idx == 0:
                ax.legend(loc='upper left', fontsize=8)
    
    plt.suptitle('Prediction Comparison: Minimal (top) vs Medium (bottom)', fontsize=14, y=0.98)
    plt.tight_layout()
    plt.show()

# Load validation dataset (without augmentation)
val_dataset_compare = FaceLandmarkDataset(
    f"{DATA_DIR}/val/images",
    f"{DATA_DIR}/val/landmarks",
    transform=get_val_transforms(256),
    image_size=256
)

# Compare predictions on the same images
compare_predictions(model_minimal, model_medium, val_dataset_compare, device, 
                   num_samples=4, image_size=256, heatmap_size=256)