#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Create Comparison Grid
Combines unwarped images into a single comparison chart using OpenCV
"""
import sys
import os
import cv2
import numpy as np

def create_grid(input_dir, output_path):
    images = []
    
    # Sort files by factor
    files = sorted([f for f in os.listdir(input_dir) if f.endswith(".png")])
    
    # Filter for reasonable range 1.5 to 4.5
    selected_files = []
    for f in files:
        try:
            factor = float(f.split("_")[1].replace(".png", ""))
            if 1.5 <= factor <= 4.5:
                selected_files.append((factor, f))
        except: pass
        
    selected_files.sort(key=lambda x: x[0])
    
    for factor, filename in selected_files:
        path = os.path.join(input_dir, filename)
        img = cv2.imread(path)
        if img is None: continue
        
        # Resize to fixed height for consistency
        disp_h = 80
        disp_w = int(img.shape[1] * (disp_h / img.shape[0]))
        img_disp = cv2.resize(img, (disp_w, disp_h))
        
        # Add labels
        img_disp = cv2.copyMakeBorder(img_disp, 40, 0, 0, 0, cv2.BORDER_CONSTANT, value=(255, 255, 255))
        cv2.putText(img_disp, f"Factor {factor:.1f}", (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2)
        
        images.append(img_disp)
        
    if not images:
        print("No images found.")
        return

    # Stack vertically
    max_w = max(i.shape[1] for i in images)
    combined = np.zeros((sum(i.shape[0] for i in images), max_w, 3), dtype=np.uint8)
    combined[:] = 255 # White
    
    y = 0
    for img in images:
        h, w = img.shape[:2]
        combined[y:y+h, 0:w] = img
        y += h
        
    cv2.imwrite(output_path, combined)
    print(f"Saved grid to {output_path}")

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python create_comparison_grid.py <input_dir> <output_path>")
        sys.exit(1)
    create_grid(sys.argv[1], sys.argv[2])