"""
模型推理服务 - 统一管理 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
import onnxruntime as ort
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__))))



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/model.onnx',
                'preprocessor_path': './models/checkpoints/dlinear_classification_2026/preprocessor.pkl'
            },
            'nbeats': {
                'model_path': './models/checkpoints/nbeats_btc_1h/model.onnx',
                'preprocessor_path': './models/checkpoints/nbeats_btc_1h/preprocessor.pkl'
            },
            'ttm': {
                'model_path': './models/checkpoints/ttm_r2_onnx/ttm_r2_crypto_channel_mixing.onnx'
            }
        }
        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

        try:
            self.models['dlinear'] = ort.InferenceSession(config['model_path'])
            
            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 (ONNX) 已加载")
        except Exception as e:
            print(f"DLinear 加载失败: {e}")

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

        try:
            self.models['nbeats'] = ort.InferenceSession(config['model_path'])
            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 (ONNX) 已加载")
        except Exception as e:
            print(f"N-BEATS 加载失败: {e}")

    def _load_ttm(self):
        config = self.configs['ttm']
        try:
            self.models['ttm'] = ort.InferenceSession(config['model_path'])
            print("  ✓ TTM-R2 (ONNX) 已加载")
        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
            
        session = self.models['dlinear']
        preprocessor = self.preprocessors['dlinear']
        
        # 预处理
        X, _, _ = preprocessor.preprocess_for_inference(df)
        
        if len(X) == 0:
            return None, None
            
        # 取最后一个窗口
        # ONNX 需要 float32
        X_last = X[-1:].astype(np.float32)
        
        try:
            input_name = session.get_inputs()[0].name
            output_name = session.get_outputs()[0].name
            
            logits = session.run([output_name], {input_name: X_last})[0]
            prob = 1 / (1 + np.exp(-logits))
            prediction = int(prob.item() > 0.5)
            probability = float(prob.item())
            
            return prediction, probability
        except Exception as e:
            print(f"DLinear 预测出错: {e}")
            return None, None

    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
            
        session = self.models['nbeats']
        preprocessor = self.preprocessors['nbeats']
        
        # 预处理
        X, _, _ = preprocessor.preprocess(df, fit=False)
        
        if len(X) == 0:
            return None, None
            
        # 取最后一个窗口
        X_last = X[-1:].astype(np.float32)
        current_price_scaled = X[-1, -1] # 窗口中最后一个点的缩放价格
        
        try:
            input_name = session.get_inputs()[0].name
            output_name = session.get_outputs()[0].name
            
            output = session.run([output_name], {input_name: X_last})[0]
            
            # 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
        except Exception as e:
            print(f"N-BEATS 预测出错: {e}")
            return None, None

    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
            
        session = 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)
            
        # ONNX input
        context_data = context_data.astype(np.float32)
        base_price = float(df['close'].iloc[-1])
        
        try:
            input_name = session.get_inputs()[0].name
            output_name = session.get_outputs()[0].name
            
            output = session.run([output_name], {input_name: context_data})[0]
            # output shape might be (batch, prediction_length, features)
            # prediction_length=96 usually for TTM zero-shot
            # We want the first prediction step [0, 0, 0]
            pred_price = output[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)
        except Exception as e:
            print(f"TTM 预测出错: {e}")
            return None, None

    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()