
import pandas as pd  
import backtrader as bt
import numpy as np
import pickle
import tushare as ts
token='2876ea85cb005fb5fa17c809a98174f2d5aae8b1f830110a5ead6211'
pro=ts.pro_api(token)

# 回测结果分析模块
from analyzer import analyzer

# 全局变量
class _a():
    pass
g = _a()
    
# 策略类
class strategy(bt.Strategy):
    def __init__(self):
        self.order_list = [] # 记录以往订单
        self.trades = [] # 记录成交成功的历史订单

    def log(self, txt, dt=None):
        ''' 策略日志打印函数'''
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))  # 这里是打印输出
        
    def next_open(self):
        # 在调仓之前，取消之前所下的没成交也未到期的订单
        if len(self.order_list) > 0:
            for od in self.order_list:
                self.cancel(od) # 如果订单未完成，则撤销订单
        self.order_list = [] #重置订单列表  

        dt = self.datas[0].datetime.date(0)       
        dt_str = dt.strftime('%Y%m%d')        
        # 今天需要持有的ETF
        stk_tohold = [g.etf_tohold[dt_str]]
        # 今日持有的ETF
        stk_hold = [d._name for d, pos in self.getpositions().items() if pos]
        # 今日需要买入的ETF
        stk_buy = [stk for stk in stk_tohold if stk not in stk_hold]
        # 今日需要卖出的ETF
        stk_sell = [stk for stk in stk_hold if stk not in stk_tohold]
        
        # 先卖后买
        if len(stk_sell)>0:
            for stk in stk_sell:
                data = self.getdatabyname(stk)
                od = self.close(data=data) 
                self.order_list.append(od)
        
        if len(stk_buy)>0:
            for stk in stk_buy:
                data = self.getdatabyname(stk)
                volume_tobuy = (self.broker.getvalue()*0.98/data.open[0]//100)*100
                od = self.order_target_size(data=data, target=volume_tobuy) 
                self.order_list.append(od)
            
        # 总结每天的持仓
        stk_hold = [d._name for d, pos in self.getpositions().items() if pos]
        stk_holdvolume = [pos.size for d, pos in self.getpositions().items() if pos]  
        if stk_hold: 
            g.df.loc[dt_str]['hold'] = stk_hold[0]
            g.df.loc[dt_str]['volume'] = stk_holdvolume[0]
            g.df.loc[dt_str]['close'] = self.getdatabyname(stk_hold[0]).close[0]
        g.df.loc[dt_str]['cash'] = self.broker.getcash()
        g.df.loc[dt_str]['value'] = self.broker.getvalue()
        
    def notify_trade(self, trade):
        if trade.isclosed:
            self.trades.append(trade.pnl)  # 记录盈亏金额
            
    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(
                        '买单完成(BUY), 订单编号: %.0f, %s, 成交价:%.3f, 成交额:%.1f, 佣金 %.1f, 成交量:%.0f' %
                        (order.ref,            # 订单编号
                         order.data._name,     # 股票名称
                         order.executed.price, # 成交价
                         order.executed.value, # 成交额
                         order.executed.comm, # 佣金
                         order.executed.size # 成交量
                         )) 
            elif order.issell():
                self.log('卖单完成(SELL), 订单编号: %.0f, %s, 成交价:%.3f, 成交额:%.1f, 佣金 %.1f, 成交量:%.0f' %
                            (order.ref,
                             order.data._name,
                             order.executed.price,
                             order.executed.value,
                             order.executed.comm,
                             order.executed.size
                             ))
        # 未成交的订单
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected, Ref: %.0f, %s, Price:%.3f, Cost:%.1f, Comm %.1f, Size:%.0f' % 
                     (order.ref,
                      order.data._name, 
                      order.executed.price, 
                      order.executed.value,
                      order.executed.comm, 
                      order.executed.size))


# if __name__ == '__main__':
    
# 回测起止日期，这里的起止日期也可以根据需要修改
start_date, end_date = '20180207', '20260517'

# 回测初始资金
startcash = 100000
# 交易日
trade_day = pro.trade_cal(exchange='', start_date=start_date, end_date=end_date)
trade_day = trade_day[trade_day['is_open'] == 1]
g.trade_day = trade_day.sort_values(by='cal_date', ascending=True)

# 初始化变量，用于储存每日持仓信息
g.df = pd.DataFrame(index=g.trade_day['cal_date'], columns = ['hold','volume','close','cash', 'value'])
g.etf_tohold = pd.read_pickle('data_ts/etf_tohold.pkl')
    
# 初始化cerebro回测系统设置  
cerebro = bt.Cerebro(cheat_on_open=True) 

#将行情数据传入回测系统
with open('data_ts/daily_data_ETF_bt.pkl', "rb") as file:
    prices = pickle.load(file)   
etf_libs = list(prices.keys())
for code in etf_libs:
    data = prices[code]
    data = data.reindex(index=pd.to_datetime(g.trade_day['cal_date']))
    data = data[['open', 'high', 'low', 'close']].ffill()
    datafeed = bt.feeds.PandasData(dataname=data, 
                                fromdate=data.index[0],
                                todate=data.index[-1])
    cerebro.adddata(datafeed, name=code)  
print('adddata completed!')

# 设置初始资金
cerebro.broker.setcash(startcash) 
# 设置交易手续费（佣金）为万1，ETF交易无印花税、过户费等。
cerebro.broker.setcommission(commission=0.0001) 
# 滑点：暂设为万1
cerebro.broker.set_slippage_perc(perc=0.0001)

# 将交易策略加载到回测系统中
cerebro.addstrategy(strategy) 
cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='_TimeReturn')
result = cerebro.run()

# 提取收益序列
pnl = pd.Series(result[0].analyzers._TimeReturn.get_analysis())

# 基准指数收益
Ashare = pro.index_daily(ts_code = '000300.SH', start_date = start_date, end_date = end_date)
Ashare.index = pd.to_datetime(Ashare.trade_date)
Ashare = Ashare.sort_index(axis=0)
Ashare = Ashare['pct_chg']/100

# 分析结果
stats_all = analyzer(pnl, Ashare, result[0])
g.df.to_csv('回测结果汇总.csv')

    

