tests/test_stats_engine.py

"""
Unit tests for core/utils/stats_engine.py

All functions are pure statistical calculations.  We test them with
deterministic synthetic data — no external dependencies, no DB, no network.
"""

import pandas as pd
import pytest

from core.utils.stats_engine import (
    aggregate_breakout_times,
    aggregate_bucket_stats,
    aggregate_day_of_week,
    aggregate_gap_analysis,
    aggregate_gap_metrics,
    aggregate_hod_lod_times,
    aggregate_premarket_buckets,
    aggregate_premarket_metrics,
    aggregate_price_bucket,
    aggregate_price_buckets,
    aggregate_range_buckets,
    aggregate_range_metrics,
    aggregate_rel_vol_buckets,
    aggregate_return_metrics,
    aggregate_sector_performance,
    aggregate_volume_by_time,
    aggregate_volume_distribution,
    aggregate_volume_metrics,
    bucket_gap,
    bucket_premium_pct,
    bucket_price,
    bucket_range,
    bucket_rel_vol,
    bucket_time,
    calculate_mean_time,
    calculate_median_time,
    calculate_mode_time,
    minutes_to_time,
    time_to_minutes,
)

# =============================================================================
# Time utilities
# =============================================================================


class TestTimeUtils:
    def test_time_to_minutes(self):
        assert time_to_minutes("00:00") == 0
        assert time_to_minutes("09:30") == 570
        assert time_to_minutes("16:00") == 960
        assert time_to_minutes("23:59") == 1439

    def test_minutes_to_time(self):
        assert minutes_to_time(0) == "00:00"
        assert minutes_to_time(570) == "09:30"
        assert minutes_to_time(960) == "16:00"
        assert minutes_to_time(1439) == "23:59"


# =============================================================================
# Bucket helpers
# =============================================================================


class TestBucketGap:
    @pytest.mark.parametrize(
        "value,expected",
        [
            (-10.0, "gap_n5_p"),
            (-5.0, "gap_n5_0"),
            (-0.1, "gap_n5_0"),
            (0.0, "gap_0_5"),
            (4.9, "gap_0_5"),
            (5.0, "gap_5_10"),
            (9.9, "gap_5_10"),
            (10.0, "gap_10_p"),
            (50.0, "gap_10_p"),
            (None, None),
        ],
    )
    def test_bucket_gap(self, value, expected):
        assert bucket_gap(value) == expected


class TestBucketRange:
    @pytest.mark.parametrize(
        "value,expected",
        [
            (0.0, "rng_0_2"),
            (1.9, "rng_0_2"),
            (2.0, "rng_2_5"),
            (4.9, "rng_2_5"),
            (5.0, "rng_5_10"),
            (9.9, "rng_5_10"),
            (10.0, "rng_10_p"),
            (None, None),
        ],
    )
    def test_bucket_range(self, value, expected):
        assert bucket_range(value) == expected


class TestBucketPrice:
    @pytest.mark.parametrize(
        "value,expected",
        [
            (0.0, "prc_lt_5"),
            (4.99, "prc_lt_5"),
            (5.0, "prc_5_20"),
            (19.99, "prc_5_20"),
            (20.0, "prc_20_50"),
            (49.99, "prc_20_50"),
            (50.0, "prc_50_100"),
            (99.99, "prc_50_100"),
            (100.0, "prc_gt_100"),
            (None, None),
        ],
    )
    def test_bucket_price(self, value, expected):
        assert bucket_price(value) == expected


class TestBucketRelVol:
    @pytest.mark.parametrize(
        "value,expected",
        [
            (1.0, "rvol_lt_2"),
            (1.9, "rvol_lt_2"),
            (2.0, "rvol_2_5"),
            (4.9, "rvol_2_5"),
            (5.0, "rvol_5_10"),
            (9.9, "rvol_5_10"),
            (10.0, "rvol_gt_10"),
            (None, None),
        ],
    )
    def test_bucket_rel_vol(self, value, expected):
        assert bucket_rel_vol(value) == expected


class TestBucketPremiumPct:
    @pytest.mark.parametrize(
        "pm,day,expected",
        [
            (40.0, 100.0, "pm_pct_0_50"),
            (50.0, 100.0, "pm_pct_50_75"),
            (60.0, 100.0, "pm_pct_50_75"),
            (80.0, 100.0, "pm_pct_75_90"),
            (90.0, 100.0, "pm_pct_90_100"),  # 90% should be in 90-100 bucket
            (95.0, 100.0, "pm_pct_90_100"),
            (99.0, 100.0, "pm_pct_90_100"),
            (100.0, 100.0, "pm_pct_100_p"),
            (110.0, 100.0, "pm_pct_100_p"),
            (None, 100.0, None),
            (50.0, None, None),
            (50.0, 0, None),
        ],
    )
    def test_bucket_premium_pct(self, pm, day, expected):
        assert bucket_premium_pct(pm, day) == expected


# =============================================================================
# Time bucketing
# =============================================================================


class TestBucketTime:
    def test_bucket_time_trading_hours(self):
        times = [
            {"time": "09:31", "symbol": "A"},
            {"time": "09:45", "symbol": "B"},
            {"time": "10:00", "symbol": "C"},
            {"time": "10:15", "symbol": "D"},
            {"time": "14:45", "symbol": "E"},
        ]
        buckets, examples = bucket_time(times)

        assert "09:30" in buckets
        assert buckets["09:30"] == 2  # 09:31 and 09:45
        assert "10:00" in buckets
        assert "14:30" in buckets

        # Examples: one per bucket max
        assert isinstance(examples, dict)
        assert len(examples) <= len(buckets)

    def test_bucket_time_outside_trading(self):
        times = [{"time": "08:00", "symbol": "A"}, {"time": "20:00", "symbol": "B"}]
        buckets, examples = bucket_time(times)

        # Outside 9-16, we store as HH:MM directly
        assert "08:00" in buckets or "20:00" in buckets


class TestTimeStats:
    def test_calculate_mean_time(self):
        times = [{"time": "09:00"}, {"time": "10:00"}, {"time": "11:00"}]
        assert calculate_mean_time(times) == "10:00"

    def test_calculate_mean_time_empty(self):
        assert calculate_mean_time([]) is None

    def test_calculate_median_time(self):
        times = [{"time": "09:00"}, {"time": "10:00"}, {"time": "11:00"}, {"time": "12:00"}]
        assert calculate_median_time(times) == "10:30"

    def test_calculate_median_time_odd(self):
        times = [{"time": "09:00"}, {"time": "10:00"}, {"time": "11:00"}]
        assert calculate_median_time(times) == "10:00"

    def test_calculate_median_time_empty(self):
        assert calculate_median_time([]) is None

    def test_calculate_mode_time(self):
        distribution = {"09:30": 3, "10:00": 5, "10:30": 2}
        assert calculate_mode_time(distribution) == "10:00"

    def test_calculate_mode_time_empty(self):
        assert calculate_mode_time({}) is None


# =============================================================================
# Return metrics
# =============================================================================


class TestReturnMetrics:
    def test_aggregate_return_basic(self):
        daily_df = pd.DataFrame(
            {
                "timestamp": pd.to_datetime(["2025-01-02", "2025-01-03"]),
                "open": [100.0, 101.0],
                "close": [105.0, 102.0],
                "volume": [1_000_000, 1_100_000],
            }
        )
        ret, green, dow = aggregate_return_metrics(daily_df, 1, 101.0, 102.0, 100.0, "AAPL", "2025-01-03")
        assert ret == pytest.approx(0.990099, rel=1e-4)
        assert green == 1
        assert dow is not None
        assert dow["dow"] in ["mon", "tue", "wed", "thu", "fri"]

    def test_aggregate_return_invalid_prices(self):
        daily_df = pd.DataFrame({"timestamp": pd.to_datetime(["2025-01-02"]), "open": [0], "close": [0]})
        ret, green, dow = aggregate_return_metrics(daily_df, 0, 0, 0, None, "AAPL", "2025-01-02")
        assert ret is None
        assert green is None
        assert dow is None


# =============================================================================
# Range metrics
# =============================================================================


class TestRangeMetrics:
    def test_aggregate_range_basic(self):
        result = aggregate_range_metrics(
            open_price=100.0, high_price=110.0, low_price=95.0, symbol="AAPL", date_to_test="2025-01-02"
        )
        assert result["o2h"] == 10.0
        assert result["l2h"] == pytest.approx(15.789, rel=1e-2)
        assert result["o2l"] == -5.0
        # 10% exactly -> rng_10_p bucket (bucket_range uses <10 for 5-10)
        assert result["o2h_bucket"] == "rng_10_p"
        assert result["l2h_bucket"] == "rng_10_p"
        assert result["o2l_bucket"] == "rng_0_2"
        assert result["o2h_example"] is not None
        assert result["l2h_example"] is not None
        assert result["o2l_example"] is not None

    def test_aggregate_range_invalid(self):
        result = aggregate_range_metrics(
            open_price=0, high_price=110, low_price=95, symbol="AAPL", date_to_test="2025-01-02"
        )
        assert result == {}


# =============================================================================
# Volume metrics
# =============================================================================


class TestVolumeMetrics:
    def test_aggregate_volume_basic(self):
        daily_df = pd.DataFrame(
            {
                "volume": [
                    1_000_000,
                    1_100_000,
                    1_200_000,
                    1_300_000,
                    1_400_000,
                    1_500_000,
                    1_600_000,
                    1_700_000,
                    1_800_000,
                    1_900_000,
                ],
            }
        )
        result = aggregate_volume_metrics(daily_df, 9, current_vol=2_000_000, symbol="AAPL")
        # avg_vol10 at index 9 uses the previous 9 values (shift(1) rolling)
        expected_avg_vol = (
            1_000_000 + 1_100_000 + 1_200_000 + 1_300_000 + 1_400_000 + 1_500_000 + 1_600_000 + 1_700_000 + 1_800_000
        ) / 9
        expected_rel_vol = 2_000_000 / expected_avg_vol
        assert result["rel_vol"] == pytest.approx(expected_rel_vol, rel=1e-6)
        # 1.43 < 2, so bucket is rvol_lt_2
        assert result["bucket"] == "rvol_lt_2"

    def test_aggregate_volume_insufficient_data(self):
        daily_df = pd.DataFrame({"volume": [1_000_000]})
        result = aggregate_volume_metrics(daily_df, 0, current_vol=1_000_000, symbol="AAPL")
        # avg_vol10 returns None for day_idx=0 on single-row df (shift(1) returns NaN)
        assert result["rel_vol"] is None or result["bucket"] is None  # depends on avg_vol10 implementation


# =============================================================================
# Gap metrics
# =============================================================================


class TestGapMetrics:
    def test_aggregate_gap_basic(self):
        result = aggregate_gap_metrics(
            open_price=105.0, prev_close_val=100.0, close_price=110.0, symbol="AAPL", date_to_test="2025-01-02"
        )
        assert result["gap_pct"] == 5.0
        # 5% exactly falls into gap_5_10 bucket (bucket_gap uses <5 for 0_5)
        assert result["bucket"] == "gap_5_10"
        assert result["trade_return"] == pytest.approx(4.7619, rel=1e-2)
        assert result["example"] is not None

    def test_aggregate_gap_no_prev_close(self):
        result = aggregate_gap_metrics(
            open_price=105.0, prev_close_val=None, close_price=110.0, symbol="AAPL", date_to_test="2025-01-02"
        )
        assert result == {}


# =============================================================================
# Price bucket
# =============================================================================


class TestPriceBucket:
    def test_aggregate_price_bucket_low(self):
        result = aggregate_price_bucket(open_price=3.0, close_price=3.5, symbol="AAPL", date_to_test="2025-01-02")
        assert result["bucket"] == "prc_lt_5"
        assert result["example"]["price"] == 3.0
        assert result["example"]["return"] == pytest.approx(16.666, rel=1e-1)

    def test_aggregate_price_bucket_high(self):
        result = aggregate_price_bucket(open_price=150.0, close_price=155.0, symbol="AAPL", date_to_test="2025-01-02")
        assert result["bucket"] == "prc_gt_100"
        assert result["example"]["price"] == 150.0

    def test_aggregate_price_bucket_none(self):
        result = aggregate_price_bucket(open_price=None, close_price=100.0, symbol="AAPL", date_to_test="2025-01-02")
        assert result is None


# =============================================================================
# Premarket metrics
# =============================================================================


class TestPremarketMetrics:
    def test_aggregate_premarket_basic(self, monkeypatch):
        # Create minute data covering premarket (04:00-09:29) and regular session (09:30-10:00)
        index = pd.date_range("2025-01-02 04:00", periods=360, freq="1min", tz="America/New_York")
        # Build high/low series: increasing values so we can identify extremes
        highs = list(range(100, 460))  # 360 values
        lows = list(range(90, 450))
        minute_data = pd.DataFrame(
            {
                "high": highs,
                "low": lows,
                "close": list(range(95, 455)),
            },
            index=index,
        )

        result = aggregate_premarket_metrics(
            minute_data, prev_close_val=100.0, symbol="AAPL", date_to_test="2025-01-02"
        )

        assert "pm_high" in result
        assert "pm_low" in result
        assert "day_high" in result
        assert "day_low" in result
        # All should be numeric (not None) because data covers full day
        assert result["pm_high"] is not None
        assert result["pm_low"] is not None
        assert result["day_high"] is not None
        assert result["day_low"] is not None
        # Premarket and day metrics present
        assert "pm_high_pct" in result
        assert "pm_low_pct" in result
        assert "pm_high_bucket" in result
        assert "pm_low_bucket" in result


# =============================================================================
# Breakout times
# =============================================================================


class TestBreakoutTimes:
    def test_aggregate_breakout_times(self):
        index = pd.date_range("2025-01-02 09:25", periods=30, freq="1min", tz="America/New_York")
        minute_data = pd.DataFrame(
            {
                "close": [100.0] * 10 + [101.0] * 10 + [102.0] * 10,
                "low": [99.0] * 10 + [100.0] * 10 + [101.0] * 10,
            },
            index=index,
        )
        # pm_low set very low to avoid triggering low breakout in test
        breakout_high, breakout_low = aggregate_breakout_times(
            minute_data, pm_high=100.5, pm_low=0, symbol="AAPL", date_to_test="2025-01-02"
        )
        assert len(breakout_high) == 1
        assert len(breakout_low) == 0
        assert breakout_high[0]["time"] in ["09:35", "09:36"]  # around when price first > 100.5


# =============================================================================
# HOD/LOD times
# =============================================================================


class TestHODLOD:
    def test_aggregate_hod_lod(self):
        index = pd.date_range("2025-01-02 09:30", periods=10, freq="1min", tz="America/New_York")
        minute_data = pd.DataFrame(
            {
                "high": [100, 101, 102, 103, 104, 103, 102, 101, 100, 99],
                "low": [99, 98, 97, 96, 95, 96, 97, 98, 99, 100],
            },
            index=index,
        )
        hod, lod = aggregate_hod_lod_times(minute_data, symbol="AAPL", date_to_test="2025-01-02")
        assert len(hod) == 1
        assert len(lod) == 1
        assert hod[0]["time"] == "09:34"  # max high at row index 4
        assert lod[0]["time"] == "09:34"  # min low at row index 4


# =============================================================================
# Volume by time
# =============================================================================


class TestVolumeByTime:
    def test_aggregate_volume_by_time(self):
        index = pd.date_range("2025-01-02 04:00", periods=120, freq="1min", tz="America/New_York")
        minute_data = pd.DataFrame({"volume": list(range(120))}, index=index)
        result = aggregate_volume_by_time(minute_data, symbol="AAPL")

        assert isinstance(result, dict)
        # Should have buckets like "04:00", "04:30", etc.
        assert "04:00" in result or "04:30" in result


# =============================================================================
# Aggregation helpers
# =============================================================================


class TestBucketTimeAggregation:
    def test_bucket_time_aggregation(self):
        times = [{"time": "09:31"}, {"time": "09:32"}, {"time": "09:45"}]
        buckets, examples = bucket_time(times)
        assert buckets["09:30"] == 3

    def test_mean_time_aggregation(self):
        times = [{"time": "09:00"}, {"time": "10:00"}, {"time": "11:00"}]
        assert calculate_mean_time(times) == "10:00"

    def test_median_time_aggregation(self):
        times = [{"time": "09:00"}, {"time": "10:00"}, {"time": "11:00"}, {"time": "12:00"}]
        assert calculate_median_time(times) == "10:30"


class TestRangeBucketAggregation:
    def test_aggregate_range_buckets(self):
        input_buckets = {
            "rng_0_2": [
                {"range_pct": 1.0, "symbol": "A", "date": "2025-01-02"},
                {"range_pct": 1.5, "symbol": "B", "date": "2025-01-02"},
            ],
            "rng_2_5": [{"range_pct": 3.0, "symbol": "C", "date": "2025-01-02"}],
        }
        result = aggregate_range_buckets(input_buckets)

        assert "rng_0_2" in result
        assert result["rng_0_2"]["count"] == 2
        assert result["rng_0_2"]["avg_range"] == pytest.approx(1.25, rel=1e-2)
        assert "rng_2_5" in result
        assert result["rng_2_5"]["avg_range"] == 3.0

    def test_aggregate_range_buckets_empty(self):
        result = aggregate_range_buckets({})
        assert result == {}


class TestGenericBucketStats:
    def test_aggregate_bucket_stats(self):
        buckets = {
            "bucket_a": [{"rel_vol": 1.5}, {"rel_vol": 2.5}],
            "bucket_b": [{"rel_vol": 5.0}],
        }
        result = aggregate_bucket_stats(buckets, value_key="rel_vol", avg_key="avg_rel_vol")

        assert result["bucket_a"]["count"] == 2
        assert result["bucket_a"]["avg_rel_vol"] == 2.0
        assert result["bucket_b"]["avg_rel_vol"] == 5.0


class TestVolumeDistribution:
    def test_aggregate_volume_distribution(self):
        vbt = {"09:30": [1000, 2000], "10:00": [3000]}
        result = aggregate_volume_distribution(vbt)

        assert result["09:30"]["total"] == 3000
        assert result["09:30"]["avg"] == 1500
        assert result["09:30"]["count"] == 2
        assert result["10:00"]["total"] == 3000


class TestPriceBuckets:
    def test_aggregate_price_buckets(self):
        pb = {
            "prc_lt_5": [{"price": 3.0, "return": 10.0}, {"price": 4.0, "return": 5.0}],
            "prc_gt_100": [{"price": 150.0, "return": 2.0}],
        }
        result = aggregate_price_buckets(pb)

        assert result["prc_lt_5"]["count"] == 2
        assert result["prc_lt_5"]["avg_price"] == 3.5
        assert result["prc_lt_5"]["avg_return"] == 7.5
        assert result["prc_gt_100"]["avg_price"] == 150.0


class TestRelVolBuckets:
    def test_aggregate_rel_vol_buckets(self):
        rvb = {
            "rvol_lt_2": [{"rel_vol": 1.2}, {"rel_vol": 1.5}],
            "rvol_gt_10": [{"rel_vol": 15.0}],
        }
        result = aggregate_rel_vol_buckets(rvb)

        assert result["rvol_lt_2"]["count"] == 2
        assert result["rvol_lt_2"]["avg_rel_vol"] == 1.35
        assert result["rvol_gt_10"]["avg_rel_vol"] == 15.0


class TestPremarketBuckets:
    def test_aggregate_premarket_buckets(self):
        buckets = {
            "pm_pct_0_50": [{"pm_high_pct": 40.0}, {"pm_high_pct": 45.0}],
            "pm_pct_90_100": [{"pm_high_pct": 95.0}],
        }
        result = aggregate_premarket_buckets(buckets, value_key="pm_high_pct")

        assert result["pm_pct_0_50"]["count"] == 2
        assert result["pm_pct_0_50"]["avg_pct"] == 42.5
        assert result["pm_pct_0_50"]["min_pct"] == 40.0
        assert result["pm_pct_0_50"]["max_pct"] == 45.0


class TestDayOfWeek:
    def test_aggregate_day_of_week(self):
        dow_data = {
            "mon": [
                {"return": 1.0, "green": 1, "symbol": "A", "date": "2025-01-06"},
                {"return": -0.5, "green": 0, "symbol": "B", "date": "2025-01-06"},
            ],
            "fri": [{"return": 2.0, "green": 1, "symbol": "C", "date": "2025-01-10"}],
        }
        result = aggregate_day_of_week(dow_data)

        assert "mon" in result
        assert result["mon"]["count"] == 2
        assert result["mon"]["avg_return"] == 0.25
        assert result["mon"]["close_green_pct"] == 50.0
        assert "fri" in result
        assert result["fri"]["avg_return"] == 2.0


class TestGapAnalysis:
    def test_aggregate_gap_analysis(self):
        gap_data = {
            "gap_0_5": {"count": 10, "returns": [1.0, 2.0, 0.5], "examples": []},
            "gap_5_10": {"count": 5, "returns": [3.0, 4.0], "examples": []},
        }
        result = aggregate_gap_analysis(gap_data)

        assert result["gap_0_5"]["count"] == 10
        assert result["gap_0_5"]["avg_return"] == pytest.approx(1.1667, rel=1e-2)
        assert result["gap_5_10"]["avg_return"] == 3.5


class TestSectorPerformance:
    def test_aggregate_sector_performance(self):
        sector_data = {
            "Technology": {"count": 5, "returns": [1.0, 2.0, 0.5, -0.5, 1.5], "examples": []},
            "Healthcare": {"count": 3, "returns": [0.5, 0.3, 0.7], "examples": []},
        }
        result = aggregate_sector_performance(sector_data)

        assert result["Technology"]["count"] == 5
        assert result["Technology"]["avg_return"] == 0.9
        assert result["Healthcare"]["avg_return"] == pytest.approx(0.5, rel=1e-2)