# -*- coding: utf-8 -*-
"""
多因子机器学习股价预测引擎 - 预测引擎
======================================
基于FT-Transformer的多因子预测，含IC/IR验证、共线性排查
纯NumPy实现，支持CPU环境
"""

import numpy as np
import pandas as pd
import json
import os
import time
import sys
import traceback
from datetime import datetime

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from multifactor_config import (
    MULTIFACTOR_MODEL_CONFIG, TRAINING_CONFIG, CHECKPOINT_DIR, REPORT_DIR
)
from multifactor_feature_engineering import MultiFactorFeatureEngineer, MultiFactorDataFetcher
from data_adapter import get_data_adapter

# 复用FT-Transformer核心（已优化SGA梯度）
from ft_transformer_model import FTTransformer, RewardMechanism, RewardWeightedLoss
from training_engine import Trainer


class FactorValidator:
    """因子有效性验证器"""
    
    def __init__(self, config=None):
        self.config = config or TRAINING_CONFIG['validation_constraints']
    
    def compute_ic(self, predictions, actual_returns):
        """计算Information Coefficient (Spearman Rank IC)"""
        n = len(predictions)
        if n < 10:
            return 0.0
        
        # 计算排名
        pred_ranks = pd.Series(predictions).rank().values
        actual_ranks = pd.Series(actual_returns).rank().values
        
        # Pearson相关系数（对排名）
        pred_z = (pred_ranks - pred_ranks.mean()) / (pred_ranks.std() + 1e-10)
        actual_z = (actual_ranks - actual_ranks.mean()) / (actual_ranks.std() + 1e-10)
        ic = np.mean(pred_z * actual_z)
        return ic
    
    def compute_ir(self, ic_series):
        """计算Information Ratio (IC均值/IC标准差)"""
        if len(ic_series) < 3:
            return 0.0
        return np.mean(ic_series) / (np.std(ic_series) + 1e-10)
    
    def compute_direction_accuracy(self, predictions, actual_returns):
        """计算方向正确率"""
        if len(predictions) < 5:
            return 0.0
        pred_direction = np.sign(predictions)
        actual_direction = np.sign(actual_returns)
        correct = np.sum(pred_direction == actual_direction)
        return correct / len(predictions)
    
    def check_collinearity(self, feature_matrix, names=None, threshold=0.7):
        """检查因子共线性，返回需要剔除的因子索引"""
        if feature_matrix.shape[1] < 2:
            return []
        
        # 标准化
        X = (feature_matrix - feature_matrix.mean(axis=0)) / (feature_matrix.std(axis=0) + 1e-10)
        
        # 相关矩阵
        corr = np.corrcoef(X.T)
        
        # 找高相关对
        to_remove = set()
        n_features = corr.shape[0]
        
        for i in range(n_features):
            for j in range(i+1, n_features):
                if abs(corr[i, j]) > threshold and i not in to_remove:
                    to_remove.add(j)
        
        removed_names = [names[i] for i in to_remove] if names else list(to_remove)
        return list(to_remove), removed_names
    
    def validate(self, predictions, actual_returns, feature_matrix=None, feature_names=None):
        """完整验证流程"""
        results = {}
        
        # 方向正确率
        dir_acc = self.compute_direction_accuracy(predictions, actual_returns)
        results['direction_accuracy'] = round(dir_acc, 4)
        results['direction_pass'] = dir_acc >= self.config['min_direction_accuracy']
        
        # IC
        ic = self.compute_ic(predictions, actual_returns)
        results['ic'] = round(ic, 4)
        results['ic_pass'] = abs(ic) >= self.config['monthly_ic_min']
        
        # IR（模拟月度IC序列）
        if len(predictions) >= 20:
            chunk_size = len(predictions) // 5  # 分5段模拟月度
            ic_series = []
            for k in range(5):
                start = k * chunk_size
                end = start + chunk_size
                if end <= len(predictions):
                    seg_ic = self.compute_ic(predictions[start:end], actual_returns[start:end])
                    ic_series.append(seg_ic)
            ir = self.compute_ir(np.array(ic_series))
            results['ir'] = round(ir, 4)
            results['monthly_ic_mean'] = round(float(np.mean(ic_series)), 4)
            results['ir_pass'] = ir >= self.config['monthly_ir_min']
        else:
            results['ir'] = 0.0
            results['monthly_ic_mean'] = ic
            results['ir_pass'] = False
        
        # 共线性排查
        if feature_matrix is not None and self.config['auto_prune']:
            remove_idx, remove_names = self.check_collinearity(
                feature_matrix, feature_names, self.config['max_corr_threshold']
            )
            results['collinear_features'] = remove_names
            results['collinear_count'] = len(remove_names)
            results['needs_pruning'] = len(remove_idx) > 0
        
        # 总评
        passed = sum([results.get('direction_pass', False), results.get('ic_pass', False), results.get('ir_pass', False)])
        results['total_checks'] = 3
        results['passed_checks'] = passed
        results['status'] = 'PASS' if passed >= 2 else 'NEEDS_IMPROVEMENT'
        
        return results


class MultiFactorEngine:
    """多因子预测引擎"""
    
    def __init__(self):
        self.model = None
        self.feature_engineer = MultiFactorFeatureEngineer(sequence_length=20)
        self.validator = FactorValidator()
        self.feature_info = self.feature_engineer.get_feature_info()
        self.is_trained = False
        self.train_report = {}
    
    def _build_model_config(self, n_features=None):
        """构建模型配置"""
        cfg = MULTIFACTOR_MODEL_CONFIG.copy()
        if n_features:
            cfg['num_continuous_features'] = n_features
        cfg['num_categorical_features'] = 1
        cfg['cat_cardinalities'] = [2]
        return cfg
    
    def train(self, df, financial_data=None, mode='quick'):
        """
        训练多因子模型
        
        mode: 'quick' (模拟数据快速验证) / 'full' (真实数据完整训练)
        """
        start_time = time.time()
        print("[多因子引擎] 开始训练...")
        
        if mode == 'quick':
            return self._quick_train()
        
        # 完整训练
        features_df = self.feature_engineer.extract_all_features(df, financial_data)
        X_cont, X_cat, labels = self.feature_engineer.create_sequences(features_df, df)
        
        if len(X_cont) < 50:
            print(f"[WARN] 样本不足: {len(X_cont)}")
            return None
        
        n_features = X_cont.shape[2]
        
        # 数据划分
        N = len(X_cont)
        train_size = int(N * 0.75)
        X_train, X_val = X_cont[:train_size], X_cont[train_size:]
        cat_train, cat_val = X_cat[:train_size], X_cat[train_size:]
        y_train, y_val = labels[:train_size], labels[train_size:]
        
        print(f"  训练集: {len(X_train)} 样本, 验证集: {len(X_val)} 样本")
        print(f"  特征维度: {n_features}")
        
        # 构建模型
        # 快速模式使用更小的模型
        cfg = self._build_model_config(n_features)
        cfg['d_token'] = 32
        cfg['n_blocks'] = 1
        cfg['n_heads'] = 2
        model_config = cfg
        self.model = FTTransformer(model_config)
        
        # 训练
        train_cfg = TRAINING_CONFIG['optimizer'].copy()
        if mode == 'fast':
            train_cfg['max_epochs'] = 10
            train_cfg['early_stopping_patience'] = 3
        
        reward = RewardMechanism({
            'trend_accuracy_weight': 0.4,
            'return_mse_weight': 0.3,
            'volatility_penalty_weight': 0.1,
            'drawdown_penalty_weight': 0.2,
        })
        loss_fn = RewardWeightedLoss(reward)
        
        trainer = Trainer(self.model, loss_fn, train_cfg, CHECKPOINT_DIR, REPORT_DIR)
        history = trainer.train(X_train, cat_train, y_train, X_val, cat_val, y_val, verbose=True)
        
        # 验证
        val_preds, _ = self.model.predict(X_val, cat_val)
        val_results = self.validator.validate(
            val_preds[:, 0], y_val[:, 0],
            X_val[:, -1, :],  # 最后时间步特征
            self.feature_info['feature_names'][:n_features] if n_features <= len(self.feature_info['feature_names']) else None
        )
        
        self.is_trained = True
        self.train_report = {
            'mode': mode,
            'train_samples': len(X_train),
            'val_samples': len(X_val),
            'n_features': n_features,
            'best_epoch': trainer.best_epoch,
            'best_val_loss': float(trainer.best_val_loss),
            'validation': val_results,
            'total_time': round(time.time() - start_time, 1),
            'feature_info': {
                k: len(v) for k, v in self.feature_info['categories'].items()
            },
        }
        
        print(f"\n[多因子引擎] 训练完成")
        print(f"  方向正确率: {val_results['direction_accuracy']:.2%}")
        print(f"  IC: {val_results['ic']:.4f}, IR: {val_results['ir']:.4f}")
        print(f"  状态: {val_results['status']}")
        
        return self.model, history, val_results
    
    def _quick_train(self):
        """快速演示模式（模拟数据）"""
        start_time = time.time()
        print("[多因子引擎] 快速演示模式（模拟数据）")
        
        np.random.seed(42)
        N = 300
        seq_len = 20
        n_features = 46  # 总特征数
        
        X_cont = np.random.randn(N, seq_len, n_features).astype(np.float32)
        X_cat = np.zeros((N, seq_len, 1), dtype=np.int32)
        
        # 模拟标签（加入信号使模型能学到东西）
        trend_signal = np.mean(X_cont[:, -5:, :3], axis=(1, 2))  # 用前3个特征做信号
        y = np.zeros((N, 3), dtype=np.float32)
        y[:, 0] = np.sign(trend_signal + np.random.randn(N) * 0.5)
        y[:, 1] = np.tanh(trend_signal * 2)
        y[:, 2] = np.clip(np.abs(np.random.randn(N)) * 0.1, 0, 0.5)
        
        train_size = int(N * 0.7)
        # 快速模式使用更小的模型
        cfg = self._build_model_config(n_features)
        cfg['d_token'] = 32
        cfg['n_blocks'] = 1
        cfg['n_heads'] = 2
        model_config = cfg
        self.model = FTTransformer(model_config)
        
        reward = RewardMechanism({
            'weights': {'wave_structure': 0.3, 'volume_signal': 0.25,
                       'volume_price_fit': 0.25, 'pattern_fit': 0.2},
            'wave_structure': {'five_wave_complete_score': 100, 'three_wave_correct_score': 80,
                              'wave_ratio_tolerance': 0.15, 'fibonacci_alignment_bonus': 10, 'wave_count_penalty': -5},
            'volume_signal': {'double_volume_threshold': 1.8, 'volume_break_confirm_days': 3,
                             'shrink_volume_ratio': 0.5, 'signal_hit_score': 100, 'false_signal_penalty': -20},
            'volume_price_fit': {'price_volume_sync_score': 100, 'divergence_detect_threshold': 0.3,
                                'trend_confirm_bonus': 15, 'false_divergence_penalty': -15},
            'pattern_fit': {'bullish_pattern_hit': 100, 'bearish_pattern_hit': 100,
                           'pattern_accuracy_weight': 0.7, 'timing_accuracy_weight': 0.3},
            'fit_threshold': {'success_threshold': 0.8, 'warning_threshold': 0.6, 'failure_threshold': 0.6},
        })
        loss_fn = RewardWeightedLoss(reward)
        
        train_cfg = {
            'batch_size': 32, 'learning_rate': 1e-2, 'max_epochs': 5,
            'early_stopping_patience': 2, 'gradient_clip_norm': 1.0,
            'warmup_epochs': 1, 'scheduler_type': 'cosine', 'weight_decay': 1e-5,
        }
        
        trainer = Trainer(self.model, loss_fn, train_cfg, CHECKPOINT_DIR, REPORT_DIR)
        history = trainer.train(
            X_cont[:train_size], X_cat[:train_size], y[:train_size],
            X_cont[train_size:], X_cat[train_size:], y[train_size:], verbose=True
        )
        
        val_preds, _ = self.model.predict(X_cont[train_size:], X_cat[train_size:])
        val_results = self.validator.validate(val_preds[:, 0], y[train_size:, 0])
        
        self.is_trained = True
        self.train_report = {
            'mode': 'quick',
            'train_samples': train_size,
            'val_samples': N - train_size,
            'n_features': n_features,
            'best_epoch': trainer.best_epoch,
            'best_val_loss': float(trainer.best_val_loss),
            'validation': val_results,
            'total_time': round(time.time() - start_time, 1),
            'feature_info': {
                k: len(v) for k, v in self.feature_info['categories'].items()
            },
        }
        
        print(f"\n[多因子引擎] 快速演示完成: {self.train_report['total_time']}s")
        return self.model, history, val_results
    
    def predict(self, df, financial_data=None):
        """对股票做预测"""
        if not self.is_trained or self.model is None:
            return None
        
        features_df = self.feature_engineer.extract_all_features(df, financial_data)
        X_cont, X_cat, _ = self.feature_engineer.create_sequences(features_df, df)
        
        if len(X_cont) == 0:
            return None
        
        # 用最后序列预测
        last_seq = X_cont[-1:].astype(np.float32)
        last_cat = X_cat[-1:].astype(np.int32)
        
        preds, attn = self.model.predict(last_seq, last_cat)
        
        # 特征重要性分析（用注意力权重）
        last_features = features_df.iloc[-len(X_cont):]
        
        trend_val = preds[0, 0]
        price_chg = preds[0, 1]
        volatility = preds[0, 2]
        
        if trend_val > 0.15:
            trend_str = '看涨'
            confidence = min(trend_val * 50, 95)
        elif trend_val < -0.15:
            trend_str = '看跌'
            confidence = min(abs(trend_val) * 50, 95)
        else:
            trend_str = '震荡'
            confidence = 50
        
        last_close = float(df['close'].iloc[-1])
        target_price = last_close * (1 + price_chg)
        
        result = {
            'trend_direction': trend_str,
            'confidence': round(confidence, 1),
            'trend_score': round(float(trend_val), 4),
            'price_change_pct': round(float(price_chg) * 100, 2),
            'target_price': round(float(target_price), 2),
            'last_close': round(last_close, 2),
            'volatility': round(float(volatility), 4),
            'data_points': len(df),
            'feature_count': X_cont.shape[2],
        }
        
        return result
    
    def get_factor_contributions(self, df, financial_data=None):
        """分析各因子类别对预测的贡献"""
        if not self.is_trained or self.model is None:
            return None
        
        features_df = self.feature_engineer.extract_all_features(df, financial_data)
        X_cont, X_cat, _ = self.feature_engineer.create_sequences(features_df, df)
        
        if len(X_cont) == 0:
            return None
        
        last_seq = X_cont[-1:].astype(np.float32)
        last_cat = X_cat[-1:].astype(np.int32)
        
        # 基线预测
        base_pred, _ = self.model.predict(last_seq, last_cat)
        base_trend = base_pred[0, 0]
        
        # 逐类因子消融
        info = self.feature_info['categories']
        contributions = {}
        offset = 0
        
        for cat_name, cat_features in info.items():
            n_cat = len(cat_features)
            if offset + n_cat > last_seq.shape[2]:
                offset += n_cat
                continue
            
            # 置零该类因子
            ablated = last_seq.copy()
            ablated[:, :, offset:offset + n_cat] = 0
            
            ablated_pred, _ = self.model.predict(ablated, last_cat)
            contribution = base_trend - ablated_pred[0, 0]
            
            contributions[cat_name] = {
                'contribution': round(float(contribution), 4),
                'n_features': n_cat,
                'importance_pct': 0,  # 后面归一化
            }
            offset += n_cat
        
        # 归一化
        total = sum(abs(v['contribution']) for v in contributions.values()) + 1e-10
        for v in contributions.values():
            v['importance_pct'] = round(abs(v['contribution']) / total * 100, 1)
        
        return contributions
    
    def save_report(self):
        """保存训练报告"""
        if not self.train_report:
            return None
        
        os.makedirs(REPORT_DIR, exist_ok=True)
        report_path = os.path.join(REPORT_DIR, 'multifactor_report.json')
        
        report = self.train_report.copy()
        report['timestamp'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        
        with open(report_path, 'w', encoding='utf-8') as f:
            json.dump(report, f, indent=2, ensure_ascii=False, default=str)
        
        return report_path