app/utils/comparison.py

"""
Comparison module for frame-by-frame analysis between user and pro swings

CRITICAL NOTE: This module preserves the original sizes and orientations of both user and professional videos.
Frames are saved as separate image files at their original resolutions without any resizing, rotation, or distortion.
"""

import os
import cv2
import numpy as np
from tqdm import tqdm
from PIL import Image


def ensure_color_frame(frame):
    """
    Ensure frame is in color format (3 channels)
    
    Args:
        frame (numpy.ndarray): Input frame
        
    Returns:
        numpy.ndarray: Color frame with 3 channels
    """
    if frame is None:
        return np.zeros((480, 640, 3), dtype=np.uint8)
    
    # If frame is grayscale (2D), convert to color (3D)
    if len(frame.shape) == 2:
        frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
    elif len(frame.shape) == 3 and frame.shape[2] == 1:
        frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
    elif len(frame.shape) == 3 and frame.shape[2] == 4:
        # Convert RGBA to BGR
        frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2BGR)
    
    return frame


def resize_frame_proportionally(frame, target_height):
    """
    Resize frame proportionally to target height while maintaining aspect ratio
    
    Args:
        frame (numpy.ndarray): Input frame
        target_height (int): Target height
        
    Returns:
        numpy.ndarray: Resized frame
    """
    # Ensure frame is in color format
    frame = ensure_color_frame(frame)
    
    h, w = frame.shape[:2]
    if h == 0:
        return np.zeros((target_height, target_height, 3), dtype=np.uint8)
    
    # Calculate new width to maintain aspect ratio
    target_width = int(w * (target_height / h))
    
    # Resize the frame
    return cv2.resize(frame, (target_width, target_height))


def extract_frames(video_path, max_frames=100):
    """
    Extract frames from a video
    
    Args:
        video_path (str): Path to the video file
        max_frames (int): Maximum number of frames to extract
        
    Returns:
        list: List of extracted frames as numpy arrays
    """
    frames = []
    
    if not os.path.exists(video_path):
        raise ValueError(f"Video file not found: {video_path}")
        
    # Use standard OpenCV VideoCapture with explicit settings to prevent any rotation
    cap = cv2.VideoCapture(video_path)
    
    if not cap.isOpened():
        raise ValueError(f"Could not open video: {video_path}")
    
    # CRITICAL: Explicitly disable ALL automatic transformations
    # This prevents OpenCV from applying any rotation based on metadata
    try:
        cap.set(cv2.CAP_PROP_ORIENTATION_AUTO, 0)  # Disable auto-orientation
        cap.set(cv2.CAP_PROP_ORIENTATION_META, 0)  # Ignore orientation metadata
        cap.set(cv2.CAP_PROP_CONVERT_RGB, 0)       # Keep BGR format
    except:
        # If properties are not supported, continue without them
        pass
    
    # Get total frame count
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    # Calculate step to get approximately max_frames
    step = max(1, total_frames // max_frames)
    
    current_frame = 0
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
            
        if current_frame % step == 0:
            # Store frame exactly as read from video - no transformations at all
            # Only verify it's a valid color frame before storing
            if frame is not None and len(frame.shape) == 3:
                frames.append(frame.copy())
            
        current_frame += 1
        
    cap.release()
    
    return frames


def extract_key_swing_frames(video_path, frames, swing_phases=None):
    """
    Extract 3 key frames from a list of processed frames.
    1. First setup frame
    2. Last backswing frame (top of backswing)
    3. First impact frame
    
    Args:
        video_path (str): Path to the original video file (used for rotation metadata).
        frames (list): List of processed video frames.
        swing_phases (dict): Dictionary mapping phase names to lists of frame indices
                             relative to the 'frames' list.
        
    Returns:
        dict: Dictionary mapping phase names to frames
    """
    key_frames = {'setup': None, 'backswing': None, 'impact': None}
    
    if not frames:
        print("Warning: No frames provided to extract_key_swing_frames.")
        return key_frames
        
    # Determine frame indices based on swing phases
    if swing_phases:
        # Get first setup frame
        setup_frames = swing_phases.get('setup', [])
        setup_idx = setup_frames[0] if setup_frames else 0
        
        # Get last backswing frame (top of backswing)
        backswing_frames = swing_phases.get('backswing', [])
        backswing_idx = backswing_frames[-1] if backswing_frames else len(frames) // 3
        
        # Get first impact frame
        impact_frames = swing_phases.get('impact', [])
        impact_idx = impact_frames[0] if impact_frames else len(frames) // 2
    else:
        # Fallback to default indices if no swing phases provided
        setup_idx = 0
        backswing_idx = len(frames) // 3
        impact_idx = len(frames) // 2
    
    print(f"Key frame indices (relative to processed frames) - Setup: {setup_idx}, Backswing: {backswing_idx}, Impact: {impact_idx}")
    print(f"These correspond to original video frames (approx) - Setup: ~{setup_idx * 1}, Backswing: ~{backswing_idx * 1}, Impact: ~{impact_idx * 1} (assuming sample_rate=1)")
    
    # Get rotation angle from the original video file
    rotation_angle = 0
    if os.path.exists(video_path):
        cap = cv2.VideoCapture(video_path)
        if cap.isOpened():
            try:
                orientation = int(cap.get(cv2.CAP_PROP_ORIENTATION_META))
                if orientation == 90:
                    rotation_angle = 270  # Rotate counterclockwise
                elif orientation == 180:
                    rotation_angle = 180
                elif orientation == 270:
                    rotation_angle = 90   # Rotate counterclockwise
                print(f"Video orientation metadata: {orientation}, applying rotation: {rotation_angle}")
            except Exception as e:
                print(f"Could not read orientation metadata: {e}")
            finally:
                cap.release()
    else:
        print(f"Warning: Video path {video_path} not found for rotation check.")

    phase_indices = {'setup': setup_idx, 'backswing': backswing_idx, 'impact': impact_idx}

    for phase_name, frame_idx in phase_indices.items():
        if 0 <= frame_idx < len(frames):
            frame = frames[frame_idx].copy()
            if rotation_angle != 0:
                frame = _apply_rotation(frame, rotation_angle)
            key_frames[phase_name] = frame
            print(f"Successfully extracted {phase_name} frame from memory.")
        else:
            print(f"Failed to extract {phase_name} frame: index {frame_idx} is out of bounds for {len(frames)} frames.")

    return key_frames


def _apply_rotation(frame, rotation_angle):
    """
    Apply rotation to a frame based on angle
    """
    if rotation_angle == 90:
        return cv2.rotate(frame, cv2.ROTATE_90_COUNTERCLOCKWISE)
    elif rotation_angle == 180:
        return cv2.rotate(frame, cv2.ROTATE_180)
    elif rotation_angle == 270:
        return cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)
    else:
        return frame


def generate_improvement_comments(phase):
    """
    Generate improvement comments for each swing phase in Professional/Comparison format
    
    Args:
        phase (str): The swing phase ('setup', 'backswing', 'impact')
        
    Returns:
        dict: Dictionary with 'pro_analysis' and 'comparison' keys
    """
    comments = {
        'setup': {
            'pro_analysis': [
                "Balanced stance with feet shoulder-width apart",
                "Even weight distribution on both feet",
                "Neutral grip with hands in proper position",
                "Athletic posture with slight forward bend",
                "Ball positioned correctly for club selection"
            ],
            'comparison': [
                "Compare your stance width to the pro's balanced setup",
                "Check if your weight is evenly distributed like the pro",
                "Ensure your grip matches the pro's neutral hand position",
                "Adjust your posture to match the pro's athletic stance",
                "Position the ball in your stance similar to the pro"
            ]
        },
        'backswing': {
            'pro_analysis': [
                "Full 90+ degree shoulder rotation",
                "Controlled hip turn with stable lower body",
                "Club on proper swing plane at top",
                "Consistent spine angle throughout",
                "Minimal weight shift to right side"
            ],
            'comparison': [
                "Increase your shoulder turn to match the pro's full rotation",
                "Control your hip movement like the pro's stable base",
                "Adjust your club position to match the pro's swing plane",
                "Maintain spine angle consistency like the professional",
                "Minimize weight shift compared to the pro's centered position"
            ]
        },
        'impact': {
            'pro_analysis': [
                "Weight shifted to front foot (70-80%)",
                "Hands ahead of ball at impact",
                "Square club face to target line",
                "Head behind ball with steady position",
                "Hips and shoulders aligned to target"
            ],
            'comparison': [
                "Shift more weight to your front foot like the pro",
                "Get your hands ahead of the ball like the professional",
                "Square your club face to match the pro's alignment",
                "Keep your head steady and behind the ball like the pro",
                "Align your body to the target like the professional"
            ]
        }
    }
    
    return comments.get(phase, {'pro_analysis': [], 'comparison': []})


def load_pro_reference_images(pro_images_dir="pro_reference"):
    """
    Load professional golfer reference images from directory
    
    Args:
        pro_images_dir (str): Directory containing professional reference images
        
    Returns:
        dict: Dictionary with phase names as keys and image arrays as values
    """
    # Get the absolute path to the pro_reference directory
    # This ensures it works regardless of the current working directory
    if not os.path.isabs(pro_images_dir):
        # Get the directory where this script is located
        script_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
        pro_images_dir = os.path.join(script_dir, pro_images_dir)
    
    pro_frames = {}
    
    # Expected filenames for the 3 phases
    phase_files = {
        'setup': 'setup.jpg',
        'backswing': 'backswing.jpg', 
        'impact': 'impact.jpg'
    }
    
    for phase, filename in phase_files.items():
        image_path = os.path.join(pro_images_dir, filename)
        if os.path.exists(image_path):
            image = cv2.imread(image_path)
            if image is not None:
                # Ensure the image is in color format
                image = ensure_color_frame(image)
                pro_frames[phase] = image
            else:
                # Create a placeholder if image can't be loaded
                pro_frames[phase] = np.zeros((480, 640, 3), dtype=np.uint8)
        else:
            # Create a placeholder if file doesn't exist
            pro_frames[phase] = np.zeros((480, 640, 3), dtype=np.uint8)
    
    return pro_frames


def save_frame_with_orientation(frame, output_path):
    """
    Save a frame using PIL after converting from BGR to RGB.
    Ensures proper color handling and orientation.
    
    Args:
        frame (numpy.ndarray): Frame in BGR format (OpenCV)
        output_path (str): Path to save the image
    """
    try:
        if frame is None or frame.size == 0:
            # Save a black image if frame is invalid
            black = np.zeros((480, 640, 3), dtype=np.uint8)
            img = Image.fromarray(black)
            img.save(output_path, format="JPEG", quality=95)
            return
        
        # Verify frame is in color (3 channels)
        if len(frame.shape) != 3 or frame.shape[2] != 3:
            raise ValueError(f"Frame is not in color format. Shape: {frame.shape}")
        
        # Convert BGR (OpenCV) to RGB (PIL)
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        # Create PIL image and save with high quality
        img = Image.fromarray(rgb_frame)
        img.save(output_path, format="JPEG", quality=95)
        
    except Exception as e:
        print(f"Warning: Error saving frame to {output_path}: {str(e)}")
        # Create a fallback black image
        try:
            black = np.zeros((480, 640, 3), dtype=np.uint8)
            img = Image.fromarray(black)
            img.save(output_path, format="JPEG", quality=95)
        except Exception as fallback_error:
            print(f"Error: Could not save fallback image: {str(fallback_error)}")
            raise


def create_key_frame_comparison(user_video_path, pro_video_path=None, user_swing_phases=None, pro_swing_phases=None, output_dir="downloads", use_pro_images=True):
    """
    Create separate images for 3 key frames from user and pro golfer swings
    
    IMPORTANT: This function preserves the original sizes of both user and professional frames.
    No resizing, rotation, or distortion is applied to either frame. Each frame is saved
    as a separate image file at its original resolution.
    
    Args:
        user_video_path (str): Path to the user's golf swing video
        pro_video_path (str): Path to the professional golfer's swing video (optional if use_pro_images=True)
        user_swing_phases (dict): Optional swing phase data for user video
        pro_swing_phases (dict): Optional swing phase data for pro video
        output_dir (str): Directory to save the separate images
        use_pro_images (bool): Whether to use provided pro reference images instead of video
        
    Returns:
        dict: Dictionary with phase names as keys and dictionaries containing 
              'user_image_path', 'pro_image_path', 'title', and 'comments' as values
    """
    # Extract key frames from user video
    user_frames = extract_key_swing_frames(user_video_path, user_frames, user_swing_phases)
    
    # Get pro frames either from provided images or video
    if use_pro_images:
        pro_frames = load_pro_reference_images()
    else:
        pro_frames = extract_key_swing_frames(pro_video_path, pro_frames, pro_swing_phases)
    
    # Create output directory with absolute path
    output_dir = os.path.abspath(output_dir)
    os.makedirs(output_dir, exist_ok=True)
    
    comparison_data = {}
    phases = ['setup', 'backswing', 'impact']
    phase_titles = ['Starting Position', 'Top of Backswing', 'Impact with Ball']
    
    for i, phase in enumerate(phases):
        # Get frames for this phase
        user_frame = user_frames.get(phase, np.zeros((480, 640, 3), dtype=np.uint8))
        pro_frame = pro_frames.get(phase, np.zeros((480, 640, 3), dtype=np.uint8))
        
        # CRITICAL: Keep user frame EXACTLY as extracted - no processing at all
        # Only ensure pro frame is in color format since it comes from reference images
        pro_frame = ensure_color_frame(pro_frame)
        
        # Save user frame with original size using PIL to ensure correct orientation and color
        video_name = os.path.splitext(os.path.basename(user_video_path))[0]
        user_output_path = os.path.join(output_dir, f"{video_name}_{phase}_user.jpg")
        pro_output_path = os.path.join(output_dir, f"{video_name}_{phase}_pro.jpg")
        
        # Save user image using PIL (handles BGR->RGB and orientation)
        try:
            save_frame_with_orientation(user_frame, user_output_path)
            user_success = True
        except Exception as e:
            print(f"Warning: Failed to save user image to {user_output_path}: {e}")
            user_success = False
        # Save pro image using OpenCV (as before)
        pro_success = cv2.imwrite(pro_output_path, pro_frame)
        
        if user_success:
            print(f"Successfully saved user image: {user_output_path}")
        if not user_success:
            print(f"Warning: Failed to save user image to {user_output_path}")
        if pro_success:
            print(f"Successfully saved pro image: {pro_output_path}")
        if not pro_success:
            print(f"Warning: Failed to save pro image to {pro_output_path}")
        
        # Get improvement comments
        comments = generate_improvement_comments(phase)
        
        comparison_data[phase] = {
            'user_image_path': user_output_path,
            'pro_image_path': pro_output_path,
            'title': phase_titles[i],
            'comments': comments
        }
    
    return comparison_data


def normalize_frames(frames, target_height=480):
    """
    Normalize frames to a consistent size while maintaining aspect ratio
    
    Args:
        frames (list): List of frames
        target_height (int): Target height for normalized frames
        
    Returns:
        list: List of normalized frames
    """
    normalized_frames = []
    
    for frame in frames:
        # Use the color-safe resize function
        resized = resize_frame_proportionally(frame, target_height)
        normalized_frames.append(resized)
        
    return normalized_frames


def create_side_by_side_comparison(user_frames, pro_frames, output_path, fps=30):
    """
    Create a side-by-side comparison video
    
    Args:
        user_frames (list): List of user swing frames
        pro_frames (list): List of pro swing frames
        output_path (str): Path to save the comparison video
        fps (int): Frames per second for output video
        
    Returns:
        str: Path to the comparison video
    """
    if not user_frames or not pro_frames:
        raise ValueError("Both user and pro frames must be provided")
        
    # Ensure all frames are in color format
    user_frames = [ensure_color_frame(frame) for frame in user_frames]
    pro_frames = [ensure_color_frame(frame) for frame in pro_frames]
    
    # Get dimensions from first frames
    user_h, user_w = user_frames[0].shape[:2]
    pro_h, pro_w = pro_frames[0].shape[:2]
    
    # Choose target height (smaller of the two, capped at 720p)
    target_height = min(user_h, pro_h, 720)
    
    # Resize both user and pro frames proportionally to the same height
    user_resized = []
    for frame in user_frames:
        resized = resize_frame_proportionally(frame, target_height)
        user_resized.append(resized)
    
    pro_resized = []
    for frame in pro_frames:
        resized = resize_frame_proportionally(frame, target_height)
        pro_resized.append(resized)
    
    # Ensure we have the same number of frames by duplicating the last frame if needed
    max_frames = max(len(user_resized), len(pro_resized))
    
    user_aligned = user_resized.copy()
    while len(user_aligned) < max_frames:
        user_aligned.append(user_aligned[-1])
        
    while len(pro_resized) < max_frames:
        pro_resized.append(pro_resized[-1])
    
    # Create output directory if it doesn't exist
    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
    
    # Get dimensions for the combined frame using original user frame dimensions
    pro_h, pro_w = pro_resized[0].shape[:2]
    
    # Create a combined frame with padding
    padding = 20  # Pixels between the two videos
    combined_width = user_w + pro_w + padding
    combined_height = max(user_h, pro_h)
    
    # Create video writer
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (combined_width, combined_height))
    
    if not out.isOpened():
        raise IOError(f"Failed to create video writer for {output_path}")
    
    # Create the combined video
    for i in tqdm(range(min(len(user_aligned), len(pro_resized))), desc="Creating comparison video"):
        # Create a blank canvas
        combined = np.ones((combined_height, combined_width, 3), dtype=np.uint8) * 255
        
        # Add title text
        font = cv2.FONT_HERSHEY_SIMPLEX
        cv2.putText(combined, "Your Swing", (user_w//2 - 60, 30), font, 1, (0, 0, 0), 2)
        cv2.putText(combined, "Pro Swing", (user_w + padding + pro_w//2 - 60, 30), font, 1, (0, 0, 0), 2)
        
        # Add frame number
        cv2.putText(combined, f"Frame: {i+1}/{min(len(user_aligned), len(pro_resized))}", 
                   (10, combined_height - 10), font, 0.5, (0, 0, 0), 1)
        
        # Paste user frame at original size and orientation
        y_offset_user = (combined_height - user_h) // 2
        combined[y_offset_user:y_offset_user + user_h, 0:user_w] = user_aligned[i]
        
        # Paste pro frame
        y_offset_pro = (combined_height - pro_h) // 2
        combined[y_offset_pro:y_offset_pro + pro_h, user_w + padding:user_w + padding + pro_w] = pro_resized[i]
        
        # Draw vertical line between frames
        cv2.line(combined, (user_w + padding//2, 0), (user_w + padding//2, combined_height), (0, 0, 0), 2)
        
        # Write to video
        out.write(combined)
    
    out.release()
    
    return output_path


def align_swings(user_frames, pro_frames, method="manual"):
    """
    Align user and pro swings based on swing phases
    
    Args:
        user_frames (list): List of user swing frames
        pro_frames (list): List of pro swing frames
        method (str): Alignment method ('manual' or 'auto')
        
    Returns:
        tuple: Aligned user frames and pro frames
    """
    # For now, we'll use a simple frame stretching approach
    # In the future, this could be enhanced with ML-based swing phase detection
    
    # Get frame counts
    user_count = len(user_frames)
    pro_count = len(pro_frames)
    
    # If almost equal, return as-is
    if abs(user_count - pro_count) <= 5:
        return user_frames, pro_frames
    
    # If user has more frames, subsample
    if user_count > pro_count:
        indices = np.linspace(0, user_count - 1, pro_count, dtype=int)
        return [user_frames[i] for i in indices], pro_frames
    
    # If pro has more frames, subsample
    indices = np.linspace(0, pro_count - 1, user_count, dtype=int)
    return user_frames, [pro_frames[i] for i in indices]


def create_frame_by_frame_comparison(user_video_path, pro_video_path, output_dir="downloads"):
    """
    Create a frame-by-frame comparison between user and pro golfer swings
    
    Args:
        user_video_path (str): Path to the user's golf swing video
        pro_video_path (str): Path to the professional golfer's swing video
        output_dir (str): Directory to save the comparison video
        
    Returns:
        str: Path to the comparison video
    """
    # Extract frames
    user_frames = extract_frames(user_video_path)
    pro_frames = extract_frames(pro_video_path)
    
    # Align swings
    aligned_user, aligned_pro = align_swings(user_frames, pro_frames)
    
    # Create output path
    video_name = os.path.splitext(os.path.basename(user_video_path))[0]
    output_path = os.path.join(output_dir, f"{video_name}_comparison.mp4")
    
    # Create side-by-side comparison
    return create_side_by_side_comparison(aligned_user, aligned_pro, output_path)