"""
简化的Regime策略 - 基于趋势强度动态调整仓位
"""

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

class SimpleRegimeStrategy(bt.Strategy):
    """
    简化版状态感知策略
    
    逻辑:
    - 计算20日趋势强度
    - 趋势强且向上: 满仓
    - 趋势弱或向下: 减仓/空仓
    - 均线金叉入场,死叉出场
    """
    
    params = (
        ('fast', 20),
        ('slow', 60),
        ('trend_threshold', 0.02),
        ('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.crossover = bt.indicators.CrossOver(self.sma_fast, self.sma_slow)
        
        # 趋势: 价格相对快均线的偏离
        self.trend = (self.dataclose - self.sma_fast) / self.sma_fast
        
    def next(self):
        if self.order:
            return
            
        # 当前趋势
        trend_val = self.trend[0] if not np.isnan(self.trend[0]) else 0
        
        # 金叉 + 趋势向上 = 买入
        if self.crossover > 0 and trend_val > -self.p.trend_threshold:
            if not self.position:
                size = int(self.broker.getcash() / self.dataclose[0] / 100) * 100
                if size > 0:
                    self.order = self.buy(size=size)
                    if self.p.printlog:
                        self.log(f'BUY @ {self.dataclose[0]:.2f}, Trend: {trend_val:.4f}')
                        
        # 死叉 或 趋势转弱 = 卖出
        elif self.crossover < 0 or (self.position and trend_val < -self.p.trend_threshold * 2):
            if self.position:
                self.order = self.close()
                if self.p.printlog:
                    reason = 'Death Cross' if self.crossover < 0 else 'Weak Trend'
                    self.log(f'SELL @ {self.dataclose[0]:.2f}, Reason: {reason}')
    
    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return
        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 not self.p.printlog:
            return
        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}% ===')


def run_simple_regime(csv_file="chinext50.csv", cash=100000.0):
    """运行简化Regime策略"""
    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(SimpleRegimeStrategy, printlog=False)
    cerebro.broker.setcash(cash)
    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('=== 简化Regime策略回测 ===')
    print(f'初始资金: {cerebro.broker.getvalue():.2f}')
    
    results = cerebro.run()
    strat = results[0]
    
    print(f'\n最终资金: {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.get('total'):
        total_trades = trades['total'].get('total', 0)
        won_trades = trades['won'].get('total', 0) if trades.get('won') else 0
        print(f"总交易次数: {total_trades}")
        print(f"盈利次数: {won_trades}")
        if total_trades > 0:
            print(f"胜率: {won_trades/total_trades:.1%}")
    
    return cerebro, strat


if __name__ == "__main__":
    run_simple_regime()