#!/usr/bin/env python3
"""
完整对比测试：验证新框架与现有实现的回测结果一致性

测试：
1. 分散化选股信号生成
2. 完整回测执行
"""

import sys
import yaml
import pandas as pd
import numpy as np
from pathlib import Path
from datetime import datetime

project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

# 现有实现
from strategies.rotation.engine import RotationStrategy

# 新框架实现
from framework.factors import FactorRegistry, FactorCombiner
from framework.execution import BacktestExecutor
from strategies.shared.factors.momentum import MomentumFactor
from strategies.shared.signals.selectors import TopNSelector


def load_config(config_path: str = "config/strategies/rotation.yaml") -> dict:
    """加载配置"""
    with open(config_path, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


def test_grouped_selection():
    """测试分散化选股信号生成"""
    print("=" * 60)
    print("测试分散化选股信号生成")
    print("=" * 60)
    
    # 模拟因子数据和分组映射
    dates = pd.date_range('2020-01-01', periods=50)
    
    factor_data = pd.DataFrame({
        '399006.SZ': [0.1] * 50,  # A股，创业板
        'H30269.CSI': [0.15] * 50,  # A股，红利低波
        'NDX': [0.2] * 50,  # 美股
        'N225': [0.18] * 50,  # 日本
        'AU.SHF': [0.12] * 50,  # 商品，黄金
    }, index=dates)
    
    # 分组映射（对应rotation.yaml中的market字段）
    group_mapping = {
        '399006.SZ': 'A',
        'H30269.CSI': 'A',
        'NDX': 'US',
        'N225': 'JP',
        'AU.SHF': 'COMMODITY',
    }
    
    # 新框架：使用TopNSelector进行分散化选股
    selector = TopNSelector(
        select_num=3,
        group_mapping=group_mapping,
        min_score=0.0,
        rebalance_threshold=0.0,
        rebalance_days=1
    )
    
    result = selector.generate(factor_data)
    
    print(f"\n信号生成结果:")
    print(f"  - 信号天数: {len(result)}")
    print(f"  - 第一天信号: {result['signal'].iloc[0]}")
    
    # 验证分散化选股逻辑
    # 每大类选Top1，然后按得分排序选Top3
    # A股：H30269.CSI(0.15) > 399006.SZ(0.1) → H30269.CSI
    # 美股：NDX(0.2) → NDX
    # 日本：N225(0.18) → N225
    # 商品：AU.SHF(0.12) → AU.SHF
    # 按得分排序：NDX(0.2) > N225(0.18) > H30269.CSI(0.15) > AU.SHF(0.12)
    # Top3: NDX, N225, H30269.CSI
    
    # 由于T+1移位，第一天信号为空
    if result['signal'].iloc[0] == '' or pd.isna(result['signal'].iloc[0]):
        # 检查第二天信号
        expected = "NDX,N225,H30269.CSI"
        actual = result['signal'].iloc[1]
        
        # 验证信号是否正确（忽略顺序）
        if actual:
            codes = set(actual.split(','))
            expected_codes = set(expected.split(','))
            
            if codes == expected_codes:
                print("✅ 分散化选股逻辑正确")
                print(f"  - 预期: {expected}")
                print(f"  - 实际: {actual}")
                return True
            else:
                print("⚠️ 分散化选股结果与预期不同")
                print(f"  - 预期: {expected}")
                print(f"  - 实际: {actual}")
                return False
    
    print("✅ 分散化选股测试通过")
    return True


def test_full_backtest_comparison():
    """测试完整回测对比"""
    print("\n" + "=" * 60)
    print("测试完整回测对比")
    print("=" * 60)
    
    config = load_config()
    
    if not config.get('end_date'):
        config['end_date'] = datetime.now().strftime('%Y-%m-%d')
    
    # 现有实现：完整回测
    old_strategy = RotationStrategy(config)
    old_strategy.run()
    
    old_result = old_strategy.backtest_result
    
    print(f"\n现有实现回测结果:")
    print(f"  - 回测天数: {len(old_result)}")
    print(f"  - 策略累计收益: {old_result['轮动策略净值'].iloc[-1] - 1:.2%}")
    print(f"  - 基准累计收益: {old_result['基准净值'].iloc[-1] - 1:.2%}")
    
    # 新框架：使用BacktestExecutor执行回测
    # 这里简化测试，使用现有策略的数据
    
    # 准备信号数据（使用现有的信号列，列名可能是中文的）
    signal_col = 'signal' if 'signal' in old_strategy.signals.columns else '信号'
    signals = old_strategy.signals[[signal_col]].copy()
    signals.columns = ['signal']  # 重命名为英文
    
    # 准备日收益率数据（需要日收益率_{code}格式的列）
    data = pd.DataFrame(index=old_strategy.data.index)
    
    for code in old_strategy.valid_codes:
        # 添加日收益率列（保持原有格式）
        if f"日收益率_{code}" in old_strategy.data.columns:
            data[f"日收益率_{code}"] = old_strategy.data[f"日收益率_{code}"]
    
    executor = BacktestExecutor(
        initial_capital=100000,
        trade_cost=config['trade_cost'],
        select_num=config['select_num'],
        benchmark_data=old_strategy.benchmark_data
    )
    
    portfolio = executor.execute(signals, data)
    
    # 检查回测结果
    if hasattr(portfolio, 'backtest_result'):
        new_result = portfolio.backtest_result
        
        print(f"\n新框架回测结果:")
        print(f"  - 回测天数: {len(new_result)}")
        print(f"  - 策略累计收益: {new_result['策略净值'].iloc[-1] - 1:.2%}")
        if '基准净值' in new_result.columns:
            print(f"  - 基准累计收益: {new_result['基准净值'].iloc[-1] - 1:.2%}")
        
        # 对比净值序列相关性
        if len(new_result) == len(old_result):
            common_dates = new_result.index.intersection(old_result.index)
            
            old_nav = old_result['轮动策略净值'].loc[common_dates]
            new_nav = new_result['策略净值'].loc[common_dates]
            
            correlation = old_nav.corr(new_nav)
            
            print(f"\n净值序列对比:")
            print(f"  - 相关系数: {correlation:.4f}")
            
            if correlation > 0.99:
                print("✅ 回测结果高度一致")
                return True
            elif correlation > 0.90:
                print("⚠️ 回测结果基本一致")
                return True
            else:
                print("❌ 回测结果差异较大")
                return False
    else:
        print("⚠️ 新框架回测结果未生成")
        return False


def main():
    """运行所有对比测试"""
    print("=" * 60)
    print("        新框架完整功能对比测试")
    print("=" * 60)
    
    results = []
    
    # 测试1：分散化选股
    try:
        r1 = test_grouped_selection()
        results.append(("分散化选股", r1))
    except Exception as e:
        print(f"❌ 分散化选股测试失败: {e}")
        results.append(("分散化选股", False))
    
    # 测试2：完整回测对比
    try:
        r2 = test_full_backtest_comparison()
        results.append(("完整回测", r2))
    except Exception as e:
        print(f"❌ 完整回测测试失败: {e}")
        results.append(("完整回测", False))
    
    # 总结
    print("\n" + "=" * 60)
    print("对比测试总结")
    print("=" * 60)
    
    for test_name, passed in results:
        status = "✅" if passed else "❌"
        print(f"{status} {test_name}")
    
    passed_count = sum(1 for _, p in results if p)
    print(f"\n通过: {passed_count}/{len(results)}")
    
    return passed_count == len(results)


if __name__ == "__main__":
    main()