# 加载模块
from datetime import datetime

from tqdm import tqdm
import rqdatac as rq

from vnpy.trader.database import DB_TZ
from vnpy.trader.datafeed import get_datafeed
from vnpy.trader.constant import Exchange, Interval
from vnpy.trader.object import HistoryRequest
from vnpy.alpha import AlphaLab, logger

# 设置下载参数
task_name = "csi300"
index_symbol = "000300.SSE"
rq_index_symbol = "000300.XSHG"

start_date = "2007-01-01"
end_date = "2024-10-31"

# 创建投研实验室
lab = AlphaLab(f"./lab/{task_name}")    # 指定数据文件夹

# 初始化数据服务（这里配置使用的RQData）
datafeed = get_datafeed()
datafeed.init()

# 下载指数成分股
data = rq.index_components(rq_index_symbol, start_date=start_date, end_date=end_date)

# 转换合约代码
index_components = {}
for dt, rq_symbols in data.items():
    vt_symbols: list = []

    for rq_symbol in rq_symbols:
        vt_symbol = rq_symbol.replace("XSHG", "SSE").replace("XSHE", "SZSE")
        vt_symbols.append(vt_symbol)

    index_components[dt.strftime("%Y-%m-%d")] = vt_symbols

# 保存到数据中心
lab.save_component_data(index_symbol, index_components)

# 加载指数成分股代码
component_symbols = lab.load_component_symbols(index_symbol, start_date, end_date)

# 转换时间格式
start = datetime.strptime(start_date, "%Y-%m-%d")
start = start.replace(tzinfo=DB_TZ)

end = datetime.strptime(end_date, "%Y-%m-%d")
end = end.replace(tzinfo=DB_TZ)

# 除了成分股，还要下载指数数据
task_symbols = component_symbols + [index_symbol]

# 遍历下载数据
for vt_symbol in tqdm(task_symbols):
    symbol, exchange_str = vt_symbol.split(".")

    req = HistoryRequest(symbol, Exchange(exchange_str), start, end, Interval.DAILY)
    bars = datafeed.query_bar_history(req)

    if bars:
        lab.save_bar_data(bars)
    else:
        logger.error(f"下载{vt_symbol}数据失败")

# 添加回测参数配置
for vt_symbol in component_symbols:
    lab.add_contract_setting(
        vt_symbol,
        long_rate=5/10000,
        short_rate=10/10000,
        size=1,
        pricetick=0.0001,
    )