#!/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 HistoricalStrategy:
    """趋势策略 - 针对真实历史数据优化"""
    
    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
        
        # 更长周期的指标（避免频繁交易）
        ma20 = np.mean(c[-20:])
        ma60 = np.mean(c[-60:])
        
        # 20日涨跌幅
        ret20 = (c[-1] / c[-20] - 1)
        
        # 买入：长期趋势向上 + 中期趋势向上
        if c[-1] > ma20 > ma60 and ret20 > 0.05:  # 5%以上动量
            return 1.0
        # 卖出：跌破60日均线或大跌
        elif c[-1] < ma60 or ret20 < -0.08:
            return 0.0
        else:
            return self.pos
    
    def generate(self, data):
        new_pos = self.signal(data)
        curr = data['close'].iloc[-1]
        
        # 更宽松的止损（15%）
        if self.pos > 0:
            if curr > self.peak:
                self.peak = curr
            if curr < self.peak * 0.85:  # 15%止损
                new_pos = 0
        
        if new_pos > 0 and self.pos == 0:
            self.entry = curr
            self.peak = curr
            state = "BUY"
        elif new_pos == 0 and self.pos > 0:
            self.entry = 0
            self.peak = 0
            state = "SELL"
        else:
            state = "HOLD" if new_pos > 0 else "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)]
    
    results = []
    nav = 1.0
    
    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,
            'close': data['close'].iloc[i]
        })
    
    df = pd.DataFrame(results).set_index('date')
    df['idx_nav'] = df['close'] / df['close'].iloc[0]
    return df

def 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
    }

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.7, 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("="*70)
    print("创业板50 - 基于真实历史节点的回测")
    print("="*70)
    
    data = load_real_data()
    print(f"\n数据: {data.index[0].date()} ~ {data.index[-1].date()}")
    print(f"价格范围: {data['close'].min():.0f} ~ {data['close'].max():.0f}")
    
    # 训练
    print("\n【训练集 2018-2023】")
    s = HistoricalStrategy()
    train = backtest(data, s, '2018-01-01', '2023-12-31')
    m = metrics(train['nav'], train['idx_nav'])
    
    print(f"  策略收益: {m['total']*100:.1f}% (年化{m['annual']*100:.1f}%)")
    print(f"  指数收益: {m['idx_total']*100:.1f}% (年化{m['idx_annual']*100:.1f}%)")
    print(f"  超额: {m['excess']*100:.1f}%")
    print(f"  最大回撤: {m['max_dd']*100:.1f}%")
    print(f"  夏普: {m['sharpe']:.2f}")
    
    plot(train, "Training (2018-2023)", "train_historical.png")
    
    # 验证
    print("\n【验证集 2024-2025】")
    s2 = HistoricalStrategy()
    val = backtest(data, s2, '2024-01-01', '2025-12-31')
    m2 = metrics(val['nav'], val['idx_nav'])
    
    print(f"  策略收益: {m2['total']*100:.1f}% (年化{m2['annual']*100:.1f}%)")
    print(f"  指数收益: {m2['idx_total']*100:.1f}% (年化{m2['idx_annual']*100:.1f}%)")
    print(f"  超额: {m2['excess']*100:.1f}%")
    print(f"  最大回撤: {m2['max_dd']*100:.1f}%")
    
    plot(val, "Validation (2024-2025)", "val_historical.png")
    
    # 保存数据
    data.to_csv('cyb50_historical_data.csv')
    print("\n真实历史数据已保存: cyb50_historical_data.csv")
    
    print("\n" + "="*70)

if __name__ == "__main__":
    main()