#!/usr/bin/env python3
"""
Test script to verify multi-panel layout splitting functionality
Tests the panel splitting with 6786505.jpg (horizontal strip layout)
"""

import sys
import os
import cv2
import numpy as np
from pathlib import Path
from panel_splitter import PanelSplitter

def test_multi_panel_splitting():
    """Test splitting the multi-panel layout image"""
    print("=" * 60)
    print("TESTING MULTI-PANEL LAYOUT SPLITTING")
    print("=" * 60)
    
    # Initialize splitter with debug mode
    splitter = PanelSplitter(debug=True)
    
    # Target layout file - this is a horizontal strip with many panels
    layout_path = Path("layouts") / "6786505.jpg"
    
    if not layout_path.exists():
        print(f"❌ ERROR: Layout file {layout_path} not found!")
        return False
    
    print(f"📁 Testing with: {layout_path.name}")
    
    # Load and examine the image
    image = cv2.imread(str(layout_path))
    if image is None:
        print(f"❌ ERROR: Could not load image {layout_path}")
        return False
    
    height, width = image.shape[:2]
    print(f"📐 Image dimensions: {width}x{height}")
    
    # Target panel count - let's first see how many OpenAI detects, then use that as target
    print("🤖 Getting OpenAI panel count first...")
    try:
        from openai_detector import OpenAIImageDetector
        temp_detector = OpenAIImageDetector()
        temp_detector.load_master_images()
        panel_result = temp_detector.count_panels_in_layout(str(layout_path))
        target_count = panel_result.get('panel_count', 10)
        print(f"🎯 OpenAI detected {target_count} panels - using this as target")
    except:
        target_count = 10  # fallback
        print(f"🎯 Using fallback target count: {target_count}")
    
    print("\n" + "─" * 40)
    print("RUNNING PANEL SPLITTING...")
    print("─" * 40)
    
    # Split the layout
    splits = splitter.split_panels(str(layout_path), target_count)
    
    print(f"\n📊 SPLITTING RESULTS:")
    print(f"Generated {len(splits)} splits (target: {target_count})")
    
    # Verify results
    success = len(splits) == target_count
    
    if success:
        print(f"✅ SUCCESS: Generated exactly {target_count} splits!")
    else:
        print(f"⚠️  WARNING: Generated {len(splits)} splits instead of {target_count}")
    
    print("\n📋 DETAILED SPLIT ANALYSIS:")
    print("─" * 40)
    
    total_area = 0
    original_area = width * height
    
    for i, split in enumerate(splits):
        x, y, w, h = split['bounds']
        area = w * h
        total_area += area
        
        print(f"Split {i+1:2d}: [{x:4d}, {y:4d}, {w:4d}, {h:4d}] "
              f"area={area:6d} conf={split['confidence']:.3f}")
        
        # Verify split bounds are reasonable
        if w < 20 or h < 20:
            print(f"    ⚠️  Split {i+1} is very small!")
        if x < 0 or y < 0 or x+w > width or y+h > height:
            print(f"    ❌ Split {i+1} bounds are out of image!")
    
    # Calculate coverage
    coverage = (total_area / original_area) * 100
    print(f"\n📈 Coverage: {coverage:.1f}% of original image")
    
    if coverage < 80:
        print("⚠️  Low coverage - some areas might be missed")
    elif coverage > 120:
        print("⚠️  High coverage - splits might be overlapping")
    
    # Check for overlaps
    print("\n🔍 CHECKING FOR OVERLAPS:")
    overlaps = 0
    for i, split1 in enumerate(splits):
        for j, split2 in enumerate(splits[i+1:], i+1):
            if rectangles_overlap(split1['bounds'], split2['bounds']):
                overlaps += 1
                print(f"    ⚠️  Splits {i+1} and {j+1} overlap!")
    
    if overlaps == 0:
        print("    ✅ No overlaps detected")
    else:
        print(f"    ❌ Found {overlaps} overlapping pairs")
    
    # Save individual split images for inspection
    print("\n💾 SAVING SPLIT IMAGES:")
    splits_dir = Path("test_splits")
    splits_dir.mkdir(exist_ok=True)
    
    for i, split in enumerate(splits):
        split_filename = splits_dir / f"6786505_split_{i+1:02d}.jpg"
        cv2.imwrite(str(split_filename), split['image'])
        print(f"    Saved: {split_filename}")
    
    print(f"\n📁 All split images saved to: {splits_dir}/")
    
    # Method analysis
    print("\n🔬 METHOD ANALYSIS:")
    method_votes = {}
    for split in splits:
        for method in split.get('method_votes', []):
            method_votes[method] = method_votes.get(method, 0) + 1
    
    for method, count in sorted(method_votes.items(), key=lambda x: x[1], reverse=True):
        print(f"    {method}: {count} votes")
    
    print("\n" + "=" * 60)
    if success:
        print(f"🎉 TEST PASSED: Successfully split {target_count}-panel layout!")
    else:
        print(f"❌ TEST FAILED: Did not produce exactly {target_count} splits")
    print("=" * 60)
    
    return success

def rectangles_overlap(rect1, rect2):
    """Check if two rectangles overlap"""
    x1, y1, w1, h1 = rect1
    x2, y2, w2, h2 = rect2
    
    return not (x1 + w1 <= x2 or x2 + w2 <= x1 or y1 + h1 <= y2 or y2 + h2 <= y1)

def test_with_openai_guidance():
    """Test splitting with OpenAI panel count guidance"""
    print("\n" + "=" * 60)
    print("TESTING WITH OPENAI PANEL COUNT GUIDANCE")
    print("=" * 60)
    
    try:
        from openai_detector import OpenAIImageDetector
        
        # Initialize OpenAI detector
        detector = OpenAIImageDetector(split_mode=True)
        detector.load_master_images()
        
        layout_path = str(Path("layouts") / "6786505.jpg")
        
        print("🤖 Getting OpenAI panel count...")
        panel_result = detector.count_panels_in_layout(layout_path)
        openai_count = panel_result.get('panel_count', 1)
        confidence = panel_result.get('confidence', 'unknown')
        
        print(f"OpenAI detected: {openai_count} panels (confidence: {confidence})")
        
        # Test full split_layout_and_match functionality
        print("\n🔄 Testing full split_layout_and_match...")
        master_ids = list(detector.master_images.keys())[:10]  # Test with first 10 masters
        
        result = detector.splitter.split_layout_and_match(layout_path, master_ids, detector)
        
        print(f"Split and match result:")
        print(f"  Splits generated: {result.get('splits_generated', 0)}")
        print(f"  Panel count: {result.get('panel_count', 'unknown')}")
        print(f"  Detected masters: {len(result.get('detected_masters', []))}")
        
        return openai_count >= 5  # Success if we detect at least 5 panels
        
    except Exception as e:
        print(f"❌ Error testing with OpenAI guidance: {e}")
        return False

def main():
    """Main test function"""
    print("🧪 STARTING MULTI-PANEL LAYOUT SPLITTING TEST")
    
    # Test basic splitting
    basic_success = test_multi_panel_splitting()
    
    # Test with OpenAI guidance
    openai_success = test_with_openai_guidance()
    
    print(f"\n📊 FINAL RESULTS:")
    print(f"Basic splitting: {'✅ PASSED' if basic_success else '❌ FAILED'}")
    print(f"OpenAI guidance: {'✅ PASSED' if openai_success else '❌ FAILED'}")
    
    if basic_success and openai_success:
        print("\n🎉 ALL TESTS PASSED! The multi-panel splitting is working correctly.")
        return 0
    else:
        print("\n❌ Some tests failed. Please check the output above.")
        return 1

if __name__ == "__main__":
    sys.exit(main())