#!/usr/bin/env python3
"""
Script to create a 4-panel video from front, side, and top view images.
Layout:
- Top Left: Front view
- Top Right: Side view  
- Bottom Left: Top view
- Bottom Right: Empty (black) for future content
"""

import cv2
import numpy as np
import os
import glob
from pathlib import Path

def create_4panel_video():
    # Define paths
    base_path = Path("/home/warlock/Projects/ScfHeatmapGen/js/2025-10-06_18-00-42")
    front_dir = base_path / "front"
    side_dir = base_path / "side"
    top_dir = base_path / "top"
    output_path = base_path / "combined_4panel_video.mp4"
    
    # Get all image files from each directory
    front_images = sorted(glob.glob(str(front_dir / "*.png")))
    side_images = sorted(glob.glob(str(side_dir / "*.png")))
    top_images = sorted(glob.glob(str(top_dir / "*.png")))
    
    print(f"Found {len(front_images)} front images")
    print(f"Found {len(side_images)} side images")
    print(f"Found {len(top_images)} top images")
    
    # Check if all directories have the same number of images
    if not (len(front_images) == len(side_images) == len(top_images)):
        print("Warning: Different number of images in directories!")
        min_count = min(len(front_images), len(side_images), len(top_images))
        front_images = front_images[:min_count]
        side_images = side_images[:min_count]
        top_images = top_images[:min_count]
        print(f"Using {min_count} images from each directory")
    
    # Read first image to get dimensions
    first_img = cv2.imread(front_images[0])
    if first_img is None:
        print(f"Error: Could not read image {front_images[0]}")
        return
    
    img_height, img_width = first_img.shape[:2]
    print(f"Image dimensions: {img_width}x{img_height}")
    
    # Calculate panel dimensions (2x2 grid)
    panel_width = img_width
    panel_height = img_height
    
    # Create video writer
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    fps = 30  # Adjust frame rate as needed
    
    # Create output video with 2x2 panel layout
    output_width = panel_width * 2
    output_height = panel_height * 2
    
    out = cv2.VideoWriter(str(output_path), fourcc, fps, (output_width, output_height))
    
    print(f"Creating video: {output_width}x{output_height} at {fps} FPS")
    print(f"Output file: {output_path}")
    
    # Process each frame
    for i, (front_img_path, side_img_path, top_img_path) in enumerate(zip(front_images, side_images, top_images)):
        # Read images
        front_img = cv2.imread(front_img_path)
        side_img = cv2.imread(side_img_path)
        top_img = cv2.imread(top_img_path)
        
        if front_img is None or side_img is None or top_img is None:
            print(f"Warning: Could not read images for frame {i+1}")
            continue
        
        # Resize images to panel size if needed
        front_img = cv2.resize(front_img, (panel_width, panel_height))
        side_img = cv2.resize(side_img, (panel_width, panel_height))
        top_img = cv2.resize(top_img, (panel_width, panel_height))
        
        # Create empty panel for bottom right
        empty_panel = np.zeros((panel_height, panel_width, 3), dtype=np.uint8)
        
        # Create 2x2 grid
        # Top row
        top_row = np.hstack([front_img, side_img])
        # Bottom row  
        bottom_row = np.hstack([top_img, empty_panel])
        # Combine rows
        combined_frame = np.vstack([top_row, bottom_row])
        
        # Add labels to each panel
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 1.0
        color = (255, 255, 255)  # White
        thickness = 2
        
        # Add labels
        cv2.putText(combined_frame, "Front View", (10, 30), font, font_scale, color, thickness)
        cv2.putText(combined_frame, "Side View", (panel_width + 10, 30), font, font_scale, color, thickness)
        cv2.putText(combined_frame, "Top View", (10, panel_height + 30), font, font_scale, color, thickness)
        cv2.putText(combined_frame, "Reserved", (panel_width + 10, panel_height + 30), font, font_scale, color, thickness)
        
        # Write frame
        out.write(combined_frame)
        
        # Progress indicator
        if (i + 1) % 10 == 0:
            print(f"Processed {i + 1}/{len(front_images)} frames")
    
    # Release everything
    out.release()
    cv2.destroyAllWindows()
    
    print(f"\nVideo creation completed!")
    print(f"Output saved to: {output_path}")
    print(f"Total frames: {len(front_images)}")
    print(f"Duration: {len(front_images)/fps:.2f} seconds")

if __name__ == "__main__":
    create_4panel_video()