"""导出 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.models.tinytimemixer import TinyTimeMixerForPrediction

class TTMWrapper(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
    
    def forward(self, x):
        output = self.model(x)
        return output.prediction_outputs

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.load_state_dict(checkpoint['model_state_dict'])
        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=14,
            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 = TinyTimeMixerForPrediction.from_pretrained(
            model_name,
            trust_remote_code=True
        )
        
        # Wrap model
        model = TTMWrapper(base_model)
        model.eval()
        
        # TTM input: (batch_size, context_length, num_features)
        # Context length 512, features 1
        dummy_input = torch.randn(1, 512, 1)
        
        print(f"导出 ONNX...")
        torch.onnx.export(
            model,
            dummy_input,
            str(output_path),
            export_params=True,
            opset_version=14,
            do_constant_folding=True,
            input_names=['past_values'],
            output_names=['prediction_outputs'],
            dynamic_axes={'past_values': {0: 'batch_size'}, 'prediction_outputs': {0: 'batch_size'}}
        )
        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()