#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
创业板50指数真实数据回测 - 高收益优化版
"""

import pandas as pd
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')

# 加载真实数据
def load_real_data():
    """加载创业板50指数真实数据 - cyb50_baostock.csv"""
    df = pd.read_csv('cyb50_baostock.csv')
    df['date'] = pd.to_datetime(df['date'])
    df = df.set_index('date').sort_index()
    
    # 转换数据类型
    for col in ['open', 'high', 'low', 'close', 'volume']:
        df[col] = pd.to_numeric(df[col], errors='coerce')
    
    print(f"真实数据加载成功: {df.index[0].date()} ~ {df.index[-1].date()}")
    return df

# ==================== 高性能策略 ====================

class HighPerformanceStrategy:
    """
    高收益策略：趋势跟踪 + 动量加速 + 智能止盈
    """
    
    def __init__(self, params=None):
        self.params = params or {
            'fast_ma': 5,       # 超短均线，快速响应
            'slow_ma': 20,      # 月均线
            'trend_ma': 60,     # 季均线
            'momentum_period': 10,
            'volatility_period': 20,
            'max_position': 1.0,
            'profit_take': 0.15,    # 15%止盈
            'trailing_stop': 0.08,  # 8%移动止损
        }
        # 确保所有参数都有默认值
        default_params = {
            'fast_ma': 5,
            'slow_ma': 20,
            'trend_ma': 60,
            'momentum_period': 10,
            'volatility_period': 20,
            'max_position': 1.0,
            'profit_take': 0.15,
            'trailing_stop': 0.08,
        }
        if params:
            for key, val in default_params.items():
                if key not in params:
                    self.params[key] = val
        self.position = 0
        self.entry_price = None
        self.max_price = None
    
    def generate_signal(self, data):
        """生成交易信号"""
        close = data['close']
        p = self.params
        
        # 计算指标
        ma_fast = close.rolling(p['fast_ma']).mean().iloc[-1]
        ma_slow = close.rolling(p['slow_ma']).mean().iloc[-1]
        ma_trend = close.rolling(p['trend_ma']).mean().iloc[-1]
        
        # 动量
        momentum = (close.iloc[-1] / close.iloc[-p['momentum_period']] - 1) * 100
        
        # 波动率
        returns = close.pct_change()
        vol = returns.rolling(p['volatility_period']).std().iloc[-1] * np.sqrt(252) * 100
        
        curr_price = close.iloc[-1]
        
        # 趋势强度
        trend_strong = (curr_price > ma_fast) and (ma_fast > ma_slow) and (ma_slow > ma_trend)
        trend_weak = (curr_price < ma_fast) and (ma_fast < ma_slow)
        
        # 信号生成
        if trend_strong and momentum > 2:
            # 强势上涨，满仓
            target_pos = p['max_position']
            state = "STRONG_UP"
        elif trend_strong and momentum > 0:
            # 趋势向上但动量一般，80%仓位
            target_pos = p['max_position'] * 0.8
            state = "UP"
        elif trend_weak or momentum < -3:
            # 趋势转弱，空仓
            target_pos = 0
            state = "DOWN"
        else:
            # 震荡，50%仓位
            target_pos = p['max_position'] * 0.5
            state = "OSCILLATE"
        
        # 移动止盈
        if self.position > 0 and self.max_price:
            current_return = (curr_price - self.entry_price) / self.entry_price
            
            # 更新最高价
            if curr_price > self.max_price:
                self.max_price = curr_price
            
            # 移动止损：从最高点回撤8%离场
            drawdown_from_peak = (curr_price - self.max_price) / self.max_price
            if drawdown_from_peak < -p['trailing_stop']:
                target_pos = 0
                state = "TRAILING_STOP"
            
            # 固定止盈15%
            elif current_return > p['profit_take']:
                target_pos = 0.5  # 减半仓，锁定利润
                state = "PROFIT_TAKE"
        
        # 更新状态
        if target_pos > 0 and self.position == 0:
            self.entry_price = curr_price
            self.max_price = curr_price
        elif target_pos == 0:
            self.entry_price = None
            self.max_price = None
        
        self.position = target_pos
        return target_pos, state

# ==================== 回测引擎 ====================

def backtest(data, strategy, start_date=None, end_date=None, warmup=60):
    """回测引擎"""
    if start_date:
        data = data[data.index >= start_date]
    if end_date:
        data = data[data.index <= end_date]
    
    results = []
    nav = 1.0
    
    for i in range(warmup, len(data)):
        curr_data = data.iloc[:i+1]
        position, state = strategy.generate_signal(curr_data)
        
        if i > warmup:
            daily_return = data['close'].iloc[i] / data['close'].iloc[i-1] - 1
            strategy_return = daily_return * results[-1]['position'] if results else 0
            nav *= (1 + strategy_return)
        
        results.append({
            'date': data.index[i],
            'position': position,
            'nav': nav,
            'state': state,
            'close': data['close'].iloc[i]
        })
    
    df = pd.DataFrame(results).set_index('date')
    df['index_nav'] = df['close'] / df['close'].iloc[0]
    
    metrics = calculate_metrics(df['nav'], df['index_nav'])
    return df, metrics

def calculate_metrics(strategy_nav, index_nav):
    """计算绩效指标"""
    s_returns = strategy_nav.pct_change().dropna()
    
    total_return = strategy_nav.iloc[-1] - 1
    days = len(strategy_nav)
    annual_return = (1 + total_return) ** (252 / days) - 1
    
    index_return = index_nav.iloc[-1] - 1
    index_annual = (1 + index_return) ** (252 / days) - 1
    
    running_max = strategy_nav.expanding().max()
    max_dd = ((strategy_nav - running_max) / running_max).min()
    
    volatility = s_returns.std() * np.sqrt(252)
    sharpe = (annual_return - 0.03) / volatility if volatility > 0 else 0
    calmar = annual_return / abs(max_dd) if max_dd != 0 else 0
    win_rate = (s_returns > 0).mean()
    
    return {
        'annual_return': annual_return,
        'index_annual': index_annual,
        'excess_annual': annual_return - index_annual,
        'max_drawdown': max_dd,
        'sharpe': sharpe,
        'calmar': calmar,
        'win_rate': win_rate,
        'total_return': total_return,
        'index_return': index_return
    }

def plot_results(results, title, filename):
    """绘制回测图表"""
    fig, axes = plt.subplots(3, 1, figsize=(14, 10))
    
    ax1 = axes[0]
    ax1.plot(results.index, results['nav'], label='Strategy', linewidth=2, color='red')
    ax1.plot(results.index, results['index_nav'], label='Index', linewidth=1, color='gray', alpha=0.7)
    ax1.set_title(f'{title}', fontsize=14)
    ax1.set_ylabel('NAV')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    ax2 = axes[1]
    colors = {'STRONG_UP': 'green', 'UP': 'lightgreen', 'DOWN': 'red', 
              'OSCILLATE': 'yellow', 'TRAILING_STOP': 'orange', 'PROFIT_TAKE': 'blue'}
    pos_colors = [colors.get(s, 'gray') for s in results['state']]
    ax2.fill_between(results.index, 0, results['position'], alpha=0.5, color='green')
    ax2.set_ylabel('Position')
    ax2.set_ylim(0, 1.1)
    ax2.grid(True, alpha=0.3)
    
    ax3 = axes[2]
    running_max = results['nav'].expanding().max()
    drawdown = (results['nav'] - running_max) / running_max
    ax3.fill_between(results.index, drawdown, 0, alpha=0.3, color='red')
    ax3.set_ylabel('Drawdown')
    ax3.set_xlabel('Date')
    ax3.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig(filename, dpi=150, bbox_inches='tight')
    print(f"  图表已保存: {filename}")

# ==================== 主程序 ====================

def main():
    print("="*70)
    print("创业板50指数高收益策略回测")
    print("="*70)
    
    # 加载真实数据
    print("\n[1] 加载真实数据...")
    data = load_real_data()
    
    print(f"    数据区间: {data.index[0].date()} ~ {data.index[-1].date()}")
    
    # 训练阶段
    print("\n[2] 训练阶段 (2018-2023) - 优化参数...")
    
    # 测试多组参数，找最优
    best_params = None
    best_score = -999
    
    test_configs = [
        {'fast_ma': 5, 'slow_ma': 20, 'profit_take': 0.15, 'trailing_stop': 0.08},
        {'fast_ma': 3, 'slow_ma': 15, 'profit_take': 0.12, 'trailing_stop': 0.06},
        {'fast_ma': 10, 'slow_ma': 30, 'profit_take': 0.20, 'trailing_stop': 0.10},
    ]
    
    for cfg in test_configs:
        strategy = HighPerformanceStrategy(cfg)
        results, metrics = backtest(data, strategy, start_date='2018-01-01', end_date='2023-12-31')
        
        # 评分：收益优先
        score = metrics['annual_return'] * 0.5 + metrics['calmar'] * 0.3 + metrics['sharpe'] * 0.2
        
        print(f"\n  参数: {cfg}")
        print(f"  年化: {metrics['annual_return']*100:.1f}%, 回撤: {metrics['max_drawdown']*100:.1f}%, 评分: {score:.2f}")
        
        if score > best_score and metrics['max_drawdown'] > -0.40:
            best_score = score
            best_params = cfg
    
    print(f"\n  最优参数: {best_params}")
    
    # 用最优参数重新回测训练集
    strategy = HighPerformanceStrategy(best_params)
    train_results, train_metrics = backtest(data, strategy, start_date='2018-01-01', end_date='2023-12-31')
    
    print(f"\n  训练集最终表现:")
    print(f"  ┌─────────────────────────────────────┐")
    print(f"  │ 策略年化收益: {train_metrics['annual_return']*100:>8.2f}%      │")
    print(f"  │ 指数年化收益: {train_metrics['index_annual']*100:>8.2f}%      │")
    print(f"  │ 超额收益:     {train_metrics['excess_annual']*100:>8.2f}%      │")
    print(f"  │ 最大回撤:     {train_metrics['max_drawdown']*100:>8.2f}%      │")
    print(f"  │ 夏普比率:     {train_metrics['sharpe']:>8.2f}          │")
    print(f"  │ 卡玛比率:     {train_metrics['calmar']:>8.2f}          │")
    print(f"  │ 胜率:         {train_metrics['win_rate']*100:>8.1f}%        │")
    print(f"  └─────────────────────────────────────┘")
    
    plot_results(train_results, "Training Set (2018-2023)", "train_high_perf.png")
    
    # 验证阶段
    print(f"\n[3] 验证阶段 (2024-2025) - 样本外测试...")
    strategy_val = HighPerformanceStrategy(best_params)
    val_results, val_metrics = backtest(data, strategy_val, start_date='2024-01-01', end_date='2025-12-31')
    
    print(f"\n  验证集最终表现:")
    print(f"  ┌─────────────────────────────────────┐")
    print(f"  │ 策略年化收益: {val_metrics['annual_return']*100:>8.2f}%      │")
    print(f"  │ 指数年化收益: {val_metrics['index_annual']*100:>8.2f}%      │")
    print(f"  │ 超额收益:     {val_metrics['excess_annual']*100:>8.2f}%      │")
    print(f"  │ 最大回撤:     {val_metrics['max_drawdown']*100:>8.2f}%      │")
    print(f"  │ 夏普比率:     {val_metrics['sharpe']:>8.2f}          │")
    print(f"  │ 卡玛比率:     {val_metrics['calmar']:>8.2f}          │")
    print(f"  └─────────────────────────────────────┘")
    
    plot_results(val_results, "Validation Set (2024-2025)", "val_high_perf.png")
    
    # 过拟合检测
    print(f"\n[4] 过拟合检测:")
    return_decay = (train_metrics['annual_return'] - val_metrics['annual_return']) / train_metrics['annual_return'] if train_metrics['annual_return'] != 0 else 0
    print(f"    年化收益衰减: {return_decay*100:.1f}%")
    if return_decay > 0.5:
        print("    ⚠️  策略在验证集表现下降明显")
    else:
        print("    ✓ 策略稳健性良好")
    
    print("\n" + "="*70)
    print("回测完成！")
    print("="*70)

if __name__ == "__main__":
    main()