data_pipeline/gp_dataset_builder.py

"""Build qlib-compatible handler/dataset from GP mining outputs."""

from __future__ import annotations

from pathlib import Path

import pandas as pd

from config.settings import load_settings
from data_pipeline.init_qlib import init_qlib
from qlib.contrib.data.handler import check_transform_proc
from qlib.data import D
from qlib.data.dataset import DatasetH
from qlib.data.dataset.handler import DataHandlerLP
from qlib.data.dataset.loader import StaticDataLoader


def _normalize_symbol(code: str) -> str:
    code = str(code).upper()
    if code.startswith(("SH", "SZ", "BJ")):
        return code
    if code[0] == "6":
        return f"SH{code}"
    return f"SZ{code}"


def load_gp_feature_frame(run_id: str | None = None) -> pd.DataFrame:
    settings = load_settings()
    run_id = run_id or settings.raw.get("experiment", {}).get("run_id", "qlib_gp_run_0")
    gp_dir = settings.gp_output_dir(run_id)

    parquet_path = gp_dir / "ML_Features_qlib.parquet"
    csv_path = gp_dir / "ML_Features_qlib.csv"
    if parquet_path.exists():
        df = pd.read_parquet(parquet_path)
    elif csv_path.exists():
        df = pd.read_csv(csv_path)
    else:
        raise FileNotFoundError(f"GP features not found under {gp_dir}")

    df["date"] = pd.to_datetime(df["date"])
    df["symbol"] = df["symbol"].map(_normalize_symbol)
    return df


def build_qlib_label(instruments, start: str, end: str, label_expr: str) -> pd.Series:
    settings = load_settings()
    label_df = D.features(
        instruments,
        [label_expr],
        start_time=start,
        end_time=end,
        freq=settings.freq,
    )
    label_df.columns = ["LABEL0"]
    return label_df["LABEL0"]


def gp_features_to_qlib_df(
    gp_df: pd.DataFrame,
    label: pd.Series | None = None,
    target_col: str = "target_return",
) -> pd.DataFrame:
    factor_cols = [c for c in gp_df.columns if c.startswith("factor_")]
    if not factor_cols:
        raise ValueError("No factor_* columns found in GP feature frame")

    panel = gp_df[["date", "symbol", *factor_cols]].copy()
    panel["date"] = pd.to_datetime(panel["date"])
    panel = panel.set_index(["date", "symbol"]).sort_index()
    panel.index.names = ["datetime", "instrument"]

    if label is None:
        aligned = gp_df.set_index(["date", "symbol"])[target_col]
        aligned.index.names = ["datetime", "instrument"]
        panel["LABEL0"] = aligned.reindex(panel.index)
    else:
        label = label.copy()
        label.index.names = ["datetime", "instrument"]
        panel = panel.join(label.rename("LABEL0"), how="left")

    out = pd.concat(
        {
            "feature": panel[factor_cols],
            "label": panel[["LABEL0"]],
        },
        axis=1,
    )
    out.index.names = ["datetime", "instrument"]
    out = out.sort_index()
    out = out.sort_index(level=["datetime", "instrument"], sort_remaining=True)
    return out


def build_gp_handler(run_id: str | None = None) -> DataHandlerLP:
    settings = load_settings()
    init_qlib()
    market = settings.market
    inst_config = D.instruments(market)

    gp_df = load_gp_feature_frame(run_id)
    start = settings.raw["data"]["start_time"]
    end = settings.raw["data"]["end_time"]
    label_expr = settings.raw["data"].get("label_expr", "Ref($close, -2)/Ref($close, -1) - 1")

    inst_list = D.list_instruments(inst_config, start_time=start, end_time=end, as_list=True)
    gp_df = gp_df[gp_df["symbol"].isin(inst_list)]

    try:
        label = build_qlib_label(inst_list, start, end, label_expr)
    except Exception:
        label = None

    qlib_df = gp_features_to_qlib_df(gp_df, label=label)
    fit_start, fit_end = settings.fit_segment

    infer_processors = check_transform_proc(
        [
            {"class": "ProcessInf"},
            {"class": "ZScoreNorm"},
            {"class": "Fillna"},
        ],
        fit_start,
        fit_end,
    )
    learn_processors = [
        {"class": "DropnaLabel"},
        {"class": "CSZScoreNorm", "kwargs": {"fields_group": "label"}},
    ]

    handler = DataHandlerLP(
        instruments=None,
        start_time=start,
        end_time=end,
        data_loader=StaticDataLoader(qlib_df),
        infer_processors=infer_processors,
        learn_processors=learn_processors,
        process_type=DataHandlerLP.PTYPE_A,
    )
    return handler


def build_gp_dataset(run_id: str | None = None) -> DatasetH:
    settings = load_settings()
    handler = build_gp_handler(run_id)
    dataset = DatasetH(handler=handler, segments=settings.segments)
    return dataset


def export_gp_artifacts(run_id: str | None = None) -> dict[str, Path]:
    settings = load_settings()
    out_dir = settings.gp_output_dir(run_id)
    out_dir.mkdir(parents=True, exist_ok=True)

    handler = build_gp_handler(run_id)
    handler_path = out_dir / "gp_qlib_handler.pkl"
    handler.to_pickle(str(handler_path), dump_all=True)

    dataset = DatasetH(handler=handler, segments=settings.segments)
    dataset_path = out_dir / "gp_qlib_dataset.pkl"
    dataset.config(dump_all=True, recursive=True)
    dataset.to_pickle(str(dataset_path))

    gp_df = load_gp_feature_frame(run_id)
    qlib_df = gp_features_to_qlib_df(gp_df, label=None)
    feature_path = out_dir / "gp_qlib_features.parquet"
    qlib_df.to_parquet(feature_path)

    paths = {
        "handler": handler_path,
        "dataset": dataset_path,
        "features": feature_path,
    }
    print("Exported GP qlib artifacts:")
    for k, p in paths.items():
        print(f"  {k}: {p}")
    return paths