#!/usr/bin/env python3
"""
创业板50指数Regime策略 - 最终工作版
基于双均线(20/60)的金叉/死叉信号
"""

import backtrader as bt
import pandas as pd
import numpy as np


class Chinext50RegimeStrategy(bt.Strategy):
    """
    创业板50双均线Regime策略
    
    参数:
        fast: 短期均线周期 (默认20)
        slow: 长期均线周期 (默认60)
    """
    params = (
        ('fast', 20),
        ('slow', 60),
    )
    
    def __init__(self):
        self.sma_fast = bt.indicators.SMA(period=self.p.fast)
        self.sma_slow = bt.indicators.SMA(period=self.p.slow)
        self.trade_count = 0
        
    def next(self):
        """策略主逻辑 - 每个交易日调用"""
        if len(self) < 2:
            return
            
        # 获取当前和前一周期的均线值
        fast_now = self.sma_fast[0]
        fast_prev = self.sma_fast[-1]
        slow_now = self.sma_slow[0]
        slow_prev = self.sma_slow[-1]
        
        # 检查有效值
        if np.isnan(fast_now):
            return
        
        # 金叉检测: 前一周fast<=slow, 这一周fast>slow
        golden_cross = fast_prev <= slow_prev and fast_now > slow_now
        
        # 死叉检测: 前一周fast>=slow, 这一周fast<slow
        death_cross = fast_prev >= slow_prev and fast_now < slow_now
        
        # 买入逻辑: 金叉 + 空仓
        if golden_cross and not self.position:
            close_price = self.datas[0].close[0]
            size = int(self.broker.getcash() / close_price)
            if size > 0:
                self.buy(size=size)
                self.trade_count += 1
                
        # 卖出逻辑: 死叉 + 有持仓
        elif death_cross and self.position:
            self.close()


def run_backtest(csv_file="chinext50.csv", cash=100000.0, commission=0.001):
    """
    运行回测
    
    参数:
        csv_file: 数据文件路径
        cash: 初始资金
        commission: 手续费率
    """
    cerebro = bt.Cerebro()
    
    # 加载数据
    df = pd.read_csv(csv_file, parse_dates=['datetime'], index_col='datetime')
    data = bt.feeds.PandasData(dataname=df)
    cerebro.adddata(data)
    
    # 添加策略
    cerebro.addstrategy(Chinext50RegimeStrategy)
    
    # 设置资金
    cerebro.broker.setcash(cash)
    cerebro.broker.setcommission(commission=commission)
    
    # 添加分析器
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='sharpe', riskfreerate=0.02)
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='drawdown')
    cerebro.addanalyzer(bt.analyzers.Returns, _name='returns')
    cerebro.addanalyzer(bt.analyzers.TradeAnalyzer, _name='trades')
    
    print('=' * 50)
    print('创业板50 Regime策略回测')
    print('=' * 50)
    print(f'初始资金: {cerebro.broker.getvalue():.2f}')
    
    # 运行回测
    results = cerebro.run()
    strat = results[0]
    
    # 输出结果
    print(f'最终资金: {cerebro.broker.getvalue():.2f}')
    print(f'交易次数: {strat.trade_count}')
    
    returns = strat.analyzers.returns.get_analysis()
    print(f"年化收益: {returns.get('rnorm100', 0):.2f}%")
    
    sharpe = strat.analyzers.sharpe.get_analysis()
    sharpe_val = sharpe.get('sharperatio', 0)
    print(f"夏普比率: {sharpe_val:.3f}" if sharpe_val else "夏普比率: N/A")
    
    drawdown = strat.analyzers.drawdown.get_analysis()
    print(f"最大回撤: {drawdown.get('max', {}).get('drawdown', 0):.2f}%")
    
    trades = strat.analyzers.trades.get_analysis()
    if trades and trades.get('total'):
        total = trades['total'].get('total', 0)
        won = trades['won'].get('total', 0) if trades.get('won') else 0
        print(f'总交易: {total}, 盈利: {won}')
        if total > 0:
            print(f'胜率: {won/total:.1%}')
    
    print('=' * 50)
    return cerebro, strat


if __name__ == "__main__":
    run_backtest()