"""创业板50双均线策略 - 最终工作版"""
import backtrader as bt
import pandas as pd
import numpy as np

class Chinext50Strategy(bt.Strategy):
    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)
        
    def next(self):
        if len(self) < 2:
            return
        fast_now, fast_prev = self.sma_fast[0], self.sma_fast[-1]
        slow_now, slow_prev = self.sma_slow[0], self.sma_slow[-1]
        if np.isnan(fast_now):
            return
        # 金叉买入
        if fast_prev <= slow_prev and fast_now > slow_now and not self.position:
            close = self.datas[0].close[0]
            size = (int(self.broker.getcash() / close) // 100) * 100
            if size > 0:
                self.buy(size=size)
        # 死叉卖出
        elif fast_prev >= slow_prev and fast_now < slow_now and self.position:
            self.close()

def run_strategy():
    cerebro = bt.Cerebro()
    df = pd.read_csv('chinext50.csv', parse_dates=['datetime'], index_col='datetime')
    cerebro.adddata(bt.feeds.PandasData(dataname=df))
    cerebro.addstrategy(Chinext50Strategy)
    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')
    print('初始资金: %.2f' % cerebro.broker.getvalue())
    results = cerebro.run()
    print('最终资金: %.2f' % cerebro.broker.getvalue())
    return cerebro, results[0]

if __name__ == '__main__':
    run_strategy()