#!/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 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 TrendStrategy:
    """趋势跟踪策略 - 激进高收益版"""
    
    def __init__(self):
        self.pos = 0
        self.entry = 0
        self.peak = 0
    
    def signal(self, data):
        c = data['close'].values
        if len(c) < 60:
            return 0
        
        # 技术指标 - 更短周期，更敏感
        ma3 = np.mean(c[-3:])
        ma10 = np.mean(c[-10:])
        ma30 = np.mean(c[-30:])
        
        # 价格创10日新高（更敏感）
        highest_10 = np.max(c[-10:])
        lowest_10 = np.min(c[-10:])
        
        curr = c[-1]
        
        # 突破买入：创10日新高
        breakout = (curr >= highest_10 * 0.995) and (ma3 > ma10)
        
        # 卖出：跌破10日最低点
        sell = (curr <= lowest_10 * 1.005) or (ma3 < ma10 * 0.97)
        
        if breakout and self.pos == 0:
            return 1.0  # 满仓
        elif sell and self.pos > 0:
            return 0.0  # 清仓
        else:
            return self.pos
    
    def generate(self, data):
        new_pos = self.signal(data)
        
        curr_price = data['close'].iloc[-1]
        
        # 移动止损 - 更宽松的10%
        if self.pos > 0:
            if curr_price > self.peak:
                self.peak = curr_price
            if curr_price < self.peak * 0.90:
                new_pos = 0
        
        # 更新状态
        if new_pos > 0 and self.pos == 0:
            self.entry = curr_price
            self.peak = curr_price
            state = "BUY"
        elif new_pos == 0 and self.pos > 0:
            self.entry = 0
            self.peak = 0
            state = "SELL"
        elif new_pos > 0:
            state = "HOLD"
        else:
            state = "EMPTY"
        
        self.pos = new_pos
        return new_pos, state

def backtest(data, strategy, start, end, warmup=60):
    data = data[(data.index >= start) & (data.index <= end)]
    
    nav = 1.0
    results = []
    
    for i in range(warmup, len(data)):
        curr = data.iloc[:i+1]
        pos, state = strategy.generate(curr)
        
        if i > warmup:
            ret = data['close'].iloc[i] / data['close'].iloc[i-1] - 1
            nav *= (1 + ret * results[-1]['pos'])
        
        results.append({
            'date': data.index[i],
            'pos': pos,
            'nav': nav,
            'state': state,
            '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):
    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 = nav.pct_change().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
    }

def plot(df, title, fn):
    fig, ax = plt.subplots(2, 1, figsize=(14, 8))
    
    ax[0].plot(df.index, df['nav'], 'r-', lw=2, label='Strategy')
    ax[0].plot(df.index, df['idx_nav'], 'gray', lw=1, alpha=0.6, label='Index')
    ax[0].set_title(title, fontsize=14)
    ax[0].legend()
    ax[0].grid(True, alpha=0.3)
    
    ax[1].fill_between(df.index, 0, df['pos'], alpha=0.5, color='green')
    ax[1].set_ylim(0, 1.1)
    ax[1].set_ylabel('Position')
    ax[1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig(fn, dpi=150)
    print(f"  图表: {fn}")

def main():
    print("="*60)
    print("创业板50 - 趋势突破策略")
    print("="*60)
    
    data = load_real_data()
    print(f"\n数据: {data.index[0].date()} ~ {data.index[-1].date()}")
    
    # 训练
    print("\n【训练集 2018-2023】")
    s = TrendStrategy()
    train = backtest(data, s, '2018-01-01', '2023-12-31')
    m = calc_metrics(train['nav'], train['idx_nav'])
    
    print(f"  策略收益: {m['total']*100:7.1f}%  (年化 {m['annual']*100:5.1f}%)")
    print(f"  指数收益: {m['idx_total']*100:7.1f}%  (年化 {m['idx_annual']*100:5.1f}%)")
    print(f"  超额收益: {m['excess']*100:7.1f}%")
    print(f"  最大回撤: {m['max_dd']*100:7.1f}%")
    print(f"  夏普比率: {m['sharpe']:7.2f}")
    print(f"  卡玛比率: {m['calmar']:7.2f}")
    
    plot(train, "Training 2018-2023", "train_trend.png")
    
    # 验证
    print("\n【验证集 2024-2025】")
    s2 = TrendStrategy()
    val = backtest(data, s2, '2024-01-01', '2025-12-31')
    m2 = calc_metrics(val['nav'], val['idx_nav'])
    
    print(f"  策略收益: {m2['total']*100:7.1f}%  (年化 {m2['annual']*100:5.1f}%)")
    print(f"  指数收益: {m2['idx_total']*100:7.1f}%  (年化 {m2['idx_annual']*100:5.1f}%)")
    print(f"  超额收益: {m2['excess']*100:7.1f}%")
    print(f"  最大回撤: {m2['max_dd']*100:7.1f}%")
    print(f"  夏普比率: {m2['sharpe']:7.2f}")
    
    plot(val, "Validation 2024-2025", "val_trend.png")
    
    # 评价
    print("\n【策略评价】")
    if m['annual'] > 0.30:
        print("  ✅ 训练集年化超30%，高收益潜力")
    elif m['annual'] > 0.15:
        print("  ✅ 训练集表现良好")
    else:
        print("  ⚠️  训练集收益一般")
    
    if m2['annual'] > 0:
        print("  ✅ 验证集正收益")
    else:
        print("  ❌ 验证集亏损")
    
    print("\n" + "="*60)

if __name__ == "__main__":
    main()