#!/usr/bin/env python3
"""
Test script to verify YOLOv5 postprocessing fixes

This script tests the improved YOLOv5 postprocessing configuration
to ensure positive probabilities and proper bounding box detection.
"""
import sys
import os
import numpy as np

# Add core functions to path
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.append(os.path.join(parent_dir, 'core', 'functions'))

def test_yolov5_postprocessor():
    """Test the improved YOLOv5 postprocessor with mock data"""
    from Multidongle import PostProcessorOptions, PostProcessType, PostProcessor
    
    print("=" * 60)
    print("Testing Improved YOLOv5 Postprocessor")
    print("=" * 60)
    
    # Create YOLOv5 postprocessor options
    options = PostProcessorOptions(
        postprocess_type=PostProcessType.YOLO_V5,
        threshold=0.3,
        class_names=["person", "bicycle", "car", "motorbike", "aeroplane", "bus"],
        nms_threshold=0.5,
        max_detections_per_class=50
    )
    
    postprocessor = PostProcessor(options)
    
    print(f"✓ Postprocessor created with type: {options.postprocess_type.value}")
    print(f"✓ Confidence threshold: {options.threshold}")
    print(f"✓ NMS threshold: {options.nms_threshold}")
    print(f"✓ Number of classes: {len(options.class_names)}")
    
    # Create mock YOLOv5 output data - format: [batch, detections, features]
    # Features: [x_center, y_center, width, height, objectness, class0_prob, class1_prob, ...]
    mock_output = create_mock_yolov5_output()
    
    # Test processing
    try:
        result = postprocessor.process([mock_output])
        
        print(f"\n📊 Processing Results:")
        print(f"   Result type: {type(result).__name__}")
        print(f"   Detected objects: {result.box_count}")
        print(f"   Available classes: {result.class_count}")
        
        if result.box_count > 0:
            print(f"\n📦 Detection Details:")
            for i, box in enumerate(result.box_list):
                print(f"   Detection {i+1}:")
                print(f"     Class: {box.class_name} (ID: {box.class_num})")
                print(f"     Confidence: {box.score:.3f}")
                print(f"     Bounding Box: ({box.x1}, {box.y1}) to ({box.x2}, {box.y2})")
                print(f"     Box Size: {box.x2 - box.x1} x {box.y2 - box.y1}")
        
        # Verify positive probabilities
        all_positive = all(box.score > 0 for box in result.box_list)
        print(f"\n✓ All probabilities positive: {all_positive}")
        
        # Verify reasonable coordinates
        valid_coords = all(
            box.x2 > box.x1 and box.y2 > box.y1 
            for box in result.box_list
        )
        print(f"✓ All bounding boxes valid: {valid_coords}")
        
        return result
        
    except Exception as e:
        print(f"❌ Postprocessing failed: {e}")
        import traceback
        traceback.print_exc()
        return None

def create_mock_yolov5_output():
    """Create mock YOLOv5 output data for testing"""
    # YOLOv5 output format: [batch_size, num_detections, num_features]
    # Features: [x_center, y_center, width, height, objectness, class_probs...]
    
    batch_size = 1
    num_detections = 25200  # Typical YOLOv5 output size
    num_classes = 80  # COCO classes
    num_features = 5 + num_classes  # coords + objectness + class probs
    
    # Create mock output
    mock_output = np.zeros((batch_size, num_detections, num_features), dtype=np.float32)
    
    # Add some realistic detections
    detections = [
        # Format: [x_center, y_center, width, height, objectness, class_id, class_prob]
        [320, 240, 100, 150, 0.8, 0, 0.9],   # person
        [500, 300, 80, 60, 0.7, 2, 0.85],    # car  
        [150, 100, 60, 120, 0.6, 1, 0.75],   # bicycle
    ]
    
    for i, detection in enumerate(detections):
        x_center, y_center, width, height, objectness, class_id, class_prob = detection
        
        # Set coordinates and objectness
        mock_output[0, i, 0] = x_center
        mock_output[0, i, 1] = y_center
        mock_output[0, i, 2] = width
        mock_output[0, i, 3] = height
        mock_output[0, i, 4] = objectness
        
        # Set class probabilities (one-hot style)
        mock_output[0, i, 5 + int(class_id)] = class_prob
    
    print(f"✓ Created mock YOLOv5 output: {mock_output.shape}")
    print(f"   Added {len(detections)} test detections")
    
    # Wrap in mock output object
    class MockOutput:
        def __init__(self, data):
            self.ndarray = data
    
    return MockOutput(mock_output)

def test_result_formatting():
    """Test the result formatting functions"""
    from Multidongle import ObjectDetectionResult, BoundingBox
    
    print(f"\n" + "=" * 60)
    print("Testing Result Formatting")
    print("=" * 60)
    
    # Create mock detection result
    boxes = [
        BoundingBox(x1=100, y1=200, x2=200, y2=350, score=0.85, class_num=0, class_name="person"),
        BoundingBox(x1=300, y1=150, x2=380, y2=210, score=0.75, class_num=2, class_name="car"),
        BoundingBox(x1=50, y1=100, x2=110, y2=220, score=0.65, class_num=1, class_name="bicycle"),
    ]
    
    result = ObjectDetectionResult(
        class_count=80,
        box_count=len(boxes),
        box_list=boxes
    )
    
    # Test the enhanced result string generation  
    from Multidongle import MultiDongle, PostProcessorOptions, PostProcessType
    # Create a minimal MultiDongle instance to access the method
    options = PostProcessorOptions(postprocess_type=PostProcessType.YOLO_V5)
    multidongle = MultiDongle(port_id=[1], postprocess_options=options)  # Dummy port
    result_string = multidongle._generate_result_string(result)
    
    print(f"📝 Generated result string: {result_string}")
    
    # Test individual object summaries
    print(f"\n📊 Object Summary:")
    object_counts = {}
    for box in boxes:
        if box.class_name in object_counts:
            object_counts[box.class_name] += 1
        else:
            object_counts[box.class_name] = 1
    
    for class_name, count in sorted(object_counts.items()):
        print(f"   {count} {class_name}{'s' if count > 1 else ''}")
    
    return result

def show_configuration_usage():
    """Show how to use the fixed configuration"""
    print(f"\n" + "=" * 60)
    print("Configuration Usage Instructions")
    print("=" * 60)
    
    print(f"\n🔧 Updated Configuration:")
    print(f"   1. Modified multi_series_example.mflow:")
    print(f"      - Set 'enable_postprocessing': true")
    print(f"      - Added ExactPostprocessNode with YOLOv5 settings")
    print(f"      - Connected Model → Postprocess → Output")
    
    print(f"\n⚙️ Postprocessing Settings:")
    print(f"   - postprocess_type: 'yolo_v5'")
    print(f"   - confidence_threshold: 0.3")
    print(f"   - nms_threshold: 0.5")
    print(f"   - class_names: Full COCO 80 classes")
    
    print(f"\n🎯 Expected Improvements:")
    print(f"   ✓ Positive probability values (0.0 to 1.0)")
    print(f"   ✓ Proper object detection with bounding boxes")
    print(f"   ✓ Correct class names (person, car, bicycle, etc.)")
    print(f"   ✓ Enhanced live view with corner markers")
    print(f"   ✓ Detailed terminal output with object counts")
    print(f"   ✓ Non-Maximum Suppression to reduce duplicates")
    
    print(f"\n📁 Files Modified:")
    print(f"   - core/functions/Multidongle.py (improved YOLO processing)")
    print(f"   - multi_series_example.mflow (added postprocess node)")
    print(f"   - Enhanced live view display and terminal output")

if __name__ == "__main__":
    print("YOLOv5 Postprocessing Fix Verification")
    print("=" * 60)
    
    try:
        # Test the postprocessor
        result = test_yolov5_postprocessor()
        
        if result:
            # Test result formatting
            test_result_formatting()
            
            # Show usage instructions
            show_configuration_usage()
            
            print(f"\n🎉 All tests passed! YOLOv5 postprocessing should now work correctly.")
            print(f"   Use the updated multi_series_example.mflow configuration.")
        else:
            print(f"\n❌ Tests failed. Please check the error messages above.")
        
    except Exception as e:
        print(f"\n❌ Test suite failed with error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)