import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os

# 新增：设置matplotlib支持中文和负号
import matplotlib
matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']
matplotlib.rcParams['axes.unicode_minus'] = False

# 新增：屏蔽UserWarning
import warnings
warnings.filterwarnings("ignore", category=UserWarning)

# 读取数据
DATA_PATH = os.path.join(os.path.dirname(__file__), '0020.HK.csv')
# 处理文件头乱码和尾部注释
raw = pd.read_csv(DATA_PATH, header=None, skiprows=1, names=['code','date','open','high','low','close','pct'])
raw = raw[raw['close'].apply(lambda x: str(x).replace('.','',1).isdigit())]  # 去除空行和注释
raw['date'] = pd.to_datetime(raw['date'])
raw['close'] = raw['close'].astype(float)
raw = raw.sort_values('date').reset_index(drop=True)

# 回测参数
INIT_CASH = 4000000
FEE_RATE = 0.003  # 千分之3
GRID_TYPES = ['等差网格', '等比网格', 'ATR网格']
GRID_NUM = 15  # 网格数量
PER_GRID_CASH = INIT_CASH / GRID_NUM  # 每格买入金额

# 自动确定网格上下限
price_min = raw['close'].min()
price_max = raw['close'].max()
price_mean = raw['close'].mean()
price_std = raw['close'].std()

# 网格参数生成
# 1. 等差网格
step_linear = (price_max - price_min) / GRID_NUM
linear_grids = np.arange(price_min, price_max, step_linear)
# 2. 等比网格
step_ratio = (price_max / price_min) ** (1/GRID_NUM)
ratio_grids = [price_min * (step_ratio ** i) for i in range(GRID_NUM+1)]
# 3. ATR网格
atr_window = 20
raw['tr'] = np.maximum(raw['high']-raw['low'], np.abs(raw['high']-raw['close'].shift(1)), np.abs(raw['low']-raw['close'].shift(1)))
raw['atr'] = raw['tr'].rolling(atr_window).mean()
avg_atr = raw['atr'].mean()
atr_step = avg_atr * 1.0  # 可调参数
atr_grids = np.arange(price_min, price_max, atr_step)

# 网格集合
grid_dict = {
    '等差网格': linear_grids,
    '等比网格': ratio_grids,
    'ATR网格': atr_grids
}

# 回测主函数
def grid_backtest(prices, grids, grid_type):
    cash = INIT_CASH
    position = 0
    cost = 0
    history = []  # 记录每日资金、持仓
    trade_log = []  # 买卖点
    last_grid = None
    
    def get_grid_index(price):
        """获取价格所在的网格索引"""
        grid_idx = np.searchsorted(grids, price, side='right') - 1
        if grid_idx < 0:
            grid_idx = 0
        if grid_idx >= len(grids)-1:
            grid_idx = len(grids)-2
        return grid_idx
    
    def execute_trade(price, trade_type, grid_idx, last_grid_idx):
        """执行交易"""
        nonlocal cash, position, cost
        if trade_type == 'buy' and cash >= PER_GRID_CASH:
            buy_price = price * (1 + FEE_RATE)
            buy_amount = PER_GRID_CASH / buy_price
            cash -= buy_amount * buy_price
            position += buy_amount
            cost += buy_amount * buy_price
            trade_log.append({'date': date, 'price': price, 'type': 'buy'})
            return True
        elif trade_type == 'sell' and position > 0:
            sell_price = price * (1 - FEE_RATE)
            sell_amount = PER_GRID_CASH / price
            if sell_amount > position:
                sell_amount = position
            cash += sell_amount * sell_price
            position -= sell_amount
            cost -= sell_amount * price
            trade_log.append({'date': date, 'price': price, 'type': 'sell'})
            return True
        return False

    for i, row in prices.iterrows():
        date = row['date']
        open_price = row['open']
        close_price = row['close']
        
        # 获取开盘价和收盘价所在的网格位置
        open_grid_idx = get_grid_index(open_price)
        close_grid_idx = get_grid_index(close_price)
        
        # 开盘价交易处理
        if last_grid is not None:
            # 开盘价买入处理
            if open_grid_idx < last_grid:
                # 处理多个网格的买入
                for grid_idx in range(open_grid_idx, last_grid):
                    if not execute_trade(open_price, 'buy', grid_idx, last_grid):
                        break
            # 开盘价卖出处理
            elif open_grid_idx > last_grid:
                # 处理多个网格的卖出
                for grid_idx in range(last_grid + 1, open_grid_idx + 1):
                    if not execute_trade(open_price, 'sell', grid_idx, last_grid):
                        break
        
        # 收盘价交易处理
        if last_grid is not None:
            # 收盘价买入处理
            if close_grid_idx < open_grid_idx:
                # 处理多个网格的买入
                for grid_idx in range(close_grid_idx, open_grid_idx):
                    if not execute_trade(close_price, 'buy', grid_idx, open_grid_idx):
                        break
            # 收盘价卖出处理
            elif close_grid_idx > open_grid_idx:
                # 处理多个网格的卖出
                for grid_idx in range(open_grid_idx + 1, close_grid_idx + 1):
                    if not execute_trade(close_price, 'sell', grid_idx, open_grid_idx):
                        break
        
        # 更新last_grid为收盘价所在网格
        last_grid = close_grid_idx
        
        # 记录每日净值
        total = cash + position * close_price
        history.append({
            'date': date, 
            'net': total, 
            'cash': cash, 
            'position': position, 
            'price': close_price,
            'open_grid': open_grid_idx,
            'close_grid': close_grid_idx
        })
    
    # 统计
    df_hist = pd.DataFrame(history)
    max_drawdown = ((df_hist['net'].cummax() - df_hist['net']) / df_hist['net'].cummax()).max()
    total_return = (df_hist['net'].iloc[-1] - INIT_CASH) / INIT_CASH
    
    # 胜率统计
    win, lose = 0, 0
    for i in range(1, len(trade_log)):
        if trade_log[i-1]['type']=='buy' and trade_log[i]['type']=='sell':
            if trade_log[i]['price'] > trade_log[i-1]['price']:
                win += 1
            else:
                lose += 1
    win_rate = win / (win+lose) if (win+lose)>0 else 0
    
    return df_hist, trade_log, total_return, max_drawdown, len([t for t in trade_log if t['type']=='buy']), len([t for t in trade_log if t['type']=='sell']), win_rate

# 主流程
if __name__ == '__main__':
    for grid_type in GRID_TYPES:
        grids = grid_dict[grid_type]
        df_hist, trade_log, total_return, max_drawdown, buy_cnt, sell_cnt, win_rate = grid_backtest(raw, grids, grid_type)
        print(f'\n==== {grid_type} 回测结果 ===')
        print(f'总收益率: {total_return*100:.2f}%')
        print(f'最大回撤: {max_drawdown*100:.2f}%')
        print(f'买入次数: {buy_cnt}，卖出次数: {sell_cnt}')
        print(f'胜率: {win_rate*100:.2f}%')
        # 绘图
        fig, ax1 = plt.subplots(figsize=(14,6))
        ax2 = ax1.twinx()
        ax1.plot(df_hist['date'], df_hist['price'], label='收盘价', color='gray', alpha=0.7)
        # 将净值转换为相对净值（从1开始）
        relative_net = df_hist['net'] / INIT_CASH
        ax2.plot(df_hist['date'], relative_net, label='净值', color='blue')
        for t in trade_log:
            if t['type']=='buy':
                ax1.scatter(t['date'], t['price'], marker='^', color='red', label='买入' if '买入' not in ax1.get_legend_handles_labels()[1] else "")
            else:
                ax1.scatter(t['date'], t['price'], marker='v', color='green', label='卖出' if '卖出' not in ax1.get_legend_handles_labels()[1] else "")
        for g in grids:
            ax1.axhline(g, color='orange', linestyle='--', alpha=0.2)
        plt.title(f'{grid_type} 网格回测')
        lines1, labels1 = ax1.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax1.legend(lines1 + lines2, labels1 + labels2)
        plt.tight_layout()
        plt.show() 