factor_engine/operator_builder.py

"""Build qlib factor expressions from operator trees (YAML/JSON AST)."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

import yaml

from config.settings import PROJECT_ROOT
from factor_engine.formula_registry import add_factor_to_registry, load_registry, save_registry


DEFAULT_BUILTINS_PATH = PROJECT_ROOT / "config" / "factor_builtins.yaml"


@dataclass
class BuiltFactorSpec:
    name: str
    tree: dict[str, Any]
    description: str = ""
    tags: list[str] | None = None
    enabled: bool = True


def load_builtins(path: str | Path | None = None) -> dict[str, Any]:
    path = Path(path) if path else DEFAULT_BUILTINS_PATH
    if not path.is_absolute():
        path = PROJECT_ROOT / path
    if not path.exists():
        return {"built_factors": {}, "operator_catalog": {}}
    with open(path, encoding="utf-8") as f:
        return yaml.safe_load(f) or {}


def list_built_factors(path: str | Path | None = None, enabled_only: bool = False) -> list[BuiltFactorSpec]:
    raw = load_builtins(path)
    out = []
    for name, info in raw.get("built_factors", {}).items():
        spec = BuiltFactorSpec(
            name=name,
            tree=info["tree"],
            description=info.get("description", ""),
            tags=info.get("tags") or [],
            enabled=bool(info.get("enabled", True)),
        )
        if enabled_only and not spec.enabled:
            continue
        out.append(spec)
    return out


def get_built_factor(name: str, path: str | Path | None = None) -> BuiltFactorSpec:
    for spec in list_built_factors(path):
        if spec.name == name:
            return spec
    raise KeyError(f"Built factor not found: {name}")


def _render_node(node: Any) -> str:
    if node is None:
        raise ValueError("Empty operator node")

    if isinstance(node, (int, float)):
        return str(node)

    if isinstance(node, str):
        return node

    if isinstance(node, dict):
        if "field" in node:
            return str(node["field"])
        if "const" in node:
            return str(int(node["const"]))

        op = str(node.get("op", "")).lower()
        args = node.get("args", [])

        if op == "field":
            return str(node.get("name") or node.get("value"))

        if op in ("add", "sub", "mul", "div"):
            if len(args) != 2:
                raise ValueError(f"{op} requires 2 args")
            left, right = (_render_node(a) for a in args)
            sym = {"add": "+", "sub": "-", "mul": "*", "div": "/"}[op]
            return f"({left}){sym}({right})"

        if op in ("max", "min"):
            if len(args) != 2:
                raise ValueError(f"{op} requires 2 args")
            fname = "Greater" if op == "max" else "Less"
            return f"{fname}({_render_node(args[0])}, {_render_node(args[1])})"

        if op == "abs":
            return f"Abs({_render_node(args[0])})"
        if op == "log":
            return f"Log({_render_node(args[0])})"
        if op == "rank":
            window = int(_render_node(args[1])) if len(args) > 1 else 20
            return f"Rank({_render_node(args[0])}, {window})"
        if op == "neg":
            return f"(-({_render_node(args[0])}))"

        if op == "mean":
            return f"Mean({_render_node(args[0])}, {int(_render_node(args[1]))})"
        if op == "std":
            return f"Std({_render_node(args[0])}, {int(_render_node(args[1]))})"
        if op == "sum":
            return f"Sum({_render_node(args[0])}, {int(_render_node(args[1]))})"
        if op == "ref":
            d = int(_render_node(args[1]))
            sign = -abs(d)
            return f"Ref({_render_node(args[0])}, {sign})"
        if op == "delta":
            d = int(_render_node(args[1]))
            x = _render_node(args[0])
            return f"({x}-Ref({x}, -{abs(d)}))"

        raise ValueError(f"Unsupported operator: {op}")

    raise TypeError(f"Unsupported node type: {type(node)}")


def tree_to_qlib_expression(tree: dict[str, Any]) -> str:
    return _render_node(tree)


def build_expression(name: str, path: str | Path | None = None) -> str:
    spec = get_built_factor(name, path)
    return tree_to_qlib_expression(spec.tree)


def register_built_factor(
    name: str,
    builtins_path: str | Path | None = None,
    registry_path: str | Path | None = None,
    overwrite: bool = True,
) -> str:
    """Compile operator tree → qlib expression and write into factor_registry.yaml."""
    spec = get_built_factor(name, builtins_path)
    expr = tree_to_qlib_expression(spec.tree)

    reg_path = Path(registry_path) if registry_path else PROJECT_ROOT / "config" / "factor_registry.yaml"
    if not reg_path.is_absolute():
        reg_path = PROJECT_ROOT / reg_path

    data = load_registry(reg_path) if reg_path.exists() else {"defaults": {}, "factors": {}}
    if not overwrite and name in data.get("factors", {}):
        raise FileExistsError(f"Factor already exists in registry: {name}")

    data.setdefault("factors", {})[name] = {
        "expression": expr,
        "description": spec.description or f"Built from operator tree ({name})",
        "tags": (spec.tags or []) + ["operator_built"],
        "enabled": spec.enabled,
        "source": "operator_builder",
        "operator_tree": spec.tree,
    }
    save_registry(data, reg_path)
    return expr


def register_all_built_factors(
    builtins_path: str | Path | None = None,
    registry_path: str | Path | None = None,
    enabled_only: bool = True,
) -> dict[str, str]:
    results = {}
    for spec in list_built_factors(builtins_path, enabled_only=enabled_only):
        expr = register_built_factor(spec.name, builtins_path, registry_path, overwrite=True)
        results[spec.name] = expr
    return results


def get_operator_catalog(path: str | Path | None = None) -> dict[str, Any]:
    return load_builtins(path).get("operator_catalog", {})