import os
import time
import argparse
import torch
import torchvision
import onnx
import onnxruntime as ort
import numpy as np

def export_to_onnx(model_path, onnx_path):
    print(f"Loading PyTorch checkpoint from {model_path}...")
    from data_generator import CLASSES
    num_classes = len(CLASSES) + 1
    
    # Recreate the model structure
    model = torchvision.models.detection.ssdlite320_mobilenet_v3_large(
        num_classes=num_classes,
        weights_backbone=torchvision.models.MobileNet_V3_Large_Weights.DEFAULT
    )
    
    # Load state dict
    state_dict = torch.load(model_path, map_location='cpu')
    model.load_state_dict(state_dict)
    model.eval()
    
    # Dummy input: list of 1 tensor of size [3, 320, 320]
    # SSDLite in torchvision expects a list of tensors during export
    dummy_input = [torch.randn(3, 320, 320)]
    
    print("Exporting model to ONNX...")
    os.makedirs(os.path.dirname(onnx_path), exist_ok=True)
    
    torch.onnx.export(
        model,
        (dummy_input,),
        onnx_path,
        opset_version=12,
        input_names=["images"],
        output_names=["boxes", "scores", "labels"],
        dynamic_axes={
            "images": {0: "batch_size"}
        },
        dynamo=False
    )
    
    print(f"Successfully saved ONNX model to {onnx_path}")
    
    # Verify ONNX model
    print("Verifying ONNX model graph structure...")
    onnx_model = onnx.load(onnx_path)
    onnx.checker.check_model(onnx_model)
    
    print("Inputs:")
    for input_node in onnx_model.graph.input:
        print(f"  - Name: {input_node.name}")
    print("Outputs:")
    for output_node in onnx_model.graph.output:
        print(f"  - Name: {output_node.name}")

def benchmark_onnx(onnx_path):
    print(f"Benchmarking ONNX inference speed using ONNX Runtime...")
    
    # Create inference session
    providers = ['CPUExecutionProvider']
    if 'DirectMLExecutionProvider' in ort.get_available_providers():
        providers = ['DirectMLExecutionProvider', 'CPUExecutionProvider']
        print("Using DirectML hardware acceleration!")
        
    session = ort.InferenceSession(onnx_path, providers=providers)
    
    # Prepare dummy input matching the shape of the exported model input
    # Note: torchvision SSD ONNX export expects list of images.
    # The input node shape in ONNX has dynamic axes. Let's create a float32 array
    # corresponding to the image. SSDLite internally resizes and batches.
    # Let's see the input signature. The input 'images' expects a 3D or 4D float tensor.
    # In our export script, the dummy input was [torch.randn(3, 320, 320)]
    # In ONNX, a list of Tensors gets mapped to a single concatenated batch tensor or similar.
    # Let's generate a random array of shape (1, 3, 320, 320)
    input_name = session.get_inputs()[0].name
    dummy_in = np.random.randn(3, 320, 320).astype(np.float32)
    
    # Warmup
    for _ in range(5):
        _ = session.run(None, {input_name: dummy_in})
        
    # Measure latency
    num_runs = 50
    start_time = time.time()
    for _ in range(num_runs):
        _ = session.run(None, {input_name: dummy_in})
    elapsed = time.time() - start_time
    avg_latency_ms = (elapsed / num_runs) * 1000
    
    print(f"Average ONNX Runtime inference latency: {avg_latency_ms:.2f} ms")
    if avg_latency_ms < 50.0:
        print("Success: Inference latency is under the 50 ms limit!")
    else:
        print("Warning: Inference latency is above the 50 ms limit. Optimize hardware providers.")
        
    return avg_latency_ms

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Export PyTorch SSDLite model to ONNX")
    parser.add_argument("--pytorch-model", type=str, default=os.path.join("models", "model.pt"), help="Path to PyTorch model weights")
    parser.add_argument("--onnx-model", type=str, default=os.path.join("models", "model.onnx"), help="Path to save ONNX model")
    
    args = parser.parse_args()
    
    try:
        export_to_onnx(args.pytorch_model, args.onnx_model)
        benchmark_onnx(args.onnx_model)
    except Exception as e:
        print(f"Error during export: {e}")
        import sys
        sys.exit(1)