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

def detect_green_areas(video_path, sample_frames=[30, 60, 90]):
    """
    Detect green chroma areas in a video file.
    Analyzes multiple frames to get more accurate detection.
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise Exception(f"Could not open video file: {video_path}")
    
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    duration = frame_count / fps if fps > 0 else 0
    
    print(f"  - Video info: {frame_count} frames, {fps:.2f} fps, {duration:.2f}s")
    
    all_green_areas = []
    
    # Analyze multiple frames for better accuracy
    for frame_idx in sample_frames:
        if frame_idx >= frame_count:
            continue
            
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
        ret, frame = cap.read()
        
        if not ret:
            continue
            
        # Convert to HSV for better color segmentation
        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        
        # Define green range (adjust if needed)
        lower_green = np.array([35, 40, 40])  # More inclusive green range
        upper_green = np.array([85, 255, 255])
        mask = cv2.inRange(hsv, lower_green, upper_green)
        
        # Apply morphological operations to clean up the mask
        kernel = np.ones((3, 3), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        
        # Find contours of green areas
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        frame_areas = []
        for cnt in contours:
            area = cv2.contourArea(cnt)
            
            # Filter by area 
            if area > 1000:  # Minimum size requirements
                # Get bounding rectangle
                x, y, w, h = cv2.boundingRect(cnt)
                
                # Get minimum area rectangle (rotated rectangle)
                rect = cv2.minAreaRect(cnt)
                (center_x, center_y), (rect_w, rect_h), angle = rect
                
                # Convert angle to a more standard format
                # OpenCV returns angle in range [-90, 0) for rectangles
                # Convert to [0, 180) range
                if angle < -45:
                    angle = 90 + angle
                
                # Get the four corner points of the rotated rectangle
                box = cv2.boxPoints(rect)
                box = np.int0(box)
                
                frame_areas.append({
                    'x': int(x), 
                    'y': int(y), 
                    'w': int(w), 
                    'h': int(h),
                    'area': int(area),
                    'frame': frame_idx,
                    # Add rotation information
                    'rotation': {
                        'angle': float(angle),
                        'center_x': float(center_x),
                        'center_y': float(center_y),
                        'rotated_width': float(rect_w),
                        'rotated_height': float(rect_h),
                        'corners': box.tolist()
                    }
                })
        
        all_green_areas.extend(frame_areas)
    
    cap.release()
    
    # Merge similar areas across frames (simple approach)
    if not all_green_areas:
        return []
    
    # For now, return the largest area found
    largest_area = max(all_green_areas, key=lambda x: x['area'])
    return [largest_area]

def get_video_files(directory="."):
    """Get all video files in the specified directory."""
    video_extensions = ['*.mp4', '*.avi', '*.mov', '*.mkv', '*.wmv', '*.flv', '*.webm']
    video_files = []
    
    for extension in video_extensions:
        video_files.extend(glob.glob(os.path.join(directory, extension)))
        video_files.extend(glob.glob(os.path.join(directory, extension.upper())))
    
    return video_files

def save_debug_frame(video_path, frame_idx=30):
    """Save a debug frame with green area overlay for manual inspection."""
    cap = cv2.VideoCapture(video_path)
    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
    ret, frame = cap.read()
    cap.release()
    
    if not ret:
        return
    
    # Apply green detection
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    lower_green = np.array([35, 40, 40])
    upper_green = np.array([85, 255, 255])
    mask = cv2.inRange(hsv, lower_green, upper_green)
    
    # Find contours and draw both bounding and rotated rectangles
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # Create overlay
    overlay = frame.copy()
    overlay[mask > 0] = [0, 255, 0]  # Green overlay
    result = cv2.addWeighted(frame, 0.7, overlay, 0.3, 0)
    
    # Draw rectangles for significant contours
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area > 1000:  # Same filter as main detection
            # Draw regular bounding rectangle in blue
            x, y, w, h = cv2.boundingRect(cnt)
            cv2.rectangle(result, (x, y), (x+w, y+h), (255, 0, 0), 2)
            
            # Draw rotated rectangle in red
            rect = cv2.minAreaRect(cnt)
            box = cv2.boxPoints(rect)
            box = np.int0(box)
            cv2.drawContours(result, [box], 0, (0, 0, 255), 2)
            
            # Add angle text
            (center_x, center_y), (rect_w, rect_h), angle = rect
            if angle < -45:
                angle = 90 + angle
            cv2.putText(result, f'{angle:.1f}°', 
                       (int(center_x), int(center_y)), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
    
    # Save debug image
    video_name = Path(video_path).stem
    debug_filename = f"debug_{video_name}_frame{frame_idx}.jpg"
    cv2.imwrite(debug_filename, result)
    print(f"  - Debug frame saved: {debug_filename} (Blue=bounding box, Red=rotated box)")

def main():
    print("🎬 Green Chroma Detection Script")
    print("=" * 40)
    
    # Get all video files in current directory
    video_files = get_video_files()
    
    if not video_files:
        print("❌ No video files found in current directory!")
        return
    
    print(f"📁 Found {len(video_files)} video file(s):")
    for video in video_files:
        print(f"  - {video}")
    
    print("\n🔍 Processing videos...")
    results = {}
    
    for video_path in video_files:
        try:
            video_name = os.path.basename(video_path)
            print(f"\n📹 Processing: {video_name}")
            
            # Detect green areas
            areas = detect_green_areas(video_path)
            results[video_name] = areas
            
            if areas:
                print(f"  ✅ Found {len(areas)} green area(s)")
                for i, area in enumerate(areas):
                    print(f"     Area {i+1}: x={area['x']}, y={area['y']}, w={area['w']}, h={area['h']}, size={area['area']}px")
            else:
                print(f"  ⚠️  No green areas detected")
            
            # Save debug frame for manual inspection
            save_debug_frame(video_path)
            
        except Exception as e:
            print(f"  ❌ Error processing {video_path}: {str(e)}")
            results[video_name] = []
    
    # Save results
    output_file = 'green_areas.json'
    with open(output_file, 'w') as f:
        json.dump(results, f, indent=2)
    
    print(f"\n✅ Detection complete!")
    print(f"📄 Results saved to: {output_file}")
    print(f"🖼️  Debug frames saved for manual inspection")
    
    # Summary
    total_areas = sum(len(areas) for areas in results.values())
    print(f"\n📊 Summary: {len(results)} videos processed, {total_areas} green areas detected")

if __name__ == "__main__":
    main()
