import backtrader as bt
import pandas as pd
from datetime import datetime

class SmaCrossStrategy(bt.Strategy):
    """双均线交叉策略 - 创业板50示例"""
    params = (
        ('fast', 20),
        ('slow', 60),
        ('printlog', False),
    )
    
    def __init__(self):
        self.dataclose = self.datas[0].close
        self.order = None
        self.buyprice = None
        self.buycomm = None
        
        # 双均线
        self.sma_fast = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.fast)
        self.sma_slow = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.slow)
        
        # 交叉信号
        self.crossover = bt.indicators.CrossOver(self.sma_fast, self.sma_slow)
        
    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return
            
        if order.status in [order.Completed]:
            if order.isbuy():
                if self.params.printlog:
                    self.log(f'BUY EXECUTED, Price: {order.executed.price:.2f}, '
                            f'Cost: {order.executed.value:.2f}, '
                            f'Comm: {order.executed.comm:.2f}')
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
            else:
                if self.params.printlog:
                    self.log(f'SELL EXECUTED, Price: {order.executed.price:.2f}, '
                            f'Cost: {order.executed.value:.2f}, '
                            f'Comm: {order.executed.comm:.2f}')
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            if self.params.printlog:
                self.log('Order Canceled/Margin/Rejected')
                
        self.order = None
        
    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        if self.params.printlog:
            self.log(f'OPERATION PROFIT, GROSS: {trade.pnl:.2f}, NET: {trade.pnlcomm:.2f}')
            
    def next(self):
        if self.order:
            return
            
        # 金叉买入
        if self.crossover > 0:
            if not self.position:
                self.order = self.buy()
                if self.params.printlog:
                    self.log(f'BUY CREATE, {self.dataclose[0]:.2f}')
        # 死叉卖出
        elif self.crossover < 0:
            if self.position:
                self.order = self.sell()
                if self.params.printlog:
                    self.log(f'SELL CREATE, {self.dataclose[0]:.2f}')
                    
    def log(self, txt, dt=None):
        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.broker.startingcash:.2f}')
        print(f'最终资金: {self.broker.getvalue():.2f}')


def run_backtest(csv_file="chinext50.csv", cash=100000.0, commission=0.001):
    """运行回测"""
    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(SmaCrossStrategy, fast=20, slow=60, printlog=True)
    
    # 资金与手续费
    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')
    
    print(f'初始资金: {cerebro.broker.getvalue():.2f}')
    
    # 运行
    results = cerebro.run()
    strat = results[0]
    
    # 输出指标
    print(f'\n=== 回测指标 ===')
    print(f"年化收益: {strat.analyzers.returns.get_analysis()['rnorm100']:.2f}%")
    print(f"夏普比率: {strat.analyzers.sharpe.get_analysis()['sharperatio']:.3f}")
    print(f"最大回撤: {strat.analyzers.drawdown.get_analysis()['max']['drawdown']:.2f}%")
    
    return cerebro, strat


if __name__ == "__main__":
    # 先运行 fetch_data.py 获取数据
    cerebro, strat = run_backtest()
    # cerebro.plot()  # 如需图表，取消注释