stonks-oracle/services/intelligence_pipeline_v3/impact/integration.py

"""Impact output integration — connects impact predictions to legacy consumers.

Provides the ImpactPrediction Pydantic model, compatibility adapter mapping,
feature flag for v3 mode, and comparison metrics placeholder.

Design reference: Section I & K in design.md.
Requirement 12.1, 12.7, 12.9.
"""

from __future__ import annotations

import logging
import os
from datetime import datetime, timezone
from typing import Literal

from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Feature flags
# ---------------------------------------------------------------------------


class ImpactPipelineConfig(BaseModel):
    """Configuration for the impact pipeline integration.

    Controls whether generative impact/novelty/confidence are removed
    from aggregation inputs when v3 mode is active.
    """

    v3_mode_enabled: bool = Field(
        default=False,
        description="When True, removes generative impact/novelty/confidence from aggregation inputs.",
    )
    use_trained_model: bool = Field(
        default=False,
        description="When True, uses trained model if approved. Otherwise uses deterministic baseline.",
    )
    legacy_compatibility: bool = Field(
        default=True,
        description="When True, maps impact predictions to legacy impact_score/impact_horizon.",
    )
    comparison_metrics_enabled: bool = Field(
        default=False,
        description="When True, stores comparison metrics between v3 and generative predictions.",
    )


def get_impact_config() -> ImpactPipelineConfig:
    """Get impact pipeline configuration from environment."""
    return ImpactPipelineConfig(
        v3_mode_enabled=os.environ.get("IMPACT_V3_MODE_ENABLED", "false").lower() == "true",
        use_trained_model=os.environ.get("IMPACT_USE_TRAINED_MODEL", "false").lower() == "true",
        legacy_compatibility=os.environ.get("IMPACT_LEGACY_COMPATIBILITY", "true").lower() == "true",
        comparison_metrics_enabled=os.environ.get("IMPACT_COMPARISON_METRICS", "false").lower() == "true",
    )


# ---------------------------------------------------------------------------
# Impact prediction output model (Pydantic)
# ---------------------------------------------------------------------------


class DirectionProbabilities(BaseModel):
    """Probability distribution over market direction outcomes."""

    positive: float = Field(ge=0.0, le=1.0, default=0.0)
    negative: float = Field(ge=0.0, le=1.0, default=0.0)
    neutral: float = Field(ge=0.0, le=1.0, default=0.0)


class HorizonProbabilities(BaseModel):
    """Probability distribution over impact horizons."""

    intraday: float = Field(ge=0.0, le=1.0, default=0.0)
    one_day: float = Field(ge=0.0, le=1.0, default=0.0)
    seven_day: float = Field(ge=0.0, le=1.0, default=0.0)
    thirty_day: float = Field(ge=0.0, le=1.0, default=0.0)
    ninety_day: float = Field(ge=0.0, le=1.0, default=0.0)


class ImpactPredictionOutput(BaseModel):
    """Complete impact prediction output for persistence and downstream use.

    Contains the full probability distributions, not just point estimates.
    This is richer than the legacy scalar fields.
    """

    direction_probs: DirectionProbabilities = Field(
        description="Probability distribution over market direction.",
    )
    expected_magnitude: float = Field(
        ge=0.0,
        description="Expected absolute magnitude of market response.",
    )
    signed_magnitude: float = Field(
        description="Direction-weighted expected magnitude.",
    )
    horizon_probs: HorizonProbabilities = Field(
        description="Probability distribution over response horizons.",
    )
    uncertainty: float = Field(
        ge=0.0,
        le=1.0,
        description="Model uncertainty (higher = less confident).",
    )
    model_source: str = Field(
        description="Which model produced this prediction.",
    )
    feature_snapshot_id: str | None = Field(
        default=None,
        description="ID of the immutable feature snapshot used for this prediction.",
    )
    prediction_time: datetime = Field(
        default_factory=lambda: datetime.now(tz=timezone.utc),
    )


# ---------------------------------------------------------------------------
# Legacy compatibility adapter
# ---------------------------------------------------------------------------


class LegacyImpactMapping(BaseModel):
    """Mapping from v3 impact prediction to legacy impact_score/impact_horizon."""

    impact_score: float = Field(
        ge=-1.0,
        le=1.0,
        description="Legacy impact score mapped from v3 signed magnitude.",
    )
    impact_horizon: Literal["intraday", "1d", "7d", "30d", "90d"] = Field(
        description="Legacy horizon mapped from most probable v3 horizon.",
    )
    mapping_version: str = "1.0.0"


def map_to_legacy_impact(prediction: ImpactPredictionOutput) -> LegacyImpactMapping:
    """Map a v3 impact prediction to legacy impact_score and impact_horizon.

    Parameters
    ----------
    prediction
        The full v3 impact prediction output.

    Returns
    -------
    LegacyImpactMapping
        Legacy-compatible fields for downstream consumers.
    """
    # impact_score: clamp signed_magnitude to [-1, 1]
    impact_score = max(-1.0, min(1.0, prediction.signed_magnitude))

    # impact_horizon: argmax of horizon probabilities
    horizon_map: dict[str, float] = {
        "intraday": prediction.horizon_probs.intraday,
        "1d": prediction.horizon_probs.one_day,
        "7d": prediction.horizon_probs.seven_day,
        "30d": prediction.horizon_probs.thirty_day,
        "90d": prediction.horizon_probs.ninety_day,
    }
    impact_horizon = max(horizon_map, key=lambda k: horizon_map[k])

    return LegacyImpactMapping(
        impact_score=impact_score,
        impact_horizon=impact_horizon,
    )


# ---------------------------------------------------------------------------
# V3 mode signal filtering
# ---------------------------------------------------------------------------


def filter_generative_scores(
    signal_dict: dict,
    config: ImpactPipelineConfig | None = None,
) -> dict:
    """Remove generative impact/novelty/confidence from aggregation inputs in v3 mode.

    When v3 mode is enabled, these fields are replaced by calibrated v3 values.
    The original generative values are removed to prevent double-counting.

    Parameters
    ----------
    signal_dict
        Dictionary of signal fields from the extraction pipeline.
    config
        Pipeline configuration. Uses env-based default if None.

    Returns
    -------
    dict
        Signal dict with generative scores removed if v3 mode is active.
    """
    if config is None:
        config = get_impact_config()

    if not config.v3_mode_enabled:
        return signal_dict

    # Fields produced by generative model that v3 replaces
    generative_fields = {
        "impact_score",
        "impact_horizon",
        "novelty_score",
        "confidence",
    }

    filtered = {k: v for k, v in signal_dict.items() if k not in generative_fields}

    logger.debug(
        "V3 mode: removed generative fields %s from signal",
        generative_fields & set(signal_dict.keys()),
    )

    return filtered


# ---------------------------------------------------------------------------
# Comparison metrics (placeholder for dashboard integration)
# ---------------------------------------------------------------------------


class ComparisonMetric(BaseModel):
    """Single comparison data point between v3 prediction and realized outcome."""

    ticker: str
    event_time: datetime
    prediction_source: str
    predicted_direction: str
    predicted_magnitude: float
    predicted_horizon: str
    realized_return_1d: float | None = None
    realized_return_7d: float | None = None
    realized_return_30d: float | None = None
    direction_correct: bool | None = None
    magnitude_error: float | None = None


# In-memory store for comparison metrics (production would use database)
_COMPARISON_METRICS: list[ComparisonMetric] = []


def record_comparison_metric(metric: ComparisonMetric) -> None:
    """Record a comparison metric for later dashboard display.

    Only records if comparison metrics are enabled in config.
    """
    config = get_impact_config()
    if not config.comparison_metrics_enabled:
        return
    _COMPARISON_METRICS.append(metric)


def get_comparison_metrics(
    ticker: str | None = None,
    limit: int = 100,
) -> list[ComparisonMetric]:
    """Retrieve stored comparison metrics, optionally filtered by ticker."""
    metrics = _COMPARISON_METRICS
    if ticker:
        metrics = [m for m in metrics if m.ticker == ticker]
    return metrics[:limit]


def clear_comparison_metrics() -> None:
    """Clear all stored comparison metrics (for testing only)."""
    _COMPARISON_METRICS.clear()