
import pandas as pd
import pymysql
import backtrader as bt
import config

# 从 MySQL 提取数据
def fetch_data_from_mysql():
    connection = pymysql.connect(**config.db_config)
    #    ['datetime', 'open', 'high', 'low', 'close', 'volume', 'openinterest']]
    #query = "SELECT code, time_key, qfq_close as close FROM hs300_ajust_kline_202410 where time_key>='2023-01-01 00:00:00' "
    query = "SELECT code, time_key, open, high, low, close, volume FROM hs300_his_kline_none where time_key>='2023-01-01 00:00:00' "
    data = pd.read_sql(query, connection)
    data['time_key'] = pd.to_datetime(data['time_key'], errors='coerce')  # 确保为datetime格式
    data['openinterest'] = 0  # 添加 openinterest 列并赋值为 0
    connection.close()
    return data

# 格式化数据为Backtrader兼容的格式
data = fetch_data_from_mysql()
data_by_code = {code: df for code, df in data.groupby('code')}


class BestPortfolioStrategy(bt.Strategy):
    params = dict(
        max_stocks=30,
        rebalance_monthly=True,
        commission=0.0003,
        slippage=0.005,
        maperiod = 15,
        printlog = False
    )    
    def __init__(self):
        self.stocks = []  # 持有的股票
        self.rebalance_date = None

    def next(self):
        # 判断是否是月初
        if self.rebalance_date and self.data.datetime.date(0) < self.rebalance_date.date():
            return

        print(f"rebalance date: {self.rebalance_date}")
        # 设置下次调仓日
        self.rebalance_date = self.data.datetime.date(0).replace(day=1) + pd.DateOffset(months=1)
        
        # 选股逻辑：计算过去一段时间的收益率，选取最高的30只股票
        returns = {}
        for data in self.datas:
            returns[data._name] = (data.close[0] / data.close[-20]) - 1  # 上个月的收益率

        # 按收益率排序并选择最佳组合
        sorted_stocks = sorted(returns, key=returns.get, reverse=True)[:self.params.max_stocks]
        print(sorted_stocks)
        
        # 调仓：卖出非选中股票，买入新选股票
        self.rebalance(sorted_stocks)
    
    def rebalance(self, selected_stocks):
        for stock in self.stocks:
            if stock not in selected_stocks:
                self.close(self.getdatabyname(stock))

        for stock in selected_stocks:
            if stock not in self.stocks:
                self.buy(self.getdatabyname(stock), size=100)

        self.stocks = selected_stocks


    # 交易记录日志（可省略，默认不输出结果）

    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()},{txt}')

    def notify_order(self, order):
        # 未被处理的订单
        if order.status in [order.Submitted, order.Accepted]:
            return
        # 已经处理的订单
        if order.status in [order.Completed, order.Canceled, order.Margin]:
            if order.isbuy():
                self.log(
                    'BUY EXECUTED, ref:%.0f，Price: %.2f, Cost: %.2f, Comm %.2f, Size: %.2f, Stock: %s' %
                    (order.ref,  # 订单编号
                     order.executed.price,  # 成交价
                     order.executed.value,  # 成交额
                     order.executed.comm,  # 佣金
                     order.executed.size,  # 成交量
                     order.data._name))  # 股票名称
            else:  # Sell
                self.log('SELL EXECUTED, ref:%.0f, Price: %.2f, Cost: %.2f, Comm %.2f, Size: %.2f, Stock: %s' %
                         (order.ref,
                          order.executed.price,
                          order.executed.value,
                          order.executed.comm,
                          order.executed.size,
                          order.data._name))


cerebro = bt.Cerebro()
cerebro.broker.setcash(1_000_000)  # 初始资金 100万
cerebro.broker.setcommission(commission=0.0003)  # 设置交易费率 万分之三
cerebro.broker.set_slippage_perc(0.005)  # 设置滑点 0.005

# 加载数据
for code, df in data_by_code.items():
    # 删除缺失值
    df = df.dropna(subset=['time_key'])

    # 设置DataFrame为Backtrader的数据格式
    data_feed = bt.feeds.PandasData(
        dataname=df,
        datetime='time_key',
        openinterest=None  # 如果没有openinterest列，则为None
    )
    #bt_data = bt.feeds.PandasData(dataname=df)
    cerebro.adddata(data_feed, name=code)


cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='pnl')  # 返回收益率时序数据
cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='_AnnualReturn')  # 年化收益率
cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='_SharpeRatio')  # 夏普比率
cerebro.addanalyzer(bt.analyzers.DrawDown, _name='_DrawDown')  # 回撤

# 加载策略
cerebro.addstrategy(BestPortfolioStrategy, printlog=True)
# 运行回测
results = cerebro.run()


# 获取分析结果
results = cerebro.run()
strat = results[0]

# 返回日度收益率序列
daily_return = pd.Series(strat.analyzers.pnl.get_analysis())
# 打印评价指标
print("--------------- AnnualReturn -----------------")
print(strat.analyzers._AnnualReturn.get_analysis())
print("--------------- SharpeRatio -----------------")
print(strat.analyzers._SharpeRatio.get_analysis())
print("--------------- DrawDown -----------------")
print(strat.analyzers._DrawDown.get_analysis())
