from turtle import width
from plotly.io import renderers
import vectorbt as vbt
import pandas as pd

# ========================
# 数据加载
# ========================
# 从本地读取 ETH/USDT 1天K线数据（来自OKX）
df = pd.read_feather("/Users/aszer/Documents/vscode/cta/user_data/data/okx/ADA_USDT-1d.feather")
# df = df[df['date'] >= '2025-07-01']
df.set_index("date", inplace=True)
# 提取收盘价作为交易价格
price = df["close"]
open = df["open"]
# ========================
# 技术指标计算
# ========================
# 计算10日指数移动平均线（EMA）
ema_10 = vbt.MA.run(price, window=10, ewm=True).ma

# 计算20日指数移动平均线（EMA）
ema_20 = vbt.MA.run(price, window=20, ewm=True).ma

# 计算6周期相对强弱指数（RSI）
rsi_6 = vbt.RSI.run(price, window=6).rsi

# 计算12周期相对强弱指数（RSI）
rsi_12 = vbt.RSI.run(price, window=12).rsi

# ========================
# 买入条件定义
# ========================
# 条件1：EMA多头排列 + 短期趋势加速
# - EMA10 > EMA20：短期趋势强于长期趋势（多头排列）
# - EMA10昨日 > 前日：短期均线继续上行，显示动量增强
ema_bullish_cross = (ema_10 > ema_20) & (ema_10.shift(1) > ema_10.shift(2))

# 条件2：双周期RSI进入深度超卖区，暗示反弹可能
# - 过去3天（含前天）的6周期RSI最高值 < 21，表示近期极度超卖
# - 过去3天的12周期RSI最高值 < 26.25，确认中周期也处于超卖状态
# 注：此条件未显式检查“今日反弹”，可后续增强
rsi_oversold_bounce = (rsi_6.shift(1).rolling(3).max() < 21) & \
                      (rsi_12.shift(1).rolling(3).max() < 26.25)

# 合并所有买入信号：任一条件满足即产生买入信号
entries = ema_bullish_cross

# ========================
# 卖出条件定义
# ========================
# 计算长期均线与短期均线的价差（空头趋势强度）
diff = ema_20 - ema_10  # 当diff扩大，表示空头趋势加强

# 条件1：EMA10连续5天下跌（动量走弱）
# - 当前EMA10 < 过去5天（不含今日）的最低EMA10值
ema10_falling = ema_10 < ema_10.shift(1).rolling(window=5).min()

# 条件2：空头趋势加速（价差达到近期最大）
# - 当前diff是过去6天（含今日）中的最大值，表示空头力量最强
diff_expanding = diff == diff.rolling(window=6).max()

# 合并卖出信号：两个条件同时满足才卖出
sell_cond = ema10_falling & diff_expanding
exits = sell_cond

# ========================
# 回测执行
# ========================
# 使用 vectorbt 的信号回测引擎，构建投资组合
pf = vbt.Portfolio.from_signals(
    price,                  # 价格序列
    entries.vbt.fshift(1),                # 买入信号
    exits.vbt.fshift(1),                  # 卖出信号
    price=open,
    init_cash=100,          # 初始资金 100 USDT
    fees=0.001,             # 交易手续费 0.1%（买卖均收）
    sl_stop=0.05,           # 止损：从最高价回撤5%时触发
    freq="1D"               # 数据频率为每日，用于复利和年化计算
)

# ========================
# 输出回测统计结果
# ========================
print(pf.stats())
# fig = pf.plot_orders(width=1200, height=800)  
# fig.show(renderer='browser')  
orders_df = pf.orders.records_readable  
print(orders_df)