#!/usr/bin/env python3
"""
ONNX 导出脚本
从 PyTorch 训练好的模型导出为 ONNX 格式
兼容 torch 2.x，使用传统 TorchScript 导出器 (dynamo=False)
"""
import argparse
import os
import torch
import torch.nn as nn


class CNNGRU(nn.Module):
    def __init__(self, n_mels=64, num_classes=2):
        super().__init__()
        self.cnn = nn.Sequential(
            nn.Conv1d(n_mels, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool1d(2),
            nn.Conv1d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool1d(2),
            nn.Conv1d(128, 256, kernel_size=3, padding=1),
            nn.ReLU()
        )
        self.gru = nn.GRU(256, 128, batch_first=True, bidirectional=True)
        self.fc = nn.Sequential(
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64, num_classes)
        )

    def forward(self, x):
        # x: [B, T, n_mels]
        x = x.transpose(1, 2)
        x = self.cnn(x)
        x = x.transpose(1, 2)
        x, _ = self.gru(x)
        x = x.mean(dim=1)
        return self.fc(x)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--checkpoint', required=True, help='Path to .pth checkpoint')
    parser.add_argument('--output', default='models/gunshot_classifier.onnx')
    parser.add_argument('--n_mels', type=int, default=64)
    parser.add_argument('--num_classes', type=int, default=2)
    parser.add_argument('--opset', type=int, default=13)
    args = parser.parse_args()

    device = torch.device('cpu')
    model = CNNGRU(n_mels=args.n_mels, num_classes=args.num_classes).to(device)
    ckpt = torch.load(args.checkpoint, map_location=device)
    model.load_state_dict(ckpt['model'])
    model.eval()

    dummy_input = torch.randn(1, 63, args.n_mels, device=device)

    os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True)

    torch.onnx.export(
        model,
        dummy_input,
        args.output,
        export_params=True,
        opset_version=args.opset,
        do_constant_folding=True,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={
            'input': {0: 'batch_size'},
            'output': {0: 'batch_size'}
        },
        dynamo=False  # Use traditional exporter for GRU compatibility
    )
    print(f"ONNX model exported to {args.output}")
    print(f"Input shape: {dummy_input.shape} -> [batch, time, n_mels]")


if __name__ == '__main__':
    main()
