factor_engine/gp/config.py

"""GP mining configuration loader."""

from __future__ import annotations

import os
from pathlib import Path
from typing import Any

import yaml

PROJECT_ROOT = Path(__file__).resolve().parents[2]
DEFAULT_GP_CONFIG = PROJECT_ROOT / "config" / "gp_mining.yaml"
DEFAULT_BASE_CONFIG = PROJECT_ROOT / "config" / "base.yaml"


def load_gp_config(
    gp_config_path: str | Path | None = None,
    base_config_path: str | Path | None = None,
) -> dict[str, Any]:
    gp_path = Path(gp_config_path) if gp_config_path else DEFAULT_GP_CONFIG
    base_path = Path(base_config_path) if base_config_path else DEFAULT_BASE_CONFIG

    with open(gp_path, encoding="utf-8") as f:
        gp_cfg = yaml.safe_load(f)["gp_mining"]
    with open(base_path, encoding="utf-8") as f:
        base_cfg = yaml.safe_load(f)

    cfg = {"gp": gp_cfg, "base": base_cfg}

    # Allow runtime overrides via environment variables
    env_int_keys = {
        "POPULATION_SIZE": "population_size",
        "GENERATIONS_PER_RUN": "generations_per_run",
        "MAX_INIT_DEPTH": "max_init_depth",
        "TOP_K_EXPORT": "top_k_export",
    }
    for env_key, cfg_key in env_int_keys.items():
        if env_key in os.environ:
            gp_cfg[cfg_key] = int(os.environ[env_key])

    if "SEED" in os.environ:
        gp_cfg["seed"] = int(os.environ["SEED"])

    run_id = os.environ.get("RUN_ID", gp_cfg.get("run_id", "qlib_gp_run_0"))
    output_root = Path(base_cfg.get("output", {}).get("root", "outputs"))
    if not output_root.is_absolute():
        output_root = PROJECT_ROOT / output_root

    cfg["output_dir"] = output_root / "gp_mining" / run_id
    cfg["output_dir"].mkdir(parents=True, exist_ok=True)
    cfg["run_id"] = run_id
    return cfg