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)

# 回测参数
INIT_CASH = 4000000
FEE_RATE = 0.003  # 千分之3
GRID_NUM = 15  # 网格数量
PER_GRID_CASH = INIT_CASH / GRID_NUM  # 每格买入金额

def generate_trend_following_grids(price, ma20):
    """生成趋势跟随网格
    上升趋势：
    - 买入：每跌1%加仓
    - 卖出：每涨3%减仓
    下降趋势：
    - 买入：每跌3%加仓
    - 卖出：每涨1%减仓
    """
    if price > ma20:  # 上升趋势
        buy_steps = np.arange(price * 0.99, price * 0.85, -price * 0.01)  # 每跌1%加仓
        sell_steps = np.arange(price * 1.03, price * 1.30, price * 0.03)  # 每涨3%减仓
    else:  # 下降趋势
        buy_steps = np.arange(price * 0.97, price * 0.70, -price * 0.03)  # 每跌3%加仓
        sell_steps = np.arange(price * 1.01, price * 1.15, price * 0.01)  # 每涨1%减仓
    return np.sort(np.concatenate([buy_steps, sell_steps]))

def trend_following_grid_backtest(prices):
    cash = INIT_CASH
    position = 0
    cost = 0
    history = []  # 记录每日资金、持仓
    trade_log = []  # 买卖点
    last_grid = None
    
    # 计算MA20
    prices['MA20'] = prices['close'].rolling(window=20).mean()
    
    def get_grid_index(price, current_grids):
        """获取价格所在的网格索引"""
        if len(current_grids) == 0:
            return -1
        grid_idx = np.searchsorted(current_grids, price, side='right') - 1
        if grid_idx < 0:
            grid_idx = 0
        if grid_idx >= len(current_grids)-1:
            grid_idx = len(current_grids)-2
        return grid_idx
    
    def execute_trade(price, trade_type, grid_idx, last_grid_idx, date):
        """执行交易"""
        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():
        if pd.isna(row['MA20']):  # 跳过MA20未形成的时期
            continue
            
        date = row['date']
        open_price = row['open']
        close_price = row['close']
        ma20 = row['MA20']
        
        # 生成当前的趋势跟随网格
        open_grids = generate_trend_following_grids(open_price, ma20)
        close_grids = generate_trend_following_grids(close_price, ma20)
        
        # 获取开盘价和收盘价所在的网格位置
        open_grid_idx = get_grid_index(open_price, open_grids)
        close_grid_idx = get_grid_index(close_price, close_grids)
        
        # 开盘价交易处理
        if last_grid is not None and open_grid_idx >= 0:
            # 开盘价买入处理
            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, date):
                        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, date):
                        break
        
        # 收盘价交易处理
        if last_grid is not None and close_grid_idx >= 0:
            # 收盘价买入处理
            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, date):
                        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, date):
                        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,
            'ma20': ma20,
            '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

def plot_results(df_hist, trade_log, grids=None):
    """绘制回测结果"""
    fig, ax1 = plt.subplots(figsize=(14,6))
    ax2 = ax1.twinx()
    
    # 绘制价格、MA20和净值
    ax1.plot(df_hist['date'], df_hist['price'], label='收盘价', color='gray', alpha=0.7)
    ax1.plot(df_hist['date'], df_hist['ma20'], label='MA20', color='purple', alpha=0.5)
    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 "")
    
    # 绘制最后一个时间点的网格线
    if grids is not None and len(grids) > 0:
        for g in grids:
            ax1.axhline(g, color='orange', linestyle='--', alpha=0.2)
    
    plt.title('趋势跟随网格回测')
    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()

if __name__ == '__main__':
    # 读取数据
    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['open'] = raw['open'].astype(float)
    raw = raw.sort_values('date').reset_index(drop=True)
    
    # 运行回测
    df_hist, trade_log, total_return, max_drawdown, buy_cnt, sell_cnt, win_rate = trend_following_grid_backtest(raw)
    
    # 打印统计结果
    print('\n==== 趋势跟随网格回测结果 ===')
    print(f'总收益率: {total_return*100:.2f}%')
    print(f'最大回撤: {max_drawdown*100:.2f}%')
    print(f'买入次数: {buy_cnt}，卖出次数: {sell_cnt}')
    print(f'胜率: {win_rate*100:.2f}%')
    
    # 绘制结果
    last_price = df_hist['price'].iloc[-1]
    last_ma20 = df_hist['ma20'].iloc[-1]
    last_grids = generate_trend_following_grids(last_price, last_ma20)
    plot_results(df_hist, trade_log, last_grids) 