#!/usr/bin/env python3
"""
Export YOLO model variants with different NMS settings for shiny icon debugging
"""

from ultralytics import YOLO
import os

def export_model_variants():
    model_path = "./raw_models/best.pt"
    output_dir = "./raw_models/exports"
    
    # Create output directory
    os.makedirs(output_dir, exist_ok=True)
    
    print(f"Loading model from: {model_path}")
    model = YOLO(model_path)
    
    # Export configurations to test
    configs = [
        {
            "name": "no_nms",
            "nms": False,
            "simplify": True,
            "description": "Raw model output without NMS - for debugging shiny detection"
        },
        {
            "name": "nms_relaxed",
            "nms": True,
            "max_det": 500,  # Increase from default 300
            "conf": 0.1,     # Lower confidence threshold
            "simplify": True,
            "description": "NMS with more detections and lower confidence"
        },
        {
            "name": "nms_very_relaxed", 
            "nms": True,
            "max_det": 1000,  # Even more detections
            "conf": 0.05,     # Very low confidence
            "simplify": True,
            "description": "NMS with maximum detections for rare classes"
        }
    ]
    
    for config in configs:
        try:
            print(f"\n🚀 Exporting {config['name']}: {config['description']}")
            
            # Extract export parameters
            export_params = {k: v for k, v in config.items() 
                           if k not in ['name', 'description']}
            
            # Export model
            exported_path = model.export(
                format='onnx',
                **export_params
            )
            
            # Move to organized location
            output_file = os.path.join(output_dir, f"best_{config['name']}.onnx")
            if os.path.exists(exported_path):
                os.rename(exported_path, output_file)
                print(f"✅ Exported: {output_file}")
            else:
                print(f"❌ Export failed for {config['name']}")
                
        except Exception as e:
            print(f"❌ Error exporting {config['name']}: {e}")
    
    print(f"\n📁 All exports saved to: {output_dir}")
    print("\n📋 Summary:")
    print("- best_no_nms.onnx: Raw 8400x99 output for debugging")
    print("- best_nms_relaxed.onnx: NMS with 500 max detections")  
    print("- best_nms_very_relaxed.onnx: NMS with 1000 max detections")
    print("\nNext: Copy desired model to app/src/main/assets/ as best.onnx")

if __name__ == "__main__":
    export_model_variants()