"""
Working Regime Strategy - 基于已成功验证的cross检测逻辑
"""

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

class WorkingRegimeStrategy(bt.Strategy):
    params = (
        ('fast', 20),
        ('slow', 60),
        ('printlog', False),
    )
    
    def __init__(self):
        self.dataclose = self.datas[0].close
        self.order = None
        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 self.order:
            return
            
        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) or np.isnan(fast_prev):
            return
            
        # 金叉检测
        if fast_prev <= slow_prev and fast_now > slow_now:
            if not self.position:
                size = int(self.broker.getcash() / self.dataclose[0] / 100) * 100
                if size > 0:
                    self.order = self.buy(size=size)
                    self.trade_count += 1
                    if self.p.printlog:
                        self.log(f'BUY #{self.trade_count} @ {self.dataclose[0]:.2f}')
                        
        # 死叉检测
        elif fast_prev >= slow_prev and fast_now < slow_now:
            if self.position:
                self.order = self.close()
                if self.p.printlog:
                    self.log(f'SELL @ {self.dataclose[0]:.2f}')
    
    def notify_order(self, order):
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f'BUY EXECUTED @ {order.executed.price:.2f}')
            else:
                self.log(f'SELL EXECUTED @ {order.executed.price:.2f}')
        self.order = None
    
    def log(self, txt, dt=None):
        if self.p.printlog:
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()} {txt}')
    
    def stop(self):
        roi = (self.broker.getvalue() / self.broker.startingcash - 1) * 100
        print(f'\n=== 最终收益: {roi:.2f}% ===')
        print(f'总交易: {self.trade_count}')


def run_working_regime(csv_file="chinext50.csv"):
    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(WorkingRegimeStrategy, printlog=False)
    cerebro.broker.setcash(100000.0)
    cerebro.broker.setcommission(commission=0.001)
    
    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 Regime策略回测 ===')
    print(f'初始资金: {cerebro.broker.getvalue():.2f}')
    
    results = cerebro.run()
    strat = results[0]
    
    print(f'最终资金: {cerebro.broker.getvalue():.2f}')
    
    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%}")
    
    return cerebro, strat


if __name__ == "__main__":
    run_working_regime()