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

"""Outcome label generation for impact model training.

Computes leakage-safe abnormal returns and response labels at defined
event timestamps over multiple horizons.

Design reference: Section I (Impact and Horizon Model) — Labels.
Requirement 12.3.
"""

from __future__ import annotations

import math
from datetime import datetime, timedelta
from typing import Literal

from pydantic import BaseModel, Field

# Version label-generation code for reproducibility tracking
LABEL_GENERATOR_VERSION = "1.0.0"


# ---------------------------------------------------------------------------
# Types
# ---------------------------------------------------------------------------

HorizonName = Literal["intraday", "1d", "7d", "30d", "90d"]

HORIZON_DURATIONS: dict[HorizonName, timedelta] = {
    "intraday": timedelta(hours=6, minutes=30),  # Trading day approximation
    "1d": timedelta(days=1),
    "7d": timedelta(days=7),
    "30d": timedelta(days=30),
    "90d": timedelta(days=90),
}


class OutcomeLabel(BaseModel):
    """Single-horizon outcome label for a given event."""

    horizon: HorizonName
    signed_return: float = Field(
        description="Signed abnormal return over the horizon window.",
    )
    absolute_return: float = Field(
        ge=0.0,
        description="Absolute abnormal return over the horizon window.",
    )
    abnormal_volume: float | None = Field(
        default=None,
        description="Volume ratio vs trailing average (None if data unavailable).",
    )
    time_to_peak_hours: float | None = Field(
        default=None,
        description="Hours from event to peak response within the horizon (None if unavailable).",
    )
    data_quality: Literal["full", "partial", "insufficient"] = Field(
        default="full",
        description="Quality indicator for this label's underlying data.",
    )


class OutcomeLabelSet(BaseModel):
    """Complete label set for all horizons at a single event."""

    event_time: datetime
    ticker: str
    benchmark_ticker: str = "SPY"
    labels: list[OutcomeLabel] = Field(default_factory=list)
    label_generator_version: str = LABEL_GENERATOR_VERSION
    market_data_snapshot_id: str | None = Field(
        default=None,
        description="Reference to the market data snapshot used for label generation.",
    )


# ---------------------------------------------------------------------------
# Core computation
# ---------------------------------------------------------------------------


def compute_abnormal_return(
    price_series: list[tuple[datetime, float]],
    benchmark_series: list[tuple[datetime, float]],
    event_time: datetime,
    horizon: timedelta,
) -> float:
    """Compute abnormal return of the asset relative to benchmark over a horizon.

    Abnormal return = asset_return - benchmark_return

    Parameters
    ----------
    price_series
        Sorted list of (timestamp, price) tuples for the asset.
    benchmark_series
        Sorted list of (timestamp, price) tuples for the benchmark.
    event_time
        When the event occurred (start of measurement window).
    horizon
        Duration of the measurement window.

    Returns
    -------
    float
        Abnormal return as a fraction (e.g., 0.02 = 2%).

    Raises
    ------
    ValueError
        If series are empty or don't cover the required time range.
    """
    if not price_series:
        raise ValueError("price_series is empty")
    if not benchmark_series:
        raise ValueError("benchmark_series is empty")

    end_time = event_time + horizon

    asset_start = _get_price_at_or_before(price_series, event_time)
    asset_end = _get_price_at_or_before(price_series, end_time)
    bench_start = _get_price_at_or_before(benchmark_series, event_time)
    bench_end = _get_price_at_or_before(benchmark_series, end_time)

    if asset_start is None or asset_end is None:
        raise ValueError(
            f"Asset price series does not cover event_time to event_time+horizon "
            f"({event_time.isoformat()} to {end_time.isoformat()})"
        )
    if bench_start is None or bench_end is None:
        raise ValueError(
            f"Benchmark series does not cover event_time to event_time+horizon "
            f"({event_time.isoformat()} to {end_time.isoformat()})"
        )

    if asset_start == 0.0 or bench_start == 0.0:
        raise ValueError("Start price cannot be zero")

    asset_return = (asset_end - asset_start) / asset_start
    bench_return = (bench_end - bench_start) / bench_start

    return asset_return - bench_return


def compute_abnormal_volume(
    volume_series: list[tuple[datetime, float]],
    event_time: datetime,
    horizon: timedelta,
    lookback_days: int = 20,
) -> float | None:
    """Compute abnormal volume ratio relative to trailing average.

    Parameters
    ----------
    volume_series
        Sorted list of (timestamp, volume) tuples.
    event_time
        When the event occurred.
    horizon
        Duration window to measure event-period volume.
    lookback_days
        Number of days before event_time to compute trailing average.

    Returns
    -------
    float or None
        Volume ratio (event_volume / trailing_avg_volume), or None if insufficient data.
    """
    if not volume_series:
        return None

    lookback_start = event_time - timedelta(days=lookback_days)
    end_time = event_time + horizon

    # Trailing volume (pre-event)
    trailing_volumes = [
        v for ts, v in volume_series
        if lookback_start <= ts < event_time
    ]

    # Event-period volume
    event_volumes = [
        v for ts, v in volume_series
        if event_time <= ts <= end_time
    ]

    if not trailing_volumes or not event_volumes:
        return None

    trailing_avg = sum(trailing_volumes) / len(trailing_volumes)
    if trailing_avg == 0:
        return None

    event_avg = sum(event_volumes) / len(event_volumes)
    return event_avg / trailing_avg


def compute_time_to_peak(
    price_series: list[tuple[datetime, float]],
    event_time: datetime,
    horizon: timedelta,
) -> float | None:
    """Compute time from event to peak absolute response within horizon.

    Parameters
    ----------
    price_series
        Sorted list of (timestamp, price) tuples.
    event_time
        When the event occurred.
    horizon
        Duration window to search for peak.

    Returns
    -------
    float or None
        Hours from event to peak absolute deviation, or None if insufficient data.
    """
    if not price_series:
        return None

    end_time = event_time + horizon
    base_price = _get_price_at_or_before(price_series, event_time)
    if base_price is None or base_price == 0.0:
        return None

    # Find the point within [event_time, end_time] with max absolute deviation
    max_deviation = 0.0
    peak_time = event_time

    for ts, price in price_series:
        if ts < event_time:
            continue
        if ts > end_time:
            break
        deviation = abs((price - base_price) / base_price)
        if deviation > max_deviation:
            max_deviation = deviation
            peak_time = ts

    if max_deviation == 0.0:
        return None

    hours = (peak_time - event_time).total_seconds() / 3600.0
    return hours


# ---------------------------------------------------------------------------
# Label generation for all horizons
# ---------------------------------------------------------------------------


def generate_outcome_labels(
    ticker: str,
    event_time: datetime,
    price_series: list[tuple[datetime, float]],
    benchmark_series: list[tuple[datetime, float]],
    volume_series: list[tuple[datetime, float]] | None = None,
    benchmark_ticker: str = "SPY",
    horizons: list[HorizonName] | None = None,
    market_data_snapshot_id: str | None = None,
) -> OutcomeLabelSet:
    """Generate outcome labels for all configured horizons.

    Parameters
    ----------
    ticker
        Asset ticker symbol.
    event_time
        When the event was detected.
    price_series
        Asset price series (sorted by timestamp).
    benchmark_series
        Benchmark price series (sorted by timestamp).
    volume_series
        Optional volume series for abnormal volume labels.
    benchmark_ticker
        Benchmark identifier (default SPY).
    horizons
        Which horizons to compute. Default is all five.
    market_data_snapshot_id
        Optional reference to the market data snapshot used.

    Returns
    -------
    OutcomeLabelSet
        Complete label set for the event.
    """
    if horizons is None:
        horizons = list(HORIZON_DURATIONS.keys())

    labels: list[OutcomeLabel] = []

    for horizon_name in horizons:
        duration = HORIZON_DURATIONS[horizon_name]
        label = _compute_single_horizon_label(
            price_series=price_series,
            benchmark_series=benchmark_series,
            volume_series=volume_series,
            event_time=event_time,
            horizon_name=horizon_name,
            duration=duration,
        )
        labels.append(label)

    return OutcomeLabelSet(
        event_time=event_time,
        ticker=ticker,
        benchmark_ticker=benchmark_ticker,
        labels=labels,
        label_generator_version=LABEL_GENERATOR_VERSION,
        market_data_snapshot_id=market_data_snapshot_id,
    )


def _compute_single_horizon_label(
    price_series: list[tuple[datetime, float]],
    benchmark_series: list[tuple[datetime, float]],
    volume_series: list[tuple[datetime, float]] | None,
    event_time: datetime,
    horizon_name: HorizonName,
    duration: timedelta,
) -> OutcomeLabel:
    """Compute outcome label for a single horizon."""
    # Attempt abnormal return
    try:
        signed_return = compute_abnormal_return(
            price_series, benchmark_series, event_time, duration
        )
        data_quality: Literal["full", "partial", "insufficient"] = "full"
    except ValueError:
        signed_return = float("nan")
        data_quality = "insufficient"

    # Absolute return
    absolute_return = abs(signed_return) if not math.isnan(signed_return) else 0.0

    # Abnormal volume
    abnormal_volume = None
    if volume_series:
        abnormal_volume = compute_abnormal_volume(
            volume_series, event_time, duration
        )
        if abnormal_volume is None and data_quality == "full":
            data_quality = "partial"

    # Time to peak
    time_to_peak = None
    try:
        time_to_peak = compute_time_to_peak(price_series, event_time, duration)
    except (ValueError, ZeroDivisionError):
        pass
    if time_to_peak is None and data_quality == "full":
        data_quality = "partial"

    return OutcomeLabel(
        horizon=horizon_name,
        signed_return=signed_return if not math.isnan(signed_return) else 0.0,
        absolute_return=absolute_return,
        abnormal_volume=abnormal_volume,
        time_to_peak_hours=time_to_peak,
        data_quality=data_quality,
    )


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _get_price_at_or_before(
    series: list[tuple[datetime, float]], target: datetime
) -> float | None:
    """Get the most recent price at or before the target timestamp.

    Assumes series is sorted by timestamp ascending.
    """
    result = None
    for ts, price in series:
        if ts <= target:
            result = price
        else:
            break
    return result