crpto_predictor/models/inference.py

"""
模型推理服务 - 统一管理 DLinear, N-BEATS, TTM-R2 模型加载与预测
"""
import os
import sys
import pickle
import torch
import numpy as np
import pandas as pd
import pandas_ta as ta
from typing import Dict, List, Optional, Tuple, Union, Any
from sklearn.preprocessing import StandardScaler

# 添加项目根目录到 sys.path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from models.dlinear.model import DLinear
from models.nbeats.model import NBeats
try:
    from tsfm_public.models.tinytimemixer import TinyTimeMixerForPrediction
except (ImportError, NameError):
    # 尝试作为可选依赖处理
    TinyTimeMixerForPrediction = None



class CustomUnpickler(pickle.Unpickler):
    """自定义反序列化器，用于处理 __main__ 模块的类重定向"""
    def find_class(self, module, name):
        # 将 __main__ 中的类重定向到当前模块
        if module == "__main__":
            if name == "ClassificationPreprocessor":
                return ClassificationPreprocessor
            if name == "NBeatsPreprocessor":
                return NBeatsPreprocessor
        return super().find_class(module, name)


class ClassificationPreprocessor:
    """DLinear分类预处理器"""
    
    def __init__(self, lookback: int = 168, forecast_horizon: int = 1):
        self.lookback = lookback
        self.forecast_horizon = forecast_horizon
        self.scaler = StandardScaler()
        self.is_fitted = False
        self.feature_columns = ['open', 'high', 'low', 'close', 'volume', 
                               'rsi', 'macd', 'macd_hist', 'macd_signal']

    def add_features(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()
        df['rsi'] = ta.rsi(df['close'], length=14)
        macd = ta.macd(df['close'])
        if macd is not None and not macd.empty:
            df['macd'] = macd.iloc[:, 0]
            df['macd_hist'] = macd.iloc[:, 1]
            df['macd_signal'] = macd.iloc[:, 2]
        return df

    def create_labels(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()
        df['future_price'] = df['close'].shift(-self.forecast_horizon)
        df['future_return'] = (df['future_price'] - df['close']) / df['close']
        df['label'] = (df['future_return'] > 0).astype(int)
        return df[:-self.forecast_horizon]

    def preprocess(self, df: pd.DataFrame, fit: bool = True):
        df = self.add_features(df)
        df = self.create_labels(df)
        df = df.dropna().reset_index(drop=True)
        
        if fit:
            scaled_features = self.scaler.fit_transform(df[self.feature_columns])
            self.is_fitted = True
        else:
            scaled_features = self.scaler.transform(df[self.feature_columns])
        
        X, y = [], []
        for i in range(len(scaled_features) - self.lookback):
            X.append(scaled_features[i:i + self.lookback])
            y.append(df['label'].iloc[i + self.lookback])
        
        return np.array(X), np.array(y), df

    def preprocess_for_inference(self, df: pd.DataFrame):
        """推理专用预处理：整个df一次性处理，不生成标签"""
        df = self.add_features(df)
        df = df.dropna().reset_index(drop=True)
        
        # 确保有足够的数据
        if len(df) == 0:
             return np.array([]), df, []

        scaled_features = self.scaler.transform(df[self.feature_columns])
        
        X = []
        valid_indices = []
        
        if len(scaled_features) < self.lookback:
            return np.array([]), df, []
            
        # 注意：这里生成的窗口数量应该是 len(scaled) - lookback + 1
        # 例如 len=168, lookback=168 -> range(1) -> i=0 -> [0:168] (1个窗口)
        for i in range(len(scaled_features) - self.lookback + 1):
            window = scaled_features[i:i + self.lookback]
            X.append(window)
            valid_indices.append(i + self.lookback - 1)
            
        return np.array(X), df, valid_indices


class NBeatsPreprocessor:
    """N-BEATS预处理器"""
    def __init__(self, lookback: int = 168, forecast_horizon: int = 1):
        self.lookback = lookback
        self.forecast_horizon = forecast_horizon
        self.price_scaler = StandardScaler()
        self.is_fitted = False

    def preprocess(self, df: pd.DataFrame, fit: bool = True):
        df = df.copy()
        prices = df['close'].values.reshape(-1, 1)

        if fit:
            scaled_prices = self.price_scaler.fit_transform(prices)
            self.is_fitted = True
        else:
            scaled_prices = self.price_scaler.transform(prices)

        X = []
        valid_indices = []

        if len(scaled_prices) < self.lookback:
            return np.array([]), df, []

        for i in range(len(scaled_prices) - self.lookback + 1):
            X.append(scaled_prices[i:i + self.lookback, 0])
            valid_indices.append(i + self.lookback - 1)

        return np.array(X), df, valid_indices


class InferenceManager:
    """统一模型推理管理器"""
    
    _instance = None
    
    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
            cls._instance = super(InferenceManager, cls).__new__(cls)
            cls._instance.initialized = False
        return cls._instance

    def __init__(self, device_str: str = None):
        if self.initialized:
            return
            
        import platform
        if device_str:
            self.device = torch.device(device_str)
        elif platform.system() == 'Darwin' and torch.backends.mps.is_available():
            self.device = torch.device('mps')
        else:
            self.device = torch.device('cpu')
            
        print(f"InferenceManager 使用设备: {self.device}")
        
        self.models = {}
        self.preprocessors = {}
        
        # 配置路径
        self.configs = {
            'dlinear': {
                'model_path': './models/checkpoints/dlinear_classification_2026/best_model.pt',
                'preprocessor_path': './models/checkpoints/dlinear_classification_2026/preprocessor.pkl'
            },
            'nbeats': {
                'model_path': './models/checkpoints/nbeats_btc_1h/best_model.pt',
                'preprocessor_path': './models/checkpoints/nbeats_btc_1h/preprocessor.pkl'
            },
            'ttm': {
                'model_name': 'ibm-granite/granite-timeseries-ttm-r2'
            }
        }
        self.initialized = True

    def load_model(self, model_name: str):
        """加载指定模型"""
        if model_name in self.models:
            return

        print(f"正在加载 {model_name}...")
        try:
            if model_name == 'dlinear':
                self._load_dlinear()
            elif model_name == 'nbeats':
                self._load_nbeats()
            elif model_name == 'ttm':
                self._load_ttm()
            else:
                print(f"未知模型: {model_name}")
        except Exception as e:
            print(f"加载 {model_name} 失败: {e}")
            import traceback
            traceback.print_exc()

    def load_all_models(self):
        """加载所有模型"""
        self.load_model('dlinear')
        self.load_model('nbeats')
        self.load_model('ttm')

    def _load_dlinear(self):
        config = self.configs['dlinear']
        if not os.path.exists(config['model_path']):
            print(f"DLinear 模型文件不存在: {config['model_path']}")
            return

        checkpoint = torch.load(config['model_path'], map_location=self.device, weights_only=False)
        model = DLinear(checkpoint['config'])
        model.load_state_dict(checkpoint['model_state_dict'])
        self.models['dlinear'] = model.to(self.device).eval()
        
        with open(config['preprocessor_path'], 'rb') as f:
            try:
                self.preprocessors['dlinear'] = pickle.load(f)
            except AttributeError:
                f.seek(0)
                self.preprocessors['dlinear'] = CustomUnpickler(f).load()
        print("  ✓ DLinear 已加载")

    def _load_nbeats(self):
        config = self.configs['nbeats']
        if not os.path.exists(config['model_path']):
            print(f"N-BEATS 模型文件不存在: {config['model_path']}")
            return

        checkpoint = torch.load(config['model_path'], map_location=self.device, weights_only=False)
        model = NBeats(checkpoint['config'])
        self.models['nbeats'] = model.to(self.device).eval()
        with open(config['preprocessor_path'], 'rb') as f:
             try:
                self.preprocessors['nbeats'] = pickle.load(f)
             except AttributeError:
                f.seek(0)
                self.preprocessors['nbeats'] = CustomUnpickler(f).load()
        print("  ✓ N-BEATS 已加载")

    def _load_ttm(self):
        # 检查 TTM 类是否已导入
        if TinyTimeMixerForPrediction is None:
            print("TTM-R2 模型类未定义 (可能是依赖库缺失)")
            return

        config = self.configs['ttm']
        try:
            model = TinyTimeMixerForPrediction.from_pretrained(
                config['model_name'],
                trust_remote_code=True
            )
            self.models['ttm'] = model.to(self.device).eval()
            print("  ✓ TTM-R2 已加载")
        except Exception as e:
            print(f"TTM 加载失败: {e}")

    def predict_dlinear(self, df: pd.DataFrame) -> Tuple[Optional[int], Optional[float]]:
        """
        DLinear预测 (输入df包含历史数据，预测下一时刻)
        """
        if 'dlinear' not in self.models:
            self.load_model('dlinear')
            if 'dlinear' not in self.models:
                return None, None
            
        model = self.models['dlinear']
        preprocessor = self.preprocessors['dlinear']
        
        # 预处理
        X, _, _ = preprocessor.preprocess_for_inference(df)
        
        if len(X) == 0:
            return None, None
            
        # 取最后一个窗口
        X_last = torch.FloatTensor(X[-1:]).to(self.device)
        
        with torch.no_grad():
            logits = model(X_last).cpu().numpy()
            prob = 1 / (1 + np.exp(-logits))
            prediction = int(prob.item() > 0.5)
            probability = float(prob.item())
            
        return prediction, probability

    def predict_nbeats(self, df: pd.DataFrame) -> Tuple[Optional[int], Optional[float]]:
        """
        N-BEATS预测
        """
        if 'nbeats' not in self.models:
            self.load_model('nbeats')
            if 'nbeats' not in self.models:
                return None, None
            
        model = self.models['nbeats']
        preprocessor = self.preprocessors['nbeats']
        
        # 预处理
        X, _, _ = preprocessor.preprocess(df, fit=False)
        
        if len(X) == 0:
            return None, None
            
        # 取最后一个窗口
        X_last = torch.FloatTensor(X[-1:]).to(self.device)
        current_price_scaled = X[-1, -1] # 窗口中最后一个点的缩放价格
        
        with torch.no_grad():
            output = model(X_last).cpu().numpy()
            # output shape: (batch, forecast_horizon) or (batch, 1)
            if output.ndim == 1:
                pred_price_scaled = output[0]
            else:
                pred_price_scaled = output[0, 0]
                
            prediction = 1 if pred_price_scaled > current_price_scaled else 0
            
            # N-BEATS 回归模型没有直接概率
            probability = 0.6 if prediction == 1 else 0.4 # 占位
            
        return prediction, probability

    def predict_ttm(self, df: pd.DataFrame) -> Tuple[Optional[int], Optional[float]]:
        """
        TTM-R2预测
        """
        if 'ttm' not in self.models:
            self.load_model('ttm')
            if 'ttm' not in self.models:
                return None, None
            
        model = self.models['ttm']
        context_length = 512
        
        all_prices = df['close'].values
        if len(all_prices) == 0:
             return None, None
             
        # TTM 上下文准备
        if len(all_prices) >= context_length:
            context_data = all_prices[-context_length:].reshape(1, -1, 1)
        else:
            padded = np.full(context_length, all_prices[-1])
            padded[-len(all_prices):] = all_prices
            context_data = padded.reshape(1, -1, 1)
            
        context_tensor = torch.FloatTensor(context_data).to(self.device)
        base_price = float(df['close'].iloc[-1])
        
        with torch.no_grad():
            output = model(context_tensor)
            pred_price = output.prediction_outputs.cpu().numpy()[0, 0, 0]
            
        prediction = 1 if pred_price > base_price else 0
        
        # 简单的概率模拟
        price_pct = (pred_price - base_price) / base_price
        probability = 1 / (1 + np.exp(-price_pct * 100))
        
        return prediction, float(probability)

    def predict_all(self, df: pd.DataFrame) -> Dict[str, Any]:
        """
        使用所有加载的模型进行预测
        """
        results = {}
        
        # DLinear
        dl_pred, dl_prob = self.predict_dlinear(df)
        if dl_pred is not None:
            results['dlinear'] = {'prediction': dl_pred, 'probability': dl_prob}
            
        # N-BEATS
        nb_pred, nb_prob = self.predict_nbeats(df)
        if nb_pred is not None:
            results['nbeats'] = {'prediction': nb_pred, 'probability': nb_prob}
            
        # TTM
        ttm_pred, ttm_prob = self.predict_ttm(df)
        if ttm_pred is not None:
            results['ttm'] = {'prediction': ttm_pred, 'probability': ttm_prob}
            
        # 集成 (投票)
        votes = []
        if dl_pred is not None: votes.append(dl_pred)
        if nb_pred is not None: votes.append(nb_pred)
        if ttm_pred is not None: votes.append(ttm_pred)
        
        if votes:
            # 简单投票: > 50% 胜出
            vote_sum = sum(votes)
            if vote_sum > len(votes) / 2:
                ensemble_pred = 1
            elif vote_sum < len(votes) / 2:
                ensemble_pred = 0
            else:
                # 平票时，如果有 DLinear，优先 DLinear，否则 0 (保守)
                ensemble_pred = dl_pred if dl_pred is not None else 0
                
            results['ensemble'] = {'prediction': ensemble_pred, 'votes': votes}
        else:
            results['ensemble'] = None
            
        return results

def get_inference_manager() -> InferenceManager:
    """获取单例"""
    return InferenceManager()