#!/usr/bin/env python3
"""Continue processing synthetic images through stage0/stage1.
Skips already processed images. Uses single GPU (runs alongside training)."""

import os
import sys
import cv2
import numpy as np
import torch
from pathlib import Path
from tqdm import tqdm

# Configuration
INPUT_RAW = '/data/ecg-digitization/synthetic_0001/raw'
INPUT_GT = '/data/ecg-digitization/synthetic_0001/gt'
OUTPUT_RAW = '/data/ecg-digitization/synthetic_0001_processed/raw'
OUTPUT_GT = '/data/ecg-digitization/synthetic_0001_processed/gt'
BASELINE_PATH = '/data/ecg-digitization/hengck23'

def main():
    # Create output dirs
    Path(OUTPUT_RAW).mkdir(parents=True, exist_ok=True)
    Path(OUTPUT_GT).mkdir(parents=True, exist_ok=True)
    
    # Find images that need processing
    all_images = sorted(Path(INPUT_RAW).glob('*.png'))
    processed = set(p.name for p in Path(OUTPUT_RAW).glob('*.png'))
    
    to_process = [p for p in all_images if p.name not in processed]
    
    print(f"Total images: {len(all_images)}")
    print(f"Already processed: {len(processed)}")
    print(f"To process: {len(to_process)}")
    
    if len(to_process) == 0:
        print("All done!")
        return
    
    # Setup device - use GPU 0 only (others are for training)
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    device = torch.device('cuda')
    
    # Import baseline modules
    sys.path.insert(0, BASELINE_PATH)
    from stage0_model import Net as Stage0Net
    from stage0_common import image_to_batch, output_to_predict, normalise_by_homography, load_net
    from stage1_model import Net as Stage1Net
    from stage1_common import output_to_predict as stage1_output_to_predict, rectify_image
    
    # Load models
    print("Loading Stage 0...")
    stage0_net = Stage0Net(pretrained=False)
    stage0_net = load_net(stage0_net, f'{BASELINE_PATH}/weight/stage0-last.checkpoint.pth')
    stage0_net.to(device).eval()
    
    print("Loading Stage 1...")
    stage1_net = Stage1Net(pretrained=False)
    stage1_net = load_net(stage1_net, f'{BASELINE_PATH}/weight/stage1-last.checkpoint.pth')
    stage1_net.to(device).eval()
    
    @torch.no_grad()
    def process_image(image_rgb):
        """Process through stage0 and stage1."""
        # Stage 0
        batch = image_to_batch(image_rgb)
        with torch.amp.autocast('cuda', dtype=torch.float32):
            output = stage0_net(batch)
        rotated, keypoint = output_to_predict(image_rgb, batch, output)
        normalized, _, _ = normalise_by_homography(rotated, keypoint)
        
        # Stage 1
        batch = {'image': torch.from_numpy(np.ascontiguousarray(normalized.transpose(2, 0, 1))).unsqueeze(0)}
        with torch.amp.autocast('cuda', dtype=torch.float32):
            output = stage1_net(batch)
        gridpoint_xy, _ = stage1_output_to_predict(normalized, batch, output)
        rectified = rectify_image(normalized, gridpoint_xy)
        
        return rectified
    
    success = 0
    failed = 0
    
    for img_path in tqdm(to_process, desc="Processing"):
        gt_csv_path = Path(INPUT_GT) / f"{img_path.stem}.csv"
        out_img_path = Path(OUTPUT_RAW) / img_path.name
        out_gt_path = Path(OUTPUT_GT) / f"{img_path.stem}.csv"
        
        # Load image
        image = cv2.imread(str(img_path))
        if image is None:
            failed += 1
            continue
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        try:
            # Process through stage0/stage1
            rectified = process_image(image_rgb)
            
            # Save processed image
            cv2.imwrite(str(out_img_path), cv2.cvtColor(rectified, cv2.COLOR_RGB2BGR))
            
            # Copy GT (no need to transform - it's already in pixel coordinates relative to the final layout)
            if gt_csv_path.exists():
                import shutil
                shutil.copy(gt_csv_path, out_gt_path)
            
            success += 1
        except Exception as e:
            failed += 1
            if failed <= 5:
                print(f"  Failed {img_path.name}: {e}")
    
    print(f"\nDone: {success} success, {failed} failed")
    print(f"Total processed: {len(processed) + success}")


if __name__ == '__main__':
    main()