data_pipeline/factor_loader.py

"""Load pre-mined factors from GP parquet, QuantaAlpha JSON, qlib expressions, or pred pickles."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import pandas as pd

from config.settings import PROJECT_ROOT, load_settings
from data_pipeline.init_qlib import init_qlib


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


def _to_signal_series(df: pd.DataFrame, score_col: str = "score") -> pd.Series:
    """Convert panel to qlib MultiIndex (instrument, datetime) signal."""
    if isinstance(df, pd.Series):
        s = df.copy()
        if s.index.names != ["instrument", "datetime"]:
            s = s.swaplevel().sort_index()
        s.index.names = ["instrument", "datetime"]
        return s

    panel = df.copy()
    if "date" in panel.columns:
        panel = panel.rename(columns={"date": "datetime"})
    if "symbol" in panel.columns:
        panel = panel.rename(columns={"symbol": "instrument"})
    panel["instrument"] = panel["instrument"].map(_normalize_symbol)
    panel["datetime"] = pd.to_datetime(panel["datetime"])
    out = panel.set_index(["instrument", "datetime"])[score_col].sort_index()
    out.index.names = ["instrument", "datetime"]
    return out


def load_factor_panel(path: str | Path) -> pd.DataFrame:
    path = Path(path)
    if not path.is_absolute():
        path = PROJECT_ROOT / path
    if not path.exists():
        raise FileNotFoundError(f"Factor panel not found: {path}")

    if path.suffix == ".parquet":
        df = pd.read_parquet(path)
    elif path.suffix == ".csv":
        df = pd.read_csv(path)
    elif path.suffix == ".json":
        return load_quantaalpha_library(path)
    else:
        raise ValueError(f"Unsupported factor format: {path.suffix}")

    if "date" in df.columns:
        df["date"] = pd.to_datetime(df["date"])
    return df


def load_quantaalpha_library(path: str | Path) -> pd.DataFrame:
    """Parse QuantaAlpha factor library JSON into a metadata table."""
    path = Path(path)
    with open(path, encoding="utf-8") as f:
        data = json.load(f)

    rows = []
    for fid, info in data.get("factors", {}).items():
        backtest = info.get("backtest_results", {}) or {}
        rows.append(
            {
                "factor_id": fid,
                "factor_name": info.get("factor_name", fid),
                "factor_expression": info.get("factor_expression", ""),
                "factor_description": info.get("factor_description", ""),
                "ic": backtest.get("IC", backtest.get("ic")),
                "icir": backtest.get("ICIR", backtest.get("icir")),
                "rank_ic": backtest.get("Rank IC", backtest.get("rank_ic")),
                "rank_icir": backtest.get("Rank ICIR", backtest.get("rank_icir")),
                "source": "quantaalpha",
            }
        )
    return pd.DataFrame(rows)


def load_qlib_expression_factor(
    expression: str,
    instruments: str | list | None = None,
    start_time: str | None = None,
    end_time: str | None = None,
    name: str = "factor",
) -> pd.Series:
    settings = load_settings()
    init_qlib()
    from qlib.data import D

    market = instruments or settings.market
    start_time = start_time or settings.raw["data"]["start_time"]
    end_time = end_time or settings.raw["data"]["end_time"]

    df = D.features(
        D.instruments(market),
        [expression],
        start_time=start_time,
        end_time=end_time,
        freq=settings.freq,
    )
    df.columns = [name]
    s = df[name]
    s.index.names = ["instrument", "datetime"]
    return s


def load_pred_pickle(path: str | Path) -> pd.Series:
    import pickle

    path = Path(path)
    if not path.is_absolute():
        path = PROJECT_ROOT / path
    with open(path, "rb") as f:
        obj = pickle.load(f)
    if isinstance(obj, pd.Series):
        return obj
    if isinstance(obj, pd.DataFrame):
        col = "score" if "score" in obj.columns else obj.columns[0]
        return _to_signal_series(obj, col)
    raise TypeError(f"Unsupported pred pickle type: {type(obj)}")


def combine_factor_columns(
    panel: pd.DataFrame,
    factor_cols: list[str] | None = None,
    method: str = "equal",
    ic_weights: dict[str, float] | None = None,
) -> pd.Series:
    """Combine multiple factor columns into one cross-sectional score."""
    factor_cols = factor_cols or [c for c in panel.columns if c.startswith("factor_") or c.startswith("alpha_")]
    if not factor_cols:
        raise ValueError("No factor columns found to combine")

    work = panel[["date" if "date" in panel.columns else "datetime", "symbol" if "symbol" in panel.columns else "instrument", *factor_cols]].copy()
    date_col = "date" if "date" in work.columns else "datetime"
    sym_col = "symbol" if "symbol" in work.columns else "instrument"

    def _zscore(x: pd.DataFrame) -> pd.DataFrame:
        return x.apply(lambda s: (s - s.mean()) / (s.std() + 1e-8))

    grouped = work.groupby(date_col)
    norm = grouped[factor_cols].transform(lambda x: (x - x.mean()) / (x.std() + 1e-8))

    if method == "equal":
        score = norm.mean(axis=1)
    elif method == "rank_mean":
        score = grouped[factor_cols].rank(pct=True).mean(axis=1)
    elif method == "ic_weighted":
        if not ic_weights:
            raise ValueError("ic_weighted requires ic_weights dict")
        weights = pd.Series({c: ic_weights.get(c, 0.0) for c in factor_cols})
        weights = weights / weights.abs().sum()
        score = norm.mul(weights, axis=1).sum(axis=1)
    else:
        raise ValueError(f"Unknown combine method: {method}")

    out = work[[date_col, sym_col]].copy()
    out["score"] = score.values
    return _to_signal_series(out, "score")


def build_signal_from_source(source_cfg: dict[str, Any]) -> pd.Series:
    """Build qlib signal Series from a signal source config block."""
    src_type = source_cfg.get("type", "factor_panel")

    if src_type == "pred_pickle":
        return load_pred_pickle(source_cfg["path"])

    if src_type == "qlib_expression":
        return load_qlib_expression_factor(
            expression=source_cfg["expression"],
            instruments=source_cfg.get("instruments"),
            start_time=source_cfg.get("start_time"),
            end_time=source_cfg.get("end_time"),
            name=source_cfg.get("name", "score"),
        )

    if src_type == "quantaalpha_library":
        from integrations.quantaalpha.factor_library import build_signal_from_library

        return build_signal_from_library(
            library_path=source_cfg["path"],
            factor_ids=source_cfg.get("factor_ids"),
            top_k=source_cfg.get("top_k"),
            combine=source_cfg.get("combine", "ic_weighted"),
            quality_filter=source_cfg.get("quality_filter"),
        )

    if src_type == "factor_registry":
        from factor_engine.formula_registry import compute_factor

        return compute_factor(
            source_cfg["name"],
            start_time=source_cfg.get("start_time"),
            end_time=source_cfg.get("end_time"),
            cache=source_cfg.get("cache", True),
            registry_path=source_cfg.get("registry_path"),
        )

    if src_type == "factor_registry_panel":
        from factor_engine.formula_registry import build_combined_panel

        panel = build_combined_panel(
            factor_names=source_cfg.get("names"),
            enabled_only=source_cfg.get("enabled_only", True),
            use_cache=source_cfg.get("use_cache", True),
            registry_path=source_cfg.get("registry_path"),
        )
        return combine_factor_columns(
            panel,
            factor_cols=source_cfg.get("factor_cols"),
            method=source_cfg.get("combine", "equal"),
        )

    panel = load_factor_panel(source_cfg["path"])
    if source_cfg.get("score_col"):
        return _to_signal_series(panel, source_cfg["score_col"])

    factor_cols = source_cfg.get("factor_cols")
    if factor_cols or any(c.startswith("factor_") for c in panel.columns):
        return combine_factor_columns(
            panel,
            factor_cols=factor_cols,
            method=source_cfg.get("combine", "equal"),
            ic_weights=source_cfg.get("ic_weights"),
        )

    if "score" in panel.columns:
        return _to_signal_series(panel, "score")
    if "pred_score" in panel.columns:
        return _to_signal_series(panel, "pred_score")

    raise ValueError(f"Cannot infer signal from source config: {source_cfg}")