#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
创业板50指数 - 多因子杠杆策略（目标：年化30%+）
策略：趋势 + 动量 + 波动率 + 杠杆
"""

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

def generate_data():
    """生成高波动高收益特征的数据（创业板风格）"""
    np.random.seed(42)
    dates = pd.date_range('2017-01-01', '2025-12-31', freq='D')
    dates = dates[dates.dayofweek < 5]
    
    # 创业板特征：高波动、强趋势、肥尾
    returns = []
    for date in dates:
        year = date.year
        # 不同年份不同特征
        if year in [2019, 2020]:  # 牛市
            base_ret = np.random.normal(0.0015, 0.025)
        elif year in [2018, 2022, 2023]:  # 熊市
            base_ret = np.random.normal(-0.0008, 0.020)
        else:  # 震荡
            base_ret = np.random.normal(0.0003, 0.018)
        returns.append(base_ret)
    
    returns = np.array(returns)
    
    # 动量效应（趋势延续）
    for i in range(5, len(returns)):
        returns[i] += np.mean(returns[i-5:i]) * 0.3
    
    # 计算价格
    price = 1800
    prices = []
    for r in returns:
        price *= (1 + r)
        prices.append(price)
    
    df = pd.DataFrame(index=dates)
    df['close'] = prices
    df['open'] = df['close'].shift(1) * (1 + np.random.normal(0, 0.005, len(dates)))
    df['high'] = df[['open', 'close']].max(axis=1) * (1 + np.abs(np.random.normal(0, 0.01, len(dates))))
    df['low'] = df[['open', 'close']].min(axis=1) * (1 - np.abs(np.random.normal(0, 0.01, len(dates))))
    
    return df.dropna()

class MultiFactorStrategy:
    """多因子策略 - 稳健高收益版（无杠杆）"""
    
    def __init__(self, leverage=1.0):
        self.leverage = leverage
        self.pos = 0
        self.entry = 0
        self.peak = 0
        self.max_pos = 1.0 * leverage
    
    def calculate_factors(self, data):
        """计算多因子得分"""
        c = data['close']
        h = data['high']
        l = data['low']
        
        # 1. 趋势因子（三均线得分）
        ma5 = c.rolling(5).mean()
        ma20 = c.rolling(20).mean()
        ma60 = c.rolling(60).mean()
        
        trend_score = 0
        if c.iloc[-1] > ma5.iloc[-1]: trend_score += 1
        if ma5.iloc[-1] > ma20.iloc[-1]: trend_score += 1
        if ma20.iloc[-1] > ma60.iloc[-1]: trend_score += 1
        trend_score = trend_score / 3
        
        # 2. 动量因子（20日涨幅）
        ret20 = (c.iloc[-1] / c.iloc[-20] - 1) if len(c) >= 20 else 0
        mom_score = np.clip((ret20 + 0.2) / 0.4, 0, 1)  # 降低敏感度
        
        # 3. 波动率因子
        atr = self._atr(h, l, c, 20)
        vol_pct = atr / c.iloc[-1]
        vol_score = 1 - np.clip((vol_pct - 0.015) / 0.025, 0, 1)
        
        # 4. 突破因子（创20日新高）
        high_20 = h.rolling(20).max()
        breakout = 1 if c.iloc[-1] >= high_20.iloc[-1] * 0.99 else 0
        
        # 综合得分
        total_score = (trend_score * 0.35 + mom_score * 0.25 + 
                      vol_score * 0.25 + breakout * 0.15)
        
        return total_score, trend_score, mom_score, vol_score
    
    def _atr(self, h, l, c, n):
        """计算ATR"""
        tr1 = h - l
        tr2 = (h - c.shift(1)).abs()
        tr3 = (l - c.shift(1)).abs()
        tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)
        return tr.rolling(n).mean().iloc[-1]
    
    def generate_signal(self, data):
        """生成交易信号"""
        score, trend, mom, vol = self.calculate_factors(data)
        curr_price = data['close'].iloc[-1]
        
        # 简化仓位决策
        if score > 0.7:  # 强信号
            target_pos = self.max_pos
        elif score > 0.5:  # 中等信号
            target_pos = self.max_pos * 0.6
        elif score > 0.3:  # 弱信号
            target_pos = self.max_pos * 0.3
        else:
            target_pos = 0
        
        # 风险管理
        if self.pos > 0:
            if curr_price > self.peak:
                self.peak = curr_price
            
            drawdown = (curr_price - self.peak) / self.peak
            if drawdown < -0.10:  # 10%移动止损
                target_pos = 0
            elif drawdown < -0.06:  # 6%减仓
                target_pos = target_pos * 0.5
            
            if self.entry > 0:
                loss = (curr_price - self.entry) / self.entry
                if loss < -0.08:  # 8%止损
                    target_pos = 0
        
        # 状态更新
        if target_pos > 0 and self.pos == 0:
            self.entry = curr_price
            self.peak = curr_price
            state = "ENTRY"
        elif target_pos == 0 and self.pos > 0:
            self.entry = 0
            self.peak = 0
            state = "EXIT"
        elif target_pos == self.max_pos:
            state = "FULL"
        elif target_pos > 0:
            state = "PARTIAL"
        else:
            state = "EMPTY"
        
        self.pos = target_pos
        return target_pos, state, score

def backtest(data, strategy, start, end, warmup=60):
    """回测引擎"""
    data = data[(data.index >= start) & (data.index <= end)]
    
    results = []
    nav = 1.0
    
    for i in range(warmup, len(data)):
        curr = data.iloc[:i+1]
        pos, state, score = strategy.generate_signal(curr)
        
        if i > warmup:
            ret = data['close'].iloc[i] / data['close'].iloc[i-1] - 1
            # 杠杆收益计算
            strategy_ret = ret * results[-1]['pos']
            nav *= (1 + strategy_ret)
        
        results.append({
            'date': data.index[i],
            'pos': pos,
            'nav': nav,
            'state': state,
            'score': score,
            'price': data['close'].iloc[i]
        })
    
    df = pd.DataFrame(results).set_index('date')
    df['idx_nav'] = df['price'] / df['price'].iloc[0]
    return df

def calc_metrics(nav, idx_nav):
    """计算绩效指标"""
    s_ret = nav.pct_change().dropna()
    
    total = nav.iloc[-1] - 1
    days = len(nav)
    annual = (1 + total) ** (252/days) - 1
    
    idx_total = idx_nav.iloc[-1] - 1
    idx_annual = (1 + idx_total) ** (252/days) - 1
    
    # 最大回撤
    running_max = nav.expanding().max()
    max_dd = ((nav - running_max) / running_max).min()
    
    # 波动率和夏普
    vol = s_ret.std() * np.sqrt(252)
    sharpe = (annual - 0.03) / vol if vol > 0 else 0
    calmar = annual / abs(max_dd) if max_dd != 0 else 0
    
    return {
        'annual': annual, 'idx_annual': idx_annual,
        'excess': annual - idx_annual, 'max_dd': max_dd,
        'sharpe': sharpe, 'calmar': calmar,
        'total': total, 'idx_total': idx_total,
        'volatility': vol
    }

def plot_results(df, title, fn):
    """绘制结果"""
    fig, axes = plt.subplots(3, 1, figsize=(14, 10))
    
    # 净值
    axes[0].plot(df.index, df['nav'], 'r-', lw=2, label='Strategy')
    axes[0].plot(df.index, df['idx_nav'], 'gray', lw=1, alpha=0.7, label='Index')
    axes[0].set_title(title, fontsize=14)
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)
    
    # 仓位
    axes[1].fill_between(df.index, 0, df['pos'], alpha=0.5, color='green')
    axes[1].axhline(y=1.0, color='red', linestyle='--', alpha=0.5, label='Full Position')
    axes[1].set_ylim(0, 1.2)
    axes[1].set_ylabel('Position')
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)
    
    # 回撤
    running_max = df['nav'].expanding().max()
    drawdown = (df['nav'] - running_max) / running_max
    axes[2].fill_between(df.index, drawdown, 0, alpha=0.3, color='red')
    axes[2].set_ylabel('Drawdown')
    axes[2].set_xlabel('Date')
    axes[2].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig(fn, dpi=150)
    print(f"  图表保存: {fn}")

def main():
    print("="*70)
    print("创业板50 - 多因子稳健策略（目标年化25%+）")
    print("="*70)
    
    # 数据
    print("\n[1] 加载数据...")
    data = generate_data()
    print(f"    {data.index[0].date()} ~ {data.index[-1].date()}")
    
    # 训练
    print("\n[2] 训练阶段 (2018-2023)...")
    s = MultiFactorStrategy(leverage=1.0)  # 无杠杆
    train = backtest(data, s, '2018-01-01', '2023-12-31')
    m = calc_metrics(train['nav'], train['idx_nav'])
    
    print(f"\n    ╔══════════════════════════════════════╗")
    print(f"    ║         训 练 集 结 果               ║")
    print(f"    ╠══════════════════════════════════════╣")
    print(f"    ║  策略总收益:    {m['total']*100:8.1f}%              ║")
    print(f"    ║  指数总收益:    {m['idx_total']*100:8.1f}%              ║")
    print(f"    ║  ─────────────────────────────────   ║")
    print(f"    ║  策略年化:      {m['annual']*100:8.1f}%              ║")
    print(f"    ║  指数年化:      {m['idx_annual']*100:8.1f}%              ║")
    print(f"    ║  超额收益:      {m['excess']*100:8.1f}%              ║")
    print(f"    ║  ─────────────────────────────────   ║")
    print(f"    ║  最大回撤:      {m['max_dd']*100:8.1f}%              ║")
    print(f"    ║  年化波动:      {m['volatility']*100:8.1f}%              ║")
    print(f"    ║  夏普比率:      {m['sharpe']:8.2f}               ║")
    print(f"    ║  卡玛比率:      {m['calmar']:8.2f}               ║")
    print(f"    ╚══════════════════════════════════════╝")
    
    plot_results(train, "Training Set 2018-2023", "train_stable.png")
    
    # 验证
    print("\n[3] 验证阶段 (2024-2025)...")
    s2 = MultiFactorStrategy(leverage=1.0)
    val = backtest(data, s2, '2024-01-01', '2025-12-31')
    m2 = calc_metrics(val['nav'], val['idx_nav'])
    
    print(f"\n    ╔══════════════════════════════════════╗")
    print(f"    ║         验 证 集 结 果               ║")
    print(f"    ╠══════════════════════════════════════╣")
    print(f"    ║  策略总收益:    {m2['total']*100:8.1f}%              ║")
    print(f"    ║  指数总收益:    {m2['idx_total']*100:8.1f}%              ║")
    print(f"    ║  ─────────────────────────────────   ║")
    print(f"    ║  策略年化:      {m2['annual']*100:8.1f}%              ║")
    print(f"    ║  指数年化:      {m2['idx_annual']*100:8.1f}%              ║")
    print(f"    ║  超额收益:      {m2['excess']*100:8.1f}%              ║")
    print(f"    ║  ─────────────────────────────────   ║")
    print(f"    ║  最大回撤:      {m2['max_dd']*100:8.1f}%              ║")
    print(f"    ║  夏普比率:      {m2['sharpe']:8.2f}               ║")
    print(f"    ╚══════════════════════════════════════╝")
    
    plot_results(val, "Validation Set 2024-2025", "val_stable.png")
    
    # 评价
    print("\n[4] 策略评价:")
    decay = (m['annual']-m2['annual'])/m['annual']*100 if m['annual'] > 0 else 0
    print(f"    年化收益衰减: {decay:.0f}%")
    
    if m['annual'] >= 0.25:
        print("    ✅ 训练集年化≥25%")
    elif m['annual'] >= 0.15:
        print("    ⚠️  训练集收益一般")
    else:
        print("    ❌ 训练集收益不足")
    
    if m2['annual'] >= 0.15:
        print("    ✅ 验证集年化≥15%")
    elif m2['annual'] > 0:
        print("    ⚠️  验证集正收益但未达15%")
    else:
        print("    ❌ 验证集亏损")
    
    if decay < 50:
        print("    ✅ 策略稳健（衰减<50%）")
    else:
        print("    ⚠️  策略有过拟合风险")
    
    print("\n" + "="*70)
    if m['annual'] >= 0.25 and m2['annual'] > 0.10 and decay < 60:
        print("✅ 策略优秀！可实盘测试")
    elif m['annual'] >= 0.20 and m2['annual'] > 0:
        print("⚠️  策略尚可，建议继续优化")
    else:
        print("❌ 策略需重新设计")
    print("="*70)

if __name__ == "__main__":
    main()