import torch
import torch.nn as nn
from models.deeplob import DeepLOB
import os

print(f"Torch Version: {torch.__version__}")
try:
    import onnx
    print(f"ONNX Version: {onnx.__version__}")
except ImportError:
    print("ONNX not installed")

try:
    print("Initializing DeepLOB...")
    model = DeepLOB(y_len=3)
    model.eval()
    
    # Dummy input: (Batch=1, Channels=2, Time=100, Features=40)
    dummy = torch.randn(1, 2, 100, 40)
    
    output_path = "debug_deeplob.onnx"
    
    print("Exporting to ONNX...")
    torch.onnx.export(
        model, dummy, output_path,
        input_names=['input'], output_names=['output'],
        opset_version=12,
        do_constant_folding=True,
        verbose=True  # Print compilation steps
    )
    print(f"✅ Export Success: {output_path}")
    
except Exception as e:
    print(f"❌ Export Failed:")
    print(e)
    import traceback
    traceback.print_exc()