#!/usr/bin/env python3
"""
Compare crop heights for V15 per-lead training.
Creates side-by-side comparison of CROP_HALF_HEIGHT=200 vs 250.
"""

import os
import random
import numpy as np
import cv2
from pathlib import Path
from tqdm import tqdm

# Constants
TARGET_HEIGHT, TARGET_WIDTH = 1696, 4352
ZERO_MV = np.array([703.5, 987.5, 1271.5, 1531.5])
T0, T1 = 235, 4161
X0, X1 = 0, 2176
Y0, Y1 = 0, 1696

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


def crop_row(image, row_idx, crop_half_height):
    """Crop a single row centered on its baseline."""
    row_height = 2 * crop_half_height
    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, :, :].copy()
    
    # Pad if necessary
    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 create_comparison(image_path, output_dir, sample_idx):
    """Create side-by-side comparison of 200 vs 250 crop heights."""
    # Load image
    image = cv2.imread(str(image_path), cv2.IMREAD_COLOR)
    if image is None:
        return False
    
    image = image[Y0:Y1, X0:X1]
    image = cv2.resize(image, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR)
    
    # Crop to signal region only
    image = image[:, T0:T1, :]
    
    for row_idx in range(4):
        # Crop at 200
        crop_200 = crop_row(image, row_idx, 200)
        
        # Crop at 250
        crop_250 = crop_row(image, row_idx, 250)
        
        # Resize 200 crop to match 250 height for comparison
        crop_200_resized = cv2.resize(crop_200, (crop_250.shape[1], crop_250.shape[0]), 
                                       interpolation=cv2.INTER_LINEAR)
        
        # Add labels
        cv2.putText(crop_200_resized, f"CROP_HALF_HEIGHT=200 (400px total)", (10, 30), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
        cv2.putText(crop_250, f"CROP_HALF_HEIGHT=250 (500px total)", (10, 30), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
        
        # Add row label
        row_labels = ['Row 0: I, aVR, V1, V4', 'Row 1: II, aVL, V2, V5', 
                      'Row 2: III, aVF, V3, V6', 'Row 3: Lead II Rhythm']
        cv2.putText(crop_200_resized, row_labels[row_idx], (10, 60), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
        cv2.putText(crop_250, row_labels[row_idx], (10, 60), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
        
        # Create white separator bar
        separator = np.ones((20, crop_250.shape[1], 3), dtype=np.uint8) * 255
        
        # Stack with white gap in between
        comparison = np.vstack([crop_200_resized, separator, crop_250])
        
        # Save
        sample_id = image_path.parent.name
        variant = image_path.stem.split('-')[-1] if '-' in image_path.stem else '0000'
        output_path = output_dir / f"{sample_idx:04d}_{sample_id}_{variant}_row{row_idx}.png"
        cv2.imwrite(str(output_path), comparison)
    
    return True


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--kaggle_data', type=str, default='/data/ecg-digitization/stage1_data/train')
    parser.add_argument('--output_dir', type=str, default=os.path.expanduser('~/tmp/crop_demo/v15'))
    parser.add_argument('--num_samples', type=int, default=500)
    parser.add_argument('--seed', type=int, default=42)
    args = parser.parse_args()
    
    random.seed(args.seed)
    
    output_dir = Path(args.output_dir)
    if output_dir.exists():
        import shutil
        shutil.rmtree(output_dir)
        print(f"Deleted old output in {output_dir}")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Find all valid images
    kaggle_dir = Path(args.kaggle_data)
    all_images = []
    
    for sample_dir in kaggle_dir.iterdir():
        if not sample_dir.is_dir():
            continue
        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:
                all_images.append(img_path)
    
    print(f"Found {len(all_images)} valid images")
    
    # Random sample
    if len(all_images) > args.num_samples:
        selected = random.sample(all_images, args.num_samples)
    else:
        selected = all_images
    
    print(f"Processing {len(selected)} images (4 rows each = {len(selected) * 4} comparisons)...")
    
    for i, img_path in enumerate(tqdm(selected)):
        try:
            create_comparison(img_path, output_dir, i)
        except Exception as e:
            print(f"Error: {e}")
    
    print(f"\nDone! Saved {len(selected) * 4} comparisons to {output_dir}")
    print(f"\nComparison:")
    print(f"  RED label (top):    CROP_HALF_HEIGHT=200 (400px total row height)")
    print(f"  GREEN label (bottom): CROP_HALF_HEIGHT=250 (500px total row height)")


if __name__ == '__main__':
    main()