"""导出 DLinear 模型为 ONNX 格式。

用法:
    python scripts/export_onnx.py

输出:
    models/checkpoints/dlinear_classification_2026/model.onnx
"""

import sys
from pathlib import Path

# 添加项目根目录到路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from models.dlinear.model import DLinear


def main():
    # 模型路径
    checkpoint_path = project_root / "models/checkpoints/dlinear_classification_2026/best_model.pt"
    output_path = checkpoint_path.parent / "model.onnx"

    print(f"加载模型: {checkpoint_path}")
    model = DLinear.load(str(checkpoint_path), device='cpu')

    print(f"模型配置: lookback={model.config.lookback}, input_dim={model.config.input_dim}")
    print(f"参数量: {model.count_parameters():,}")

    # 导出为 ONNX
    model.to_onnx(str(output_path))
    print(f"✓ ONNX 模型已保存到: {output_path}")


if __name__ == "__main__":
    main()