factor_engine/formula_registry.py

"""Factor formula registry: load, compute, cache, and export user-defined qlib expressions."""

from __future__ import annotations

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

import pandas as pd
import yaml

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


DEFAULT_REGISTRY_PATH = PROJECT_ROOT / "config" / "factor_registry.yaml"


@dataclass
class FactorSpec:
    name: str
    expression: str
    description: str = ""
    tags: list[str] | None = None
    enabled: bool = True
    market: str | None = None
    label_expr: str | None = None

    @classmethod
    def from_dict(cls, name: str, data: dict[str, Any], defaults: dict[str, Any]) -> FactorSpec:
        return cls(
            name=name,
            expression=str(data["expression"]),
            description=data.get("description", ""),
            tags=data.get("tags") or [],
            enabled=bool(data.get("enabled", True)),
            market=data.get("market"),
            label_expr=data.get("label_expr", defaults.get("label_expr")),
        )


def load_registry(path: str | Path | None = None) -> dict[str, Any]:
    path = Path(path) if path else DEFAULT_REGISTRY_PATH
    if not path.is_absolute():
        path = PROJECT_ROOT / path
    if not path.exists():
        raise FileNotFoundError(f"Factor registry not found: {path}")
    with open(path, encoding="utf-8") as f:
        return yaml.safe_load(f) or {}


def save_registry(data: dict[str, Any], path: str | Path | None = None) -> Path:
    path = Path(path) if path else DEFAULT_REGISTRY_PATH
    if not path.is_absolute():
        path = PROJECT_ROOT / path
    with open(path, "w", encoding="utf-8") as f:
        yaml.safe_dump(data, f, allow_unicode=True, sort_keys=False)
    return path


def list_factor_specs(
    path: str | Path | None = None,
    enabled_only: bool = False,
    tag: str | None = None,
) -> list[FactorSpec]:
    raw = load_registry(path)
    defaults = raw.get("defaults", {})
    specs = []
    for name, info in raw.get("factors", {}).items():
        spec = FactorSpec.from_dict(name, info, defaults)
        if enabled_only and not spec.enabled:
            continue
        if tag and tag not in (spec.tags or []):
            continue
        specs.append(spec)
    return specs


def get_factor_spec(name: str, path: str | Path | None = None) -> FactorSpec:
    for spec in list_factor_specs(path):
        if spec.name == name:
            return spec
    raise KeyError(f"Factor not found in registry: {name}")


def add_factor_to_registry(
    name: str,
    expression: str,
    description: str = "",
    tags: list[str] | None = None,
    enabled: bool = True,
    path: str | Path | None = None,
) -> FactorSpec:
    path = Path(path) if path else DEFAULT_REGISTRY_PATH
    if not path.is_absolute():
        path = PROJECT_ROOT / path

    data = load_registry(path) if path.exists() else {"defaults": {}, "factors": {}}
    data.setdefault("factors", {})[name] = {
        "expression": expression,
        "description": description,
        "tags": tags or ["custom"],
        "enabled": enabled,
    }
    save_registry(data, path)
    defaults = data.get("defaults", {})
    return FactorSpec.from_dict(name, data["factors"][name], defaults)


def registry_output_dir() -> Path:
    settings = load_settings()
    out = settings.output_root / "factors" / "registry"
    out.mkdir(parents=True, exist_ok=True)
    return out


def compute_factor(
    name: str,
    start_time: str | None = None,
    end_time: str | None = None,
    cache: bool = True,
    registry_path: str | Path | None = None,
) -> pd.Series:
    """Compute a single registered factor via qlib D.features."""
    spec = get_factor_spec(name, registry_path)
    settings = load_settings()
    start_time = start_time or settings.raw["data"]["start_time"]
    end_time = end_time or settings.raw["data"]["end_time"]

    series = load_qlib_expression_factor(
        expression=spec.expression,
        instruments=spec.market or settings.market,
        start_time=start_time,
        end_time=end_time,
        name=spec.name,
    )

    if cache:
        out_dir = registry_output_dir()
        meta = {
            "name": spec.name,
            "expression": spec.expression,
            "description": spec.description,
            "computed_at": datetime.now().isoformat(),
            "start_time": start_time,
            "end_time": end_time,
        }
        with open(out_dir / f"{spec.name}.meta.json", "w", encoding="utf-8") as f:
            json.dump(meta, f, ensure_ascii=False, indent=2)
        series.to_frame(spec.name).to_parquet(out_dir / f"{spec.name}.parquet")

    return series


def compute_all_factors(
    enabled_only: bool = True,
    start_time: str | None = None,
    end_time: str | None = None,
    cache: bool = True,
    registry_path: str | Path | None = None,
) -> dict[str, pd.Series]:
    results = {}
    for spec in list_factor_specs(registry_path, enabled_only=enabled_only):
        print(f"Computing factor: {spec.name} ...")
        results[spec.name] = compute_factor(
            spec.name,
            start_time=start_time,
            end_time=end_time,
            cache=cache,
            registry_path=registry_path,
        )
    return results


def factor_series_to_panel(series: pd.Series, factor_name: str | None = None) -> pd.DataFrame:
    """Convert qlib MultiIndex Series to long panel (date, symbol, factor)."""
    factor_name = factor_name or series.name or "factor"
    df = series.rename(factor_name).reset_index()
    df = df.rename(columns={"datetime": "date", "instrument": "symbol"})
    df["date"] = pd.to_datetime(df["date"])
    return df


def build_combined_panel(
    factor_names: list[str] | None = None,
    enabled_only: bool = True,
    use_cache: bool = True,
    registry_path: str | Path | None = None,
) -> pd.DataFrame:
    """Merge multiple registry factors into one panel for multi-factor backtest."""
    specs = list_factor_specs(registry_path, enabled_only=enabled_only)
    if factor_names:
        specs = [s for s in specs if s.name in factor_names]

    panels = []
    out_dir = registry_output_dir()
    for spec in specs:
        cache_path = out_dir / f"{spec.name}.parquet"
        if use_cache and cache_path.exists():
            part = pd.read_parquet(cache_path)
            part = part.reset_index() if isinstance(part.index, pd.MultiIndex) else part
            if "datetime" in part.columns:
                part = part.rename(columns={"datetime": "date", "instrument": "symbol"})
        else:
            s = compute_factor(spec.name, cache=True, registry_path=registry_path)
            part = factor_series_to_panel(s, spec.name)

        col = spec.name
        if col not in part.columns:
            col = [c for c in part.columns if c not in ("date", "symbol")][0]
        panels.append(part[["date", "symbol", col]])

    if not panels:
        raise ValueError("No factors to combine")

    merged = panels[0]
    for part in panels[1:]:
        merged = merged.merge(part, on=["date", "symbol"], how="outer")
    return merged.sort_values(["date", "symbol"]).reset_index(drop=True)


def load_label_panel(
    start_time: str | None = None,
    end_time: str | None = None,
    label_expr: str | None = None,
    market: str | None = None,
) -> pd.DataFrame:
    settings = load_settings()
    init_qlib()
    from qlib.data import D

    market = market or settings.market
    start_time = start_time or settings.raw["data"]["start_time"]
    end_time = end_time or settings.raw["data"]["end_time"]
    label_expr = label_expr or settings.raw["data"].get("label_expr", "Ref($close, -2)/Ref($close, -1) - 1")

    label = D.features(
        D.instruments(market),
        [label_expr],
        start_time=start_time,
        end_time=end_time,
        freq=settings.freq,
    )
    label.columns = ["label"]
    panel = label.reset_index().rename(columns={"datetime": "date", "instrument": "symbol"})
    panel["date"] = pd.to_datetime(panel["date"])
    return panel


def build_signal_source_for_factor(name: str, registry_path: str | Path | None = None) -> dict[str, Any]:
    spec = get_factor_spec(name, registry_path)
    return {
        "type": "factor_registry",
        "name": spec.name,
    }


def export_gp_seed_formulas(output_path: str | Path | None = None, enabled_only: bool = True) -> Path:
    """Export registry expressions as GP mining seed formulas (one per line)."""
    settings = load_settings()
    out = Path(output_path) if output_path else settings.output_root / "factors" / "gp_seed_formulas.txt"
    if not out.is_absolute():
        out = PROJECT_ROOT / out
    out.parent.mkdir(parents=True, exist_ok=True)

    lines = []
    for spec in list_factor_specs(enabled_only=enabled_only):
        lines.append(f"# {spec.name}: {spec.description}")
        lines.append(spec.expression)
        lines.append("")

    out.write_text("\n".join(lines), encoding="utf-8")
    return out