"""挂载 vnpy 自带 Alpha101/Alpha158 表达式因子到注册表.

vnpy 的 AlphaDataset 子类在 __init__ 里只做 add_feature 字符串注册(不触发计算),
用空 schema df 实例化即可零成本提取全部表达式.
"""
import sys
import os
import polars as pl

_VNPY_SRC = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "vnpy_v4.4.0"))
if _VNPY_SRC not in sys.path:
    sys.path.insert(0, _VNPY_SRC)

from .registry import register_factor, list_factors, _REGISTRY

_ALPHA_SCHEMA = {
    "vt_symbol": pl.Utf8,
    "datetime": pl.Datetime,
    "open": pl.Float64,
    "high": pl.Float64,
    "low": pl.Float64,
    "close": pl.Float64,
    "volume": pl.Float64,
    "turnover": pl.Float64,
    "open_interest": pl.Float64,
    "vwap": pl.Float64,
}

_DUMMY_PERIOD = ("2018-01-01", "2018-01-02")


def _extract_expressions(dataset_cls) -> dict[str, str]:
    """实例化数据集类,提取其注册的全部 {name: expression}."""
    df = pl.DataFrame(schema=_ALPHA_SCHEMA)
    ds = dataset_cls(df, _DUMMY_PERIOD, _DUMMY_PERIOD, _DUMMY_PERIOD)
    return {name: str(expr) for name, expr in ds.feature_expressions.items()}


def _mount(dataset_cls, category: str) -> int:
    """把一个数据集类的表达式挂进注册表(已存在同名跳过,幂等)."""
    mounted = 0
    for name, expression in _extract_expressions(dataset_cls).items():
        if name in _REGISTRY:
            continue
        register_factor(name, expression, category)
        mounted += 1
    return mounted


def mount_alpha101() -> int:
    """挂载 WorldQuant Alpha101(100 个),category=alpha101."""
    from vnpy.alpha.dataset.datasets.alpha_101 import Alpha101
    return _mount(Alpha101, "alpha101")


def mount_alpha158() -> int:
    """挂载 Qlib Alpha158(158 个),category=alpha158."""
    from vnpy.alpha.dataset.datasets.alpha_158 import Alpha158
    return _mount(Alpha158, "alpha158")


def mount_all() -> dict[str, int]:
    """挂载全部内置数据集,返回各类新挂数量."""
    return {"alpha101": mount_alpha101(), "alpha158": mount_alpha158()}


# 模块导入即挂载(与 library.py 内置因子同一模式,下游 import sanguo_factor 即全量可见)
mount_all()