"""导出 DLinear, N-BEATS 和 TTM-R2 模型为 ONNX 格式。"""

import sys
import os
import torch
import numpy as np
from pathlib import Path

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

from models.dlinear.model import DLinear
from models.nbeats.model import NBeats
from tsfm_public.toolkit.get_model import get_model


def export_dlinear():
    print("\n" + "="*50)
    print("导出 DLinear 模型...")
    try:
        checkpoint_path = project_root / "models/checkpoints/dlinear_classification_2026/best_model.pt"
        output_path = checkpoint_path.parent / "model.onnx"
        
        if not checkpoint_path.exists():
            print(f"跳过: 找不到检查点 {checkpoint_path}")
            return

        print(f"加载模型: {checkpoint_path}")
        model = DLinear.load(str(checkpoint_path), device='cpu')
        
        print(f"模型配置: lookback={model.config.lookback}, input_dim={model.config.input_dim}")
        
        model.to_onnx(str(output_path))
        print(f"✓ ONNX 模型已保存到: {output_path}")
    except Exception as e:
        print(f"导出 DLinear 失败: {e}")

def export_nbeats():
    print("\n" + "="*50)
    print("导出 N-BEATS 模型...")
    try:
        checkpoint_path = project_root / "models/checkpoints/nbeats_btc_1h/best_model.pt"
        output_path = checkpoint_path.parent / "model.onnx"
        
        if not checkpoint_path.exists():
            print(f"跳过: 找不到检查点 {checkpoint_path}")
            return

        print(f"加载模型: {checkpoint_path}")
        checkpoint = torch.load(str(checkpoint_path), map_location='cpu')
        model = NBeats(checkpoint['config'])
        model.eval()
        
        # N-BEATS input: (batch_size, input_size)
        # input_size default is 168
        input_size = checkpoint['config'].input_size
        dummy_input = torch.randn(1, input_size)
        
        print(f"导出 ONNX (input_size={input_size})...")
        torch.onnx.export(
            model,
            dummy_input,
            str(output_path),
            export_params=True,
            opset_version=17,
            do_constant_folding=True,
            input_names=['input'],
            output_names=['forecast'],
            dynamic_axes={'input': {0: 'batch_size'}, 'forecast': {0: 'batch_size'}}
        )
        print(f"✓ ONNX 模型已保存到: {output_path}")
    except Exception as e:
        print(f"导出 N-BEATS 失败: {e}")
        import traceback
        traceback.print_exc()

def export_ttm_r2():
    print("\n" + "="*50)
    print("导出 TTM-R2 模型...")
    try:
        model_name = "ibm-granite/granite-timeseries-ttm-r2"
        output_dir = project_root / "models/checkpoints/ttm_r2_onnx"
        output_path = output_dir / "model.onnx"
        
        output_dir.mkdir(parents=True, exist_ok=True)
        
        print(f"加载模型: {model_name}")
        base_model = get_model(
            model_path=model_name,
            context_length=512,
            prediction_length=96,
        )
        base_model.eval()
        
        print(f"TTM Config: {base_model.config}")
        
        # TTM input: (batch_size, context_length, num_features)
        # Context length 512, features 1
        dummy_input = torch.randn(1, 512, 1)
        with torch.no_grad():
            output = base_model(dummy_input)
            if hasattr(output, 'prediction_outputs'):
                test_output = output.prediction_outputs
            else:
                test_output = output
            print(f"✓ 输出形状: {test_output.shape}")
        print(f"导出 ONNX (固定 Batch Size=1)...")
        torch.onnx.export(
            base_model,
            dummy_input,
            str(output_path),
            export_params=True,
            opset_version=17,  # Revert to 14 as it is standard
            do_constant_folding=True,
            input_names=['past_values'],
            output_names=['predictions'],
            # Remove dynamic axes to fix potential shape inference issues
            dynamic_axes={'past_values': {0: 'batch_size'}, 'predictions': {0: 'batch_size'}},
            verbose=False
        )
        print(f"✓ ONNX 模型已保存到: {output_path}")
    except Exception as e:
        print(f"导出 TTM-R2 失败: {e}")
        import traceback
        traceback.print_exc()

def main():
    export_dlinear()
    export_nbeats()
    #export_ttm_r2()

if __name__ == "__main__":
    main()