#!/usr/bin/env python3
"""
Script to fix YOLOv5 postprocessing configuration issues

This script demonstrates how to properly configure YOLOv5 postprocessing
to resolve negative probability values and incorrect result formatting.
"""
import sys
import os

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

def create_yolov5_postprocessor_options():
    """Create properly configured PostProcessorOptions for YOLOv5"""
    from Multidongle import PostProcessType, PostProcessorOptions
    
    # COCO dataset class names (80 classes for YOLOv5)
    yolo_class_names = [
        "person", "bicycle", "car", "motorbike", "aeroplane", "bus", "train", "truck", "boat", 
        "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", 
        "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", 
        "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", 
        "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", 
        "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", 
        "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", 
        "sofa", "pottedplant", "bed", "diningtable", "toilet", "tvmonitor", "laptop", "mouse", 
        "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", 
        "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"
    ]
    
    # Create YOLOv5 postprocessor options
    options = PostProcessorOptions(
        postprocess_type=PostProcessType.YOLO_V5,
        threshold=0.3,                      # Confidence threshold (0.3 is good for detection)
        class_names=yolo_class_names,       # All 80 COCO classes
        nms_threshold=0.5,                  # Non-Maximum Suppression threshold
        max_detections_per_class=50         # Maximum detections per class
    )
    
    return options

def create_fire_detection_postprocessor_options():
    """Create properly configured PostProcessorOptions for Fire Detection"""
    from Multidongle import PostProcessType, PostProcessorOptions
    
    options = PostProcessorOptions(
        postprocess_type=PostProcessType.FIRE_DETECTION,
        threshold=0.5,                      # Fire detection threshold
        class_names=["No Fire", "Fire"]     # Binary classification
    )
    
    return options

def test_postprocessor_options():
    """Test both postprocessor configurations"""
    
    print("=" * 60)
    print("Testing PostProcessorOptions Configuration")
    print("=" * 60)
    
    # Test YOLOv5 configuration
    print("\n1. YOLOv5 Configuration:")
    try:
        yolo_options = create_yolov5_postprocessor_options()
        print(f"   ✓ Postprocess Type: {yolo_options.postprocess_type.value}")
        print(f"   ✓ Confidence Threshold: {yolo_options.threshold}")
        print(f"   ✓ NMS Threshold: {yolo_options.nms_threshold}")
        print(f"   ✓ Max Detections: {yolo_options.max_detections_per_class}")
        print(f"   ✓ Number of Classes: {len(yolo_options.class_names)}")
        print(f"   ✓ Sample Classes: {yolo_options.class_names[:5]}...")
    except Exception as e:
        print(f"   ✗ YOLOv5 configuration failed: {e}")
    
    # Test Fire Detection configuration  
    print("\n2. Fire Detection Configuration:")
    try:
        fire_options = create_fire_detection_postprocessor_options()
        print(f"   ✓ Postprocess Type: {fire_options.postprocess_type.value}")
        print(f"   ✓ Confidence Threshold: {fire_options.threshold}")
        print(f"   ✓ Class Names: {fire_options.class_names}")
    except Exception as e:
        print(f"   ✗ Fire Detection configuration failed: {e}")

def demonstrate_multidongle_creation():
    """Demonstrate creating MultiDongle with correct postprocessing"""
    from Multidongle import MultiDongle
    
    print("\n" + "=" * 60)
    print("Creating MultiDongle with YOLOv5 Postprocessing")
    print("=" * 60)
    
    # Create YOLOv5 postprocessor options
    yolo_options = create_yolov5_postprocessor_options()
    
    # Example configuration (adjust paths to match your setup)
    PORT_IDS = [28, 32]  # Your dongle port IDs
    MODEL_PATH = "path/to/yolov5_model.nef"  # Your YOLOv5 model path
    
    print(f"Configuration:")
    print(f"   Port IDs: {PORT_IDS}")
    print(f"   Model Path: {MODEL_PATH}")
    print(f"   Postprocess Type: {yolo_options.postprocess_type.value}")
    print(f"   Confidence Threshold: {yolo_options.threshold}")
    
    # NOTE: Uncomment below to actually create MultiDongle instance
    # (requires actual dongle hardware and valid paths)
    """
    try:
        multidongle = MultiDongle(
            port_id=PORT_IDS,
            model_path=MODEL_PATH,
            auto_detect=True,
            postprocess_options=yolo_options  # This is the key fix!
        )
        
        print("   ✓ MultiDongle created successfully with YOLOv5 postprocessing")
        print("   ✓ This should resolve negative probability issues")
        
        # Initialize and start
        multidongle.initialize()
        multidongle.start()
        
        print("   ✓ MultiDongle initialized and started")
        
        # Don't forget to stop when done
        multidongle.stop()
        
    except Exception as e:
        print(f"   ✗ MultiDongle creation failed: {e}")
    """
    
    print(f"\n   📝 To fix your current issue:")
    print(f"      1. Change postprocess_type from 'fire_detection' to 'yolo_v5'")
    print(f"      2. Set proper class names (80 COCO classes)")
    print(f"      3. Adjust confidence threshold to 0.3 (instead of 0.5)")
    print(f"      4. Set NMS threshold to 0.5")

def show_configuration_summary():
    """Show summary of configuration changes needed"""
    
    print("\n" + "=" * 60)
    print("CONFIGURATION FIX SUMMARY")
    print("=" * 60)
    
    print("\n🔧 Current Issue:")
    print("   - YOLOv5 model with FIRE_DETECTION postprocessing")
    print("   - Results in negative probabilities like -0.39")
    print("   - Incorrect result formatting")
    
    print("\n✅ Solution:")
    print("   1. Use PostProcessType.YOLO_V5 instead of FIRE_DETECTION")
    print("   2. Set confidence threshold to 0.3 (good for object detection)")
    print("   3. Use 80 COCO class names for YOLOv5")
    print("   4. Set NMS threshold to 0.5 for proper object filtering")
    
    print("\n📁 File Changes Needed:")
    print("   - multi_series_example.mflow: Add ExactPostprocessNode")
    print("   - Set 'enable_postprocessing': true in model node")
    print("   - Configure postprocess_type: 'yolo_v5'")
    
    print("\n🚀 Expected Result After Fix:")
    print("   - Positive probabilities (0.0 to 1.0)")
    print("   - Object detection results with bounding boxes")
    print("   - Proper class names like 'person', 'car', etc.")
    print("   - Multiple objects detected per frame")

if __name__ == "__main__":
    print("YOLOv5 Postprocessing Fix Utility")
    print("=" * 60)
    
    try:
        test_postprocessor_options()
        demonstrate_multidongle_creation()
        show_configuration_summary()
        
        print("\n🎉 Configuration examples completed successfully!")
        print("   Use the fixed .mflow file or update your configuration.")
        
    except Exception as e:
        print(f"\n❌ Script failed with error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)