prepare_cn_data.py

from __future__ import annotations

from datetime import datetime
from pathlib import Path
import time
from typing import Iterable

import baostock as bs
import pandas as pd
from dateutil.relativedelta import relativedelta
from tqdm import tqdm

from harness_exp.config import get_baseline_backtest_spec

DATASET_LOCK = get_baseline_backtest_spec()
RAW_DATA_DIR = Path("~/.qlib/qlib_data/cn_data/raw_data_back_adjust").expanduser()
UNIVERSE_CACHE_DIR = Path("~/.qlib/qlib_data/cn_data").expanduser()
TARGET_UNIVERSE = "CSI500"
BENCHMARK_CODE = "sh.000905"
REQUEST_RETRY_LIMIT = 5
RETRY_SLEEP_SECONDS = 2.0


def _iter_month_starts(start_date: str, end_date: str) -> Iterable[str]:
    current = datetime.strptime(start_date, "%Y-%m-%d").replace(day=1)
    end = datetime.strptime(end_date, "%Y-%m-%d")
    while current <= end:
        yield current.strftime("%Y-%m-%d")
        current += relativedelta(months=1)


def _query_to_dataframe(query_result) -> pd.DataFrame:
    rows: list[list[str]] = []
    while query_result.error_code == "0" and query_result.next():
        rows.append(query_result.get_row_data())
    return pd.DataFrame(rows, columns=query_result.fields)


def get_csi500_stock_union(start_date: str, end_date: str) -> list[str]:
    cache_file = UNIVERSE_CACHE_DIR / f"csi500_union_{DATASET_LOCK.dataset_version}_{start_date}_{end_date}.txt"
    if cache_file.exists():
        cached_codes = [line.strip() for line in cache_file.read_text(encoding="utf-8").splitlines() if line.strip()]
        if cached_codes:
            print(f"使用缓存的CSI500成分并集: {cache_file}", flush=True)
            return cached_codes

    codes: set[str] = set()
    month_starts = list(_iter_month_starts(start_date, end_date))
    for query_date in tqdm(month_starts, desc="收集CSI500成分", unit="month"):
        rs = bs.query_zz500_stocks(query_date)
        if rs.error_code != "0":
            raise RuntimeError(f"Failed to query CSI500 members on {query_date}: {rs.error_msg}")
        stock_df = _query_to_dataframe(rs)
        if not stock_df.empty:
            codes.update(stock_df["code"].tolist())
    codes.add(BENCHMARK_CODE)
    sorted_codes = sorted(codes)
    cache_file.parent.mkdir(parents=True, exist_ok=True)
    cache_file.write_text("\n".join(sorted_codes) + "\n", encoding="utf-8")
    return sorted_codes


def _clean_code(code: str) -> str:
    return code.replace(".", "")


def _load_adjust_factor(code: str, start_date: str, end_date: str) -> pd.Series | None:
    rs_adj = bs.query_adjust_factor(code, start_date=start_date, end_date=end_date)
    if rs_adj.error_code != "0":
        return None
    adj_df = _query_to_dataframe(rs_adj)
    if adj_df.empty or "dividOperateDate" not in adj_df.columns or "adjustFactor" not in adj_df.columns:
        return None
    return adj_df.set_index("dividOperateDate")["adjustFactor"].rename("factor")


def _ensure_login() -> None:
    login_result = bs.login()
    if login_result.error_code != "0":
        raise RuntimeError(f"Baostock login failed: {login_result.error_msg}")


def _download_single_code(code: str, start_date: str, end_date: str, output_dir: Path, fields: str) -> str:
    output_file = output_dir / f"{_clean_code(code)}.csv"
    if output_file.exists() and output_file.stat().st_size > 0:
        return code
    last_error: Exception | None = None

    for attempt in range(1, REQUEST_RETRY_LIMIT + 1):
        try:
            rs = bs.query_history_k_data_plus(
                code,
                fields,
                start_date=start_date,
                end_date=end_date,
                frequency="d",
                adjustflag="1",
            )
            if rs.error_code != "0":
                raise RuntimeError(f"Failed to fetch {code}: {rs.error_msg}")

            new_df = _query_to_dataframe(rs)
            if new_df.empty:
                return code

            new_df = new_df.set_index("date")
            adj_factor = _load_adjust_factor(code, start_date, end_date)
            if adj_factor is not None:
                new_df = pd.concat([new_df, adj_factor], axis=1).ffill()
            if "factor" not in new_df.columns:
                new_df["factor"] = 1.0

            new_df["code"] = new_df["code"].str.replace(".", "", regex=False)
            numeric_cols = [column for column in new_df.columns if column != "code"]
            new_df[numeric_cols] = new_df[numeric_cols].apply(pd.to_numeric, errors="coerce")
            new_df.index.name = "date"
            new_df = new_df.reset_index()
            new_df.to_csv(output_file, index=False, encoding="utf-8")
            return code
        except Exception as exc:
            last_error = exc
            print(f"[retry {attempt}/{REQUEST_RETRY_LIMIT}] {code} 下载失败: {exc}", flush=True)
            try:
                bs.logout()
            except Exception:
                pass
            time.sleep(RETRY_SLEEP_SECONDS)
            _ensure_login()

    assert last_error is not None
    raise RuntimeError(f"{code} 下载失败，重试 {REQUEST_RETRY_LIMIT} 次后仍未成功: {last_error}")


def download_stock_data(codes: list[str], start_date: str, end_date: str, output_dir: Path) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    fields = "date,code,open,high,low,close,preclose,volume,amount,turn,tradestatus,pctChg"
    for code in tqdm(codes, total=len(codes), desc="下载进度", unit="stock"):
        _download_single_code(code, start_date, end_date, output_dir, fields)


def main() -> None:
    start_date = DATASET_LOCK.train_period.start
    end_date = DATASET_LOCK.test_period.end

    _ensure_login()

    try:
        print("开始收集历史CSI500成分股...", flush=True)
        codes = get_csi500_stock_union(start_date, end_date)
        print(f"Universe: {TARGET_UNIVERSE}")
        print(f"Dataset version: {DATASET_LOCK.dataset_version}")
        print(f"Date range: {start_date} -> {end_date}")
        print(f"Download target count: {len(codes)}")
        print(f"Benchmark included: {BENCHMARK_CODE}")
        print("开始下载原始CSV数据...", flush=True)
        download_stock_data(codes, start_date, end_date, RAW_DATA_DIR)
        print(f"Raw CSV saved to: {RAW_DATA_DIR}")
    finally:
        bs.logout()


if __name__ == "__main__":
    main()