core/utils/stats_engine.py

"""Stats Engine: Pure statistical calculation module.

This module centralizes all statistical helpers and aggregators used across
the codebase (ScannerService, backtester, etc.). It re-exports the classic
bucket helpers from stats_helpers and adds new pure aggregation functions.

Design goals:
- TESTABILITY: Pure functions with deterministic outputs.
- SINGLE RESPONSIBILITY: Each function does one aggregation task.
- REUSABILITY: Can be called from any context without side effects.
"""

from __future__ import annotations

from typing import Any

import pandas as pd

# Re-export classic helpers from stats_helpers (avoid duplication)
from core.utils.stats_helpers import (
    avg_vol10,
    bucket_gap,
    bucket_premium_pct,
    bucket_price,
    bucket_range,
    bucket_rel_vol,
    gap_pct,
)

# =============================================================================
# Time utilities (used by time-bucketing helpers)
# =============================================================================


def time_to_minutes(t: str) -> int:
    """Convert HH:MM string to minutes since midnight."""
    parts = t.split(":")
    return int(parts[0]) * 60 + int(parts[1])


def minutes_to_time(m: int) -> str:
    """Convert minutes since midnight to HH:MM string."""
    hour = int(m // 60)
    minute = int(m % 60)
    return f"{hour:02d}:{minute:02d}"


# =============================================================================
# Time-series aggregation helpers
# =============================================================================


def aggregate_return_metrics(
    daily_df: pd.DataFrame,
    day_idx: int,
    open_price: float,
    close_price: float,
    prev_close_val: float | None,
    symbol: str,
    date_to_test: str,
) -> tuple[float | None, int | None, dict[str, Any] | None]:
    """Calculate return, green flag, and day-of-week contribution.

    Returns:
        (return_pct, green_flag, dow_contribution)
    """
    if open_price is None or close_price is None or open_price <= 0:
        return None, None, None

    ret_pct = ((close_price - open_price) / open_price) * 100
    green = 1 if close_price > open_price else 0

    dow_contribution = None
    try:
        from datetime import datetime

        date_obj = datetime.strptime(date_to_test, "%Y-%m-%d")
        day_map = {0: "mon", 1: "tue", 2: "wed", 3: "thu", 4: "fri", 5: "sat", 6: "sun"}
        dow = day_map.get(date_obj.weekday())
        if dow:
            dow_contribution = {
                "dow": dow,
                "return": float(ret_pct),
                "green": green,
                "symbol": symbol,
                "date": date_to_test,
            }
    except Exception:
        pass

    return float(ret_pct), green, dow_contribution


def aggregate_range_metrics(
    open_price: float,
    high_price: float,
    low_price: float,
    symbol: str,
    date_to_test: str,
) -> dict[str, Any]:
    """Calculate open-high, high-low, and open-low range percentages."""
    if open_price is None or high_price is None or low_price is None or open_price <= 0:
        return {}

    o2h = ((high_price - open_price) / open_price) * 100
    l2h = ((high_price - low_price) / low_price) * 100 if low_price > 0 else 0.0
    o2l = ((low_price - open_price) / open_price) * 100

    example_info = {"symbol": symbol, "date": date_to_test}

    return {
        "o2h": float(o2h),
        "l2h": float(l2h),
        "o2l": float(o2l),
        "o2h_bucket": bucket_range(o2h),
        "l2h_bucket": bucket_range(l2h),
        "o2l_bucket": bucket_range(o2l),
        "o2h_example": {**example_info, "range_pct": round(o2h, 2)} if bucket_range(o2h) else None,
        "l2h_example": {**example_info, "range_pct": round(l2h, 2)} if bucket_range(l2h) else None,
        "o2l_example": {**example_info, "range_pct": round(o2l, 2)} if bucket_range(o2l) else None,
    }


def aggregate_volume_metrics(
    daily_df: pd.DataFrame,
    day_idx: int,
    current_vol: int | None,
    symbol: str,
) -> dict[str, Any]:
    """Calculate relative volume and bucket."""

    avg_vol_val = avg_vol10(daily_df, day_idx)
    if avg_vol_val is None or avg_vol_val <= 0 or current_vol is None:
        return {"rel_vol": None, "bucket": None}

    rel_vol = float(current_vol) / float(avg_vol_val)
    bucket = bucket_rel_vol(rel_vol)
    return {"rel_vol": rel_vol, "bucket": bucket}


def aggregate_gap_metrics(
    open_price: float | None,
    prev_close_val: float | None,
    close_price: float | None,
    symbol: str,
    date_to_test: str,
) -> dict[str, Any]:
    """Calculate gap percentage and bucket, plus trade return for gap analysis."""

    if prev_close_val is None or open_price is None or open_price <= 0:
        return {}

    gap_pct_val = gap_pct(prev_close_val, open_price)
    bucket = bucket_gap(gap_pct_val)

    trade_return = None
    if close_price is not None and open_price is not None and open_price > 0:
        trade_return = ((close_price - open_price) / open_price) * 100

    example_info = {"symbol": symbol, "date": date_to_test}
    example = None
    if bucket and trade_return is not None:
        example = {**example_info, "gap_pct": round(gap_pct_val, 2), "return": round(trade_return, 2)}

    return {"gap_pct": gap_pct_val, "bucket": bucket, "trade_return": trade_return, "example": example}


def aggregate_price_bucket(
    open_price: float | None,
    close_price: float | None,
    symbol: str,
    date_to_test: str,
) -> dict[str, Any] | None:
    """Get price bucket with example."""
    bucket = bucket_price(open_price)
    if bucket is None:
        return None

    ret_pct = None
    if open_price is not None and close_price is not None and open_price > 0:
        ret_pct = ((close_price - open_price) / open_price) * 100

    example_info = {"symbol": symbol, "date": date_to_test}
    return {
        "bucket": bucket,
        "example": {
            **example_info,
            "price": round(open_price, 2),
            "return": round(ret_pct, 2) if ret_pct is not None else 0.0,
        },
    }


def aggregate_sector_info(
    metadata_df: pd.DataFrame,
    symbol: str,
    open_price: float | None,
    close_price: float | None,
    date_to_test: str,
) -> tuple[str, float] | None:
    """Extract sector and calculate return for sector aggregation."""

    if open_price is None or close_price is None or open_price <= 0:
        return None

    sector = None
    meta = metadata_df[metadata_df["symbol"] == symbol]
    if not meta.empty:
        s = meta.iloc[0].get("sector")
        if s is not None and not pd.isna(s):
            sector = str(s)

    if not sector:
        return None

    trade_return = ((close_price - open_price) / open_price) * 100
    return sector, float(trade_return)


def aggregate_premarket_metrics(
    minute_data: pd.DataFrame,
    prev_close_val: float | None,
    symbol: str,
    date_to_test: str,
) -> dict[str, Any]:
    """Calculate all premarket-related metrics from minute data."""
    import pytz

    ny_tz = pytz.timezone("America/New_York")

    if minute_data is None or minute_data.empty:
        return {}

    # Ensure timezone-aware
    if minute_data.index.tz is None:
        minute_data.index = minute_data.index.tz_localize(pytz.UTC)
    minute_data.index = minute_data.index.tz_convert(ny_tz)

    highs = minute_data["high"].tolist()
    lows = minute_data["low"].tolist()

    pm_high = pm_low = None
    hod = lod = None

    for i, time_str in enumerate(minute_data.index.strftime("%H:%M").tolist()):
        if i >= len(highs) or i >= len(lows):
            continue
        h, low_val = highs[i], lows[i]
        if time_str < "09:30":
            if pm_high is None or h > pm_high:
                pm_high = h
            if pm_low is None or low_val < pm_low:
                pm_low = low_val
        else:
            if hod is None or h > hod:
                hod = h
            if lod is None or low_val < lod:
                lod = low_val

    result: dict[str, Any] = {
        "pm_high": pm_high,
        "pm_low": pm_low,
        "day_high": hod,
        "day_low": lod,
    }

    # Premark high/low as % of day high/low
    if pm_high is not None and hod is not None and hod > 0:
        result["pm_high_pct"] = (pm_high / hod) * 100
        result["pm_high_bucket"] = bucket_premium_pct(pm_high, hod)
    if pm_low is not None and lod is not None and lod > 0:
        result["pm_low_pct"] = (pm_low / lod) * 100
        result["pm_low_bucket"] = bucket_premium_pct(pm_low, lod)

    # Premarket range as % of prev close
    if pm_high is not None and pm_low is not None and prev_close_val and prev_close_val > 0:
        result["pm_range_pct"] = ((pm_high - pm_low) / prev_close_val) * 100

    return result


def aggregate_breakout_times(
    minute_data: pd.DataFrame,
    pm_high: float | None,
    pm_low: float | None,
    symbol: str,
    date_to_test: str,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    """Find first premarket breakout and breakdown times after 09:30."""
    if minute_data is None or minute_data.empty:
        return [], []

    ny_times = minute_data.index.strftime("%H:%M").tolist()
    closes = minute_data["close"].tolist()
    lows = minute_data["low"].tolist()

    breakout_high = []
    breakout_low = []

    for i, time_str in enumerate(ny_times):
        if time_str < "09:30":
            continue
        if i >= len(closes) or i >= len(lows):
            break

        price = closes[i]
        low_val = lows[i]

        if pm_high is not None and price > pm_high:
            breakout_high.append(
                {"time": time_str, "symbol": symbol, "date": date_to_test, "pm_high": round(pm_high, 2)}
            )
            break

        if pm_low is not None and low_val < pm_low:
            breakout_low.append({"time": time_str, "symbol": symbol, "date": date_to_test, "pm_low": round(pm_low, 2)})
            break

    return breakout_high, breakout_low


def aggregate_hod_lod_times(
    minute_data: pd.DataFrame,
    symbol: str,
    date_to_test: str,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    """Find high-of-day (HOD) and low-of-day (LOD) times."""
    if minute_data is None or minute_data.empty:
        return [], []

    hod_times = []
    lod_times = []

    hod_idx = minute_data["high"].idxmax()
    lod_idx = minute_data["low"].idxmin()

    hod_time = hod_idx.strftime("%H:%M") if hod_idx else None
    lod_time = lod_idx.strftime("%H:%M") if lod_idx else None

    if hod_time:
        hod_times.append({"time": hod_time, "symbol": symbol, "date": date_to_test})
    if lod_time:
        lod_times.append({"time": lod_time, "symbol": symbol, "date": date_to_test})

    return hod_times, lod_times


def aggregate_volume_by_time(
    minute_data: pd.DataFrame,
    symbol: str,
) -> dict[str, list[int]]:
    """Aggregate volume into 30-minute buckets (from 04:00 onwards)."""
    if minute_data is None or minute_data.empty:
        return {}

    volume_by_time: dict[str, list[int]] = {}
    ny_times = minute_data.index.strftime("%H:%M").tolist()
    volumes = minute_data["volume"].tolist()

    for i, time_str in enumerate(ny_times):
        try:
            hour, minute = int(time_str.split(":")[0]), int(time_str.split(":")[1])
            if hour >= 4:  # Include premarket from 4 AM
                bucket_hour = hour if minute < 30 else hour + (1 if minute >= 30 else 0)
                bucket_min = "00" if minute < 30 else "30"
                bucket_key = f"{bucket_hour:02d}:{bucket_min}"
                if bucket_key not in volume_by_time:
                    volume_by_time[bucket_key] = []
                if i < len(volumes):
                    volume_by_time[bucket_key].append(volumes[i])
        except Exception:
            pass

    return volume_by_time


# =============================================================================
# Distribution aggregation helpers
# =============================================================================


def bucket_time(times: list[dict[str, Any]]) -> tuple[dict[str, int], dict[str, list[dict[str, Any]]]]:
    """Bucket times into 30-minute intervals (09:30-16:00).

    Returns:
        (buckets_dict, examples_dict)
    """
    buckets: dict[str, int] = {}
    examples: dict[str, list[dict[str, Any]]] = {}

    for item in times:
        t = item.get("time") if isinstance(item, dict) else item
        try:
            hour, minute = int(t.split(":")[0]), int(t.split(":")[1])
            if 9 <= hour < 16:
                bucket_hour = hour
                bucket_min = "00" if minute < 30 else "30"
                bucket = f"{bucket_hour:02d}:{bucket_min}"
            else:
                bucket = f"{hour:02d}:{minute:02d}"

            buckets[bucket] = buckets.get(bucket, 0) + 1
            if bucket not in examples:
                examples[bucket] = []
            if len(examples[bucket]) < 5 and isinstance(item, dict):
                examples[bucket].append(item)
        except Exception:
            pass

    return buckets, examples


def calculate_mean_time(times: list[dict[str, Any]]) -> str | None:
    """Calculate the mean of all timestamps."""
    if not times:
        return None
    minutes = [time_to_minutes(t.get("time") if isinstance(t, dict) else t) for t in times]
    mean_min = sum(minutes) / len(minutes)
    return minutes_to_time(int(mean_min))


def calculate_mode_time(distribution: dict[str, int]) -> str | None:
    """Return the bucket with the highest count."""
    if not distribution:
        return None
    return max(distribution.keys(), key=lambda b: distribution[b])


def calculate_median_time(times: list[dict[str, Any]]) -> str | None:
    """Calculate the median of all timestamps (not buckets)."""
    if not times:
        return None
    minutes = sorted([time_to_minutes(t.get("time") if isinstance(t, dict) else t) for t in times])
    n = len(minutes)
    mid = n // 2
    median_min = (minutes[mid - 1] + minutes[mid]) / 2 if n % 2 == 0 else minutes[mid]
    return minutes_to_time(int(median_min))


def aggregate_range_buckets(buckets_dict: dict[str, list[dict[str, Any]]]) -> dict[str, dict[str, Any]]:
    """Aggregate range bucket items into summary statistics."""
    result = {}
    for bucket, items in buckets_dict.items():
        if items:
            ranges = [item.get("range_pct", 0) for item in items]
            result[bucket] = {
                "count": len(items),
                "avg_range": round(sum(ranges) / len(ranges), 2),
                "value": round(sum(ranges) / len(ranges), 2),  # retained for backward compatibility
                "min_range": round(min(ranges), 2) if ranges else 0,
                "max_range": round(max(ranges), 2) if ranges else 0,
                "examples": items[:5],
            }
    return result


def aggregate_bucket_stats(
    buckets_dict: dict[str, list[dict[str, Any]]],
    value_key: str,
    count_key: str = "count",
    avg_key: str = "avg",
    examples_limit: int = 5,
) -> dict[str, dict[str, Any]]:
    """Generic aggregation for any bucketed data."""
    result = {}
    for bucket, items in buckets_dict.items():
        if items:
            values = [item.get(value_key, 0) for item in items]
            result[bucket] = {
                count_key: len(items),
                avg_key: round(sum(values) / len(values), 1),
                "value": round(sum(values) / len(values), 1),  # backward compatible alias
                "min_" + value_key: round(min(values), 1) if values else 0,
                "max_" + value_key: round(max(values), 1) if values else 0,
                "examples": items[:examples_limit],
            }
    return result


def aggregate_volume_distribution(volume_by_time: dict[str, list[int]]) -> dict[str, dict[str, Any]]:
    """Aggregate volume by time bucket."""
    result = {}
    for bucket, vols in volume_by_time.items():
        if vols:
            result[bucket] = {
                "total": sum(vols),
                "avg": round(sum(vols) / len(vols), 0),
                "count": len(vols),
            }
    return result


def aggregate_price_buckets(price_buckets: dict[str, list[dict[str, Any]]]) -> dict[str, dict[str, Any]]:
    """Aggregate price bucket items."""
    result = {}
    for bucket, items in price_buckets.items():
        if items:
            returns_list = [item.get("return", 0) for item in items]
            prices_list = [item.get("price", 0) for item in items]
            result[bucket] = {
                "count": len(items),
                "avg_return": round(sum(returns_list) / len(returns_list), 2) if returns_list else 0,
                "avg_price": round(sum(prices_list) / len(prices_list), 2) if prices_list else 0,
                "examples": items[:5],
            }
    return result


def aggregate_rel_vol_buckets(rel_vol_buckets: dict[str, list[dict[str, Any]]]) -> dict[str, dict[str, Any]]:
    """Aggregate relative volume buckets."""
    result = {}
    for bucket, items in rel_vol_buckets.items():
        if items:
            rel_vols_list = [item.get("rel_vol", 0) for item in items]
            avg = sum(rel_vols_list) / len(rel_vols_list)
            result[bucket] = {
                "count": len(items),
                "avg_rel_vol": round(avg, 2),
                "value": round(avg, 2),
                "examples": items[:5],
            }
    return result


def aggregate_premarket_buckets(
    buckets_dict: dict[str, list[dict[str, Any]]],
    value_key: str,
) -> dict[str, dict[str, Any]]:
    """Aggregate premarket bucket stats."""
    result = {}
    for bucket, items in buckets_dict.items():
        if items:
            pcts = [item.get(value_key, 0) for item in items]
            result[bucket] = {
                "count": len(items),
                "avg_pct": round(sum(pcts) / len(pcts), 1),
                "value": round(sum(pcts) / len(pcts), 1),
                "min_pct": round(min(pcts), 1) if pcts else 0,
                "max_pct": round(max(pcts), 1) if pcts else 0,
                "examples": items[:5],
            }
    return result


def aggregate_day_of_week(
    days_of_week: dict[str, list[dict[str, Any]]],
) -> dict[str, dict[str, Any]]:
    """Aggregate day-of-week statistics."""
    dow_stats = {}
    for day, trades in days_of_week.items():
        if trades:
            returns_day = [t["return"] for t in trades]
            greens_day = [t["green"] for t in trades]
            examples = [{"symbol": t["symbol"], "date": t["date"], "return": t["return"]} for t in trades[:3]]
            dow_stats[day] = {
                "count": len(trades),
                "avg_return": round(sum(returns_day) / len(returns_day), 2),
                "close_green_pct": round(sum(greens_day) / len(greens_day) * 100, 1),
                "examples": examples,
            }
    return dow_stats


def aggregate_gap_analysis(gap_buckets: dict[str, dict[str, Any]]) -> dict[str, dict[str, Any]]:
    """Aggregate gap analysis from bucket data."""
    result = {}
    for bucket_name, data in gap_buckets.items():
        if data.get("count", 0) > 0:
            returns = data.get("returns", [])
            result[bucket_name] = {
                "count": data["count"],
                "avg_return": round(sum(returns) / len(returns), 2) if returns else 0,
                "examples": data.get("examples", []),
            }
    return result


def aggregate_sector_performance(sector_perf: dict[str, dict[str, Any]]) -> dict[str, dict[str, Any]]:
    """Aggregate sector performance."""
    result = {}
    for sector_name, data in sector_perf.items():
        returns = data.get("returns", [])
        if returns:
            result[sector_name] = {
                "count": data.get("count", 0),
                "avg_return": round(sum(returns) / len(returns), 2),
                "examples": data.get("examples", []),
            }
    return result


# =============================================================================
# Candidate-level metric aggregator (ScannerService integration)
# =============================================================================


def aggregate_candidate_metrics(
    daily_df: pd.DataFrame,
    day_idx: int,
    candidate: dict[str, Any],
    prev_close_val: float | None,
    minute_data: pd.DataFrame | None,
    metadata_df: pd.DataFrame | None,
) -> dict[str, Any]:
    """Calculate all per-candidate metrics using the stats engine.

    This is the single source of truth for per-candidate statistics.
    ScannerService should populate the data structures (bucket lists) by
    calling this for each candidate and then passing the results to the
    aggregation functions below.

    Args:
        daily_df: DataFrame with OHLCV data for the symbol
        day_idx: Integer index of the target day in daily_df
        candidate: Dict with at least 'symbol' and 'date' keys
        prev_close_val: Previous day's close (for gap calculations)
        minute_data: Optional minute-level DataFrame for the date
        metadata_df: Optional DataFrame with symbol metadata (sector)

    Returns:
        Dict with all computed metrics and bucket assignments.
    """
    symbol = candidate.get("symbol", "").upper()
    original_date = candidate.get("date", "")

    # Extract prices
    open_price = daily_df["open"].iloc[day_idx] if "open" in daily_df.columns and day_idx < len(daily_df) else None
    close_price = daily_df["close"].iloc[day_idx] if "close" in daily_df.columns and day_idx < len(daily_df) else None
    high_price = daily_df["high"].iloc[day_idx] if "high" in daily_df.columns and day_idx < len(daily_df) else None
    low_price = daily_df["low"].iloc[day_idx] if "low" in daily_df.columns and day_idx < len(daily_df) else None
    current_vol = (
        daily_df["volume"].iloc[day_idx]
        if "volume" in daily_df.columns and day_idx < len(daily_df) and pd.notna(daily_df["volume"].iloc[day_idx])
        else None
    )

    # 1. Return metrics
    ret_pct, green, dow_contribution = aggregate_return_metrics(
        daily_df, day_idx, open_price, close_price, prev_close_val, symbol, original_date
    )

    # 2. Volume metrics
    vol_metrics = aggregate_volume_metrics(daily_df, day_idx, current_vol, symbol)

    # 3. Gap metrics
    gap_metrics = aggregate_gap_metrics(open_price, prev_close_val, close_price, symbol, original_date)

    # 4. Price bucket
    price_bucket_info = aggregate_price_bucket(open_price, close_price, symbol, original_date)

    # 5. Range metrics
    range_metrics = aggregate_range_metrics(open_price, high_price, low_price, symbol, original_date)

    # 6. Sector info
    sector_name = None
    sector_return = None
    if metadata_df is not None and not metadata_df.empty:
        sector_result = aggregate_sector_info(metadata_df, symbol, open_price, close_price, original_date)
        if sector_result:
            sector_name, sector_return = sector_result

    # 7. Premarket metrics (requires prev_close)
    pm_metrics = (
        aggregate_premarket_metrics(minute_data, prev_close_val, symbol, original_date)
        if minute_data is not None
        else {}
    )

    # 8. HOD/LOD times
    hod_times, lod_times = (
        aggregate_hod_lod_times(minute_data, symbol, original_date) if minute_data is not None else ([], [])
    )

    # 9. Volume by time (30-min buckets)
    vol_by_time = aggregate_volume_by_time(minute_data, symbol) if minute_data is not None else {}

    # 10. Breakout times (uses pm_high/pm_low from pm_metrics)
    pm_breakout_high = []
    pm_breakout_low = []
    if minute_data is not None and "pm_high" in pm_metrics and "pm_low" in pm_metrics:
        pm_breakout_high, pm_breakout_low = aggregate_breakout_times(
            minute_data, pm_metrics["pm_high"], pm_metrics["pm_low"], symbol, original_date
        )

    return {
        "symbol": symbol,
        "date": original_date,
        "ret_pct": ret_pct,
        "green": green,
        "dow_contribution": dow_contribution,
        "rel_vol": vol_metrics.get("rel_vol"),
        "rel_vol_bucket": vol_metrics.get("bucket"),
        "gap_pct": gap_metrics.get("gap_pct"),
        "gap_bucket": gap_metrics.get("bucket"),
        "trade_return": gap_metrics.get("trade_return"),
        "gap_example": gap_metrics.get("example"),
        "price_bucket": price_bucket_info.get("bucket") if price_bucket_info else None,
        "price_example": price_bucket_info.get("example") if price_bucket_info else None,
        "range_metrics": range_metrics,
        "sector": sector_name,
        "sector_return": sector_return,
        "pm_metrics": pm_metrics,
        "hod_times": hod_times,
        "lod_times": lod_times,
        "volume_by_time": vol_by_time,
        "pm_breakout_high": pm_breakout_high,
        "pm_breakout_low": pm_breakout_low,
    }