A_Share_DP/src/ashare_dp/features/registry.py

"""Feature Registry — centralized management of all market features.

Engine code only knows about the registry; individual feature functions
are looked up by name. Adding a new feature = register it, no engine changes.
"""

from __future__ import annotations

from datetime import date
from typing import Any, Callable

from ashare_dp.domain.features import FeatureDefinition


class FeatureRegistry:
    """Central registry for all market features.

    Usage:
        from ashare_dp.features.registry import registry

        # Register a feature
        registry.register(FeatureDefinition(
            name="breadth_vector",
            category="breadth",
            description="Advance/decline ratio and related breadth metrics",
            dependencies=[],
            compute=_compute_breadth_vector,
        ))

        # Compute all features for a date
        features = registry.compute_all(db, trade_date)

        # Get features by category
        breadth_features = registry.get_by_category("breadth")
    """

    def __init__(self):
        self._features: dict[str, FeatureDefinition] = {}

    def register(self, feature: FeatureDefinition) -> None:
        """Register a feature definition. Overwrites if name exists."""
        self._features[feature.name] = feature

    def get(self, name: str) -> FeatureDefinition | None:
        """Get a feature definition by name."""
        return self._features.get(name)

    def get_by_category(self, category: str) -> list[FeatureDefinition]:
        """Get all features in a category."""
        return [f for f in self._features.values() if f.category == category]

    def compute_all(self, db: Any, trade_date: date) -> dict[str, dict[str, Any]]:
        """Compute all registered features for a given trading date.

        Args:
            db: Database instance (read_only DuckDB connection).
            trade_date: The trading date to compute features for.

        Returns:
            Dict mapping feature name → feature value dict.
            Features are computed in dependency order (simple topological sort).
        """
        results: dict[str, dict[str, Any]] = {}
        computed: set[str] = set()
        pending: set[str] = set(self._features.keys())

        while pending:
            ready = [
                name for name in pending
                if all(dep in computed for dep in self._features[name].dependencies)
            ]
            if not ready:
                # Circular dependency or missing dependency
                remaining = ", ".join(sorted(pending))
                raise RuntimeError(
                    f"Cannot resolve feature dependencies. "
                    f"Remaining: {remaining}. Computed: {computed}"
                )

            for name in ready:
                feature = self._features[name]
                if feature.compute is not None:
                    results[name] = feature.compute(db, trade_date, results)
                computed.add(name)
                pending.discard(name)

        return results

    def compute_category(
        self, db: Any, trade_date: date, category: str
    ) -> dict[str, dict[str, Any]]:
        """Compute all features in a specific category."""
        features = self.get_by_category(category)
        # Compute dependencies first
        all_needed: set[str] = set()
        for f in features:
            all_needed.add(f.name)
            all_needed.update(f.dependencies)

        full_results = self.compute_all(db, trade_date)
        return {k: v for k, v in full_results.items() if k in all_needed}

    def list_categories(self) -> list[str]:
        """List all unique feature categories."""
        return sorted({f.category for f in self._features.values()})

    def list_features(self) -> list[dict]:
        """List all registered features with metadata."""
        return [
            {
                "name": f.name,
                "category": f.category,
                "description": f.description,
                "dependencies": f.dependencies,
                "version": f.version,
            }
            for f in sorted(self._features.values(), key=lambda x: x.name)
        ]

    def __len__(self) -> int:
        return len(self._features)

    def __contains__(self, name: str) -> bool:
        return name in self._features


# Module-level singleton
registry = FeatureRegistry()