#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Generate Seal Unwarp Variations
Generates a series of unwarped images for warp factors 1.0 to 6.0
"""
import sys
import cv2
import numpy as np
import os

def generate_unwarps(image_path, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        
    print(f"Loading {image_path}...")
    img = cv2.imread(image_path)
    if img is None:
        print(f"Error: Could not load {image_path}")
        sys.exit(1)
        
    h, w = img.shape[:2]
    cx, cy = w // 2, h // 2
    max_radius = min(cx, cy)
    
    # Sweep warp factors from 1.0 to 6.0 with step 0.5
    factors = np.arange(1.0, 6.5, 0.5)
    
    for factor in factors:
        factor = float(factor)
        
        # Unwarp logic
        out_w = int(max_radius * 2 * 3.14 * factor)
        out_h = int(max_radius)
        
        # Rotate source 90 CCW to put Top at Left (180 deg in polar space)
        rotated_src = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
        
        try:
            unwarped = cv2.warpPolar(rotated_src, (out_w, out_h), (cx, cy), max_radius, cv2.WARP_FILL_OUTLIERS + cv2.WARP_POLAR_LINEAR)
            
            # Crop bottom 30% (Outer ring)
            crop_start_y = int(out_h * 0.70)
            unwarped_crop = unwarped[crop_start_y:out_h, :]
            
            # Flip vertical (feet to center correction)
            unwarped_final = cv2.flip(unwarped_crop, 0)
            
            # Save
            filename = f"factor_{factor:.1f}.png"
            path = os.path.join(output_dir, filename)
            cv2.imwrite(path, unwarped_final)
            print(f"Generated {filename}")
            
        except Exception as e:
            print(f"Error processing factor {factor}: {e}")

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python generate_unwarps.py <image_path> <output_dir>")
        sys.exit(1)
    generate_unwarps(sys.argv[1], sys.argv[2])