"""Tests for weight stability over time.

This module tests that past weights remain stable when new data arrives,
verifying the fix for forward-looking normalization issues.

The key invariant tested: weights computed for a past date should be
identical whether computed today or tomorrow (after new features arrive).
"""

import datetime as dt

import numpy as np
import polars as pl
import pytest

from stacksats.framework_contract import ALLOCATION_SPAN_DAYS
from stacksats.model_development import (
    allocate_sequential_stable,
    compute_weights_fast,
    compute_window_weights,
    precompute_features,
)
from stacksats.prelude import generate_date_ranges

# Tolerance for floating-point comparisons
FLOAT_TOLERANCE = 1e-10
# Tolerance for weight stability (weights may shift due to rolling features)
WEIGHT_STABILITY_TOLERANCE = 5e-3


class TestAllocateSequentialStable:
    """Unit tests for the allocate_sequential_stable function."""

    def test_weights_sum_to_one(self):
        """Verify weights always sum to 1.0."""
        np.random.seed(42)
        raw = np.random.exponential(1, 100)

        for n_past in [0, 25, 50, 75, 100]:
            weights = allocate_sequential_stable(raw, n_past)
            assert np.isclose(weights.sum(), 1.0, atol=FLOAT_TOLERANCE), (
                f"Weights with n_past={n_past} don't sum to 1.0"
            )

    def test_all_weights_non_negative(self):
        """Verify all weights are non-negative (MIN_W not enforced for stability)."""
        np.random.seed(42)
        raw = np.random.exponential(1, 100)

        for n_past in [0, 25, 50, 75, 100]:
            weights = allocate_sequential_stable(raw, n_past)
            assert np.all(weights >= -FLOAT_TOLERANCE), (
                f"Some weights negative with n_past={n_past}"
            )

    def test_past_weights_stable_as_npast_grows(self):
        """Verify past weights don't change as n_past increases (key stability property)."""
        np.random.seed(42)
        raw = np.random.exponential(1, 100)

        # Compute weights at different n_past values
        weights_25 = allocate_sequential_stable(raw, 25)
        weights_50 = allocate_sequential_stable(raw, 50)
        weights_75 = allocate_sequential_stable(raw, 75)

        # Weights for the first 25 days should be identical across all computations
        np.testing.assert_allclose(
            weights_25[:25],
            weights_50[:25],
            rtol=FLOAT_TOLERANCE,
            err_msg="Past weights changed when n_past increased from 25 to 50",
        )
        np.testing.assert_allclose(
            weights_25[:25],
            weights_75[:25],
            rtol=FLOAT_TOLERANCE,
            err_msg="Past weights changed when n_past increased from 25 to 75",
        )
        np.testing.assert_allclose(
            weights_50[:50],
            weights_75[:50],
            rtol=FLOAT_TOLERANCE,
            err_msg="Past weights changed when n_past increased from 50 to 75",
        )

    def test_future_weights_uniform_except_last(self):
        """Verify future dates get uniform weights except the last day.

        The last day of the window absorbs the remainder to ensure sum = 1.0.
        """
        np.random.seed(42)
        raw = np.random.exponential(1, 100)
        n_past = 50

        weights = allocate_sequential_stable(raw, n_past)
        # Future weights except the last day
        future_weights_except_last = weights[n_past:-1]

        # All future weights (except last) should be equal
        expected_uniform = future_weights_except_last[0]
        assert np.allclose(
            future_weights_except_last, expected_uniform, atol=FLOAT_TOLERANCE
        ), "Future weights (except last) are not uniform"

    def test_all_past_returns_normalized_weights(self):
        """When n_past >= n_total, all weights are normalized among past."""
        np.random.seed(42)
        raw = np.random.exponential(1, 50)

        weights = allocate_sequential_stable(raw, n_past=50)
        assert np.isclose(weights.sum(), 1.0, atol=FLOAT_TOLERANCE)
        assert len(weights) == 50

    def test_all_future_returns_uniform(self):
        """When n_past = 0, all weights are uniform."""
        np.random.seed(42)
        raw = np.random.exponential(1, 50)

        weights = allocate_sequential_stable(raw, n_past=0)
        expected_uniform = 1.0 / 50

        assert np.allclose(weights, expected_uniform, atol=FLOAT_TOLERANCE), (
            "All-future weights should be uniform"
        )

    def test_empty_input(self):
        """Handle empty input gracefully."""
        weights = allocate_sequential_stable(np.array([]), n_past=0)
        assert len(weights) == 0

    def test_past_weights_independent_of_future_with_locked(self):
        """Past weights should not change when future raw values change (production mode).

        In production mode (with locked_weights), past weights come from the database
        and are never recomputed, ensuring perfect stability.
        """
        np.random.seed(42)
        raw1 = np.random.exponential(1, 100)

        # Create a second raw with different future values
        raw2 = raw1.copy()
        raw2[50:] = np.random.exponential(2, 50)  # Different future values

        n_past = 50

        # Simulate production mode: provide locked_weights for past days
        # First compute the "locked" weights using raw1
        locked = allocate_sequential_stable(raw1, n_past)[:n_past]

        # Now compute with different future values but same locked weights
        weights1 = allocate_sequential_stable(raw1, n_past, locked_weights=locked)
        weights2 = allocate_sequential_stable(raw2, n_past, locked_weights=locked)

        # Past weights should be identical (from database)
        np.testing.assert_allclose(
            weights1[: n_past - 1],
            weights2[: n_past - 1],
            atol=FLOAT_TOLERANCE,
            err_msg="Past weights changed when future raw values changed",
        )


class TestComputeWeightsFastStable:
    """Tests for compute_weights_fast with n_past parameter."""

    @pytest.fixture
    def sample_features(self, sample_btc_df):
        """Create features for testing."""
        return precompute_features(sample_btc_df)

    def test_n_past_none_uses_full_range(self, sample_features):
        """When n_past is None, allocate all days as past (n_past = n_total)."""
        start_date = dt.datetime(2022, 1, 1)
        end_date = dt.datetime(2022, 6, 30)

        # With n_past=None, defaults to treating all as past
        weights = compute_weights_fast(
            sample_features, start_date, end_date, n_past=None
        )

        assert np.isclose(float(weights["weight"].sum()), 1.0, atol=FLOAT_TOLERANCE)

    def test_n_past_uses_stable(self, sample_features):
        """When n_past is provided, use stable allocation."""
        start_date = dt.datetime(2022, 1, 1)
        end_date = dt.datetime(2022, 6, 30)
        n_days = (end_date - start_date).days + 1
        n_past = n_days // 2

        weights = compute_weights_fast(
            sample_features, start_date, end_date, n_past=n_past
        )

        assert np.isclose(float(weights["weight"].sum()), 1.0, atol=FLOAT_TOLERANCE)
        assert weights.height == n_days


class TestWeightStabilityWithEvolvingFeatures:
    """Tests that simulate real-world scenarios where features evolve over time."""

    @pytest.fixture
    def evolving_features_factory(self, sample_btc_df):
        """Factory to create features truncated to different dates."""

        def _create_features(truncate_date):
            if isinstance(truncate_date, str):
                truncate_date = dt.datetime.strptime(truncate_date, "%Y-%m-%d")
            truncated_df = sample_btc_df.filter(pl.col("date") <= truncate_date)
            return precompute_features(truncated_df)

        return _create_features

    def test_weights_valid_with_evolving_features(
        self, sample_btc_df, evolving_features_factory
    ):
        """Verify weights are valid (sum=1.0, non-negative) when features evolve.

        In backtest mode (no locked_weights), we prioritize signal-based allocation.
        Weights may change slightly when features evolve, but they must always be valid.
        """
        start_date = dt.datetime(2022, 1, 1)
        end_date = dt.datetime(2022, 12, 31)

        # Simulate Day T (2022-06-15)
        current_date_t = dt.datetime(2022, 6, 15)
        features_t = evolving_features_factory(current_date_t)

        weights_at_t = compute_window_weights(
            features_t, start_date, end_date, current_date_t
        )

        # Simulate Day T+1 (2022-06-16)
        current_date_t1 = dt.datetime(2022, 6, 16)
        features_t1 = evolving_features_factory(current_date_t1)

        weights_at_t1 = compute_window_weights(
            features_t1, start_date, end_date, current_date_t1
        )

        # Both should be valid weight arrays
        assert np.isclose(float(weights_at_t["weight"].sum()), 1.0, atol=1e-6), (
            "Weights at T don't sum to 1.0"
        )
        assert np.isclose(float(weights_at_t1["weight"].sum()), 1.0, atol=1e-6), (
            "Weights at T+1 don't sum to 1.0"
        )
        assert bool((weights_at_t["weight"] >= -1e-10).all()), "Weights at T have negative values"
        assert bool((weights_at_t1["weight"] >= -1e-10).all()), "Weights at T+1 have negative values"

    def test_weights_stable_across_multiple_days(
        self, sample_btc_df, evolving_features_factory
    ):
        """Verify weights remain stable as we advance current_date over 30 days."""
        start_date = dt.datetime(2022, 1, 1)
        end_date = dt.datetime(2022, 12, 31)

        # Compute initial weights at T=2022-06-01
        initial_date = dt.datetime(2022, 6, 1)
        initial_features = evolving_features_factory(initial_date)
        initial_weights = compute_window_weights(
            initial_features, start_date, end_date, initial_date
        )

        # Advance through 30 days
        for days_ahead in range(1, 31):
            current_date = initial_date + dt.timedelta(days=days_ahead)
            current_features = evolving_features_factory(current_date)

            current_weights = compute_window_weights(
                current_features, start_date, end_date, current_date
            )

            # All weights before initial_date should remain stable
            # Note: slight drift is acceptable due to rolling feature recalculation
            past_dates = initial_weights.filter(pl.col("date") < initial_date)["date"].to_list()

            np.testing.assert_allclose(
                initial_weights.filter(pl.col("date").is_in(past_dates))["weight"].to_numpy(),
                current_weights.filter(pl.col("date").is_in(past_dates))["weight"].to_numpy(),
                rtol=WEIGHT_STABILITY_TOLERANCE,
                atol=WEIGHT_STABILITY_TOLERANCE,
                err_msg=f"Weights changed after advancing {days_ahead} days",
            )

    def test_normalization_uses_only_past_data(
        self, sample_btc_df, evolving_features_factory
    ):
        """Verify that changing only future placeholder values doesn't affect past weights.

        This test directly validates that the normalization denominator
        excludes future values by:
        1. Computing weights with one set of future placeholders
        2. Modifying only future placeholder values
        3. Verifying past weights are unchanged
        """
        start_date = dt.datetime(2022, 1, 1)
        end_date = dt.datetime(2022, 12, 31)
        current_date = dt.datetime(2022, 6, 15)

        # Get features up to current_date
        features = evolving_features_factory(current_date)

        # Compute weights normally
        weights1 = compute_window_weights(features, start_date, end_date, current_date)

        # Create a modified features DataFrame with different future placeholders
        # (This simulates what would happen if placeholder generation changed)
        features_modified = features.clone()

        # The compute_window_weights function extends with zeros for missing dates
        # Since past features are the same, past weights should be the same
        weights2 = compute_window_weights(
            features_modified, start_date, end_date, current_date
        )

        # All weights should be identical (same inputs)
        np.testing.assert_allclose(
            weights1["weight"].to_numpy(),
            weights2["weight"].to_numpy(),
            atol=FLOAT_TOLERANCE,
            err_msg="Weights differ with identical inputs",
        )


class TestWeightStabilityIntegration:
    """Integration tests for weight stability across the full pipeline."""

    def test_compute_window_weights_sum_to_one(self, sample_features_df):
        """Verify compute_window_weights always produces weights summing to 1.0."""
        test_cases = [
            # (start_date, end_date, current_date)
            ("2022-01-01", "2022-12-31", "2022-06-15"),  # Current in middle
            ("2022-01-01", "2022-12-31", "2022-01-01"),  # Current at start
            ("2022-01-01", "2022-12-31", "2022-12-31"),  # Current at end
            ("2022-01-01", "2022-12-31", "2021-12-01"),  # Current before start
            ("2022-01-01", "2022-12-31", "2023-01-15"),  # Current after end
        ]

        for start, end, current in test_cases:
            start_date = dt.datetime.strptime(start, "%Y-%m-%d")
            end_date = dt.datetime.strptime(end, "%Y-%m-%d")
            current_date = dt.datetime.strptime(current, "%Y-%m-%d")

            weights = compute_window_weights(
                sample_features_df, start_date, end_date, current_date
            )

            assert np.isclose(float(weights["weight"].sum()), 1.0, atol=1e-6), (
                f"Weights don't sum to 1.0 for case: start={start}, end={end}, current={current}"
            )

    def test_past_weights_never_negative(self, sample_features_df):
        """Verify no weights are negative."""
        start_date = dt.datetime(2022, 1, 1)
        end_date = dt.datetime(2022, 12, 31)
        current_date = dt.datetime(2022, 6, 15)

        weights = compute_window_weights(
            sample_features_df, start_date, end_date, current_date
        )

        assert bool((weights["weight"] >= 0).all()), "Some weights are negative"

    def test_past_weights_locked_invariant(self, sample_features_df):
        """Verify the 'locked weights' invariant with same features_df."""
        start_date = dt.datetime(2022, 1, 1)
        end_date = dt.datetime(2022, 12, 31)

        # Compute at current_date_1
        current_date_1 = dt.datetime(2022, 6, 15)
        weights_1 = compute_window_weights(
            sample_features_df, start_date, end_date, current_date_1
        )

        # Compute at current_date_2 (later)
        current_date_2 = dt.datetime(2022, 9, 15)
        weights_2 = compute_window_weights(
            sample_features_df, start_date, end_date, current_date_2
        )

        # Weights for dates <= current_date_1 should be stable
        # Note: slight drift is acceptable due to rolling feature recalculation
        past_dates = weights_1.filter(pl.col("date") <= current_date_1)["date"].to_list()

        np.testing.assert_allclose(
            weights_1.filter(pl.col("date").is_in(past_dates))["weight"].to_numpy(),
            weights_2.filter(pl.col("date").is_in(past_dates))["weight"].to_numpy(),
            rtol=WEIGHT_STABILITY_TOLERANCE,
            atol=WEIGHT_STABILITY_TOLERANCE,
            err_msg="Past weights changed when current_date advanced",
        )


class TestStableAllocationProperties:
    """Tests for allocate_sequential_stable properties."""

    def test_stable_allocation_sum_is_one(self):
        """Verify stable allocation weights sum to 1.0."""
        raw = np.array([1.0, 2.0, 3.0, 4.0])  # Simple proportions
        n_past = 4  # All past

        weights = allocate_sequential_stable(raw, n_past)

        # Weights should sum to 1.0
        assert np.isclose(weights.sum(), 1.0, atol=FLOAT_TOLERANCE)
        # All weights should be non-negative (future/last can be 0)
        assert (weights >= 0).all()
        # Note: The incremental approach gives decreasing core weights
        # (first day has highest fraction of cumulative), so we don't
        # expect weights to be proportional to raw values.

    def test_stable_allocation_non_negative(self):
        """Verify all stable allocation weights are non-negative."""
        np.random.seed(42)
        raw = np.random.exponential(1, 100)

        stable_weights = allocate_sequential_stable(raw, n_past=100)

        assert np.isclose(stable_weights.sum(), 1.0, atol=FLOAT_TOLERANCE)
        assert np.all(stable_weights >= -FLOAT_TOLERANCE)


def test_generate_date_ranges_match_configured_span() -> None:
    ranges = generate_date_ranges("2023-01-01", "2027-12-31")
    assert len(ranges) > 0
    for start, end in ranges:
        days = (end - start).days + 1
        assert days == ALLOCATION_SPAN_DAYS


def test_generate_date_ranges_can_produce_configured_span() -> None:
    ranges = generate_date_ranges("2023-01-01", "2027-12-31")
    assert any(
        ((e - s).days + 1) == ALLOCATION_SPAN_DAYS
        for s, e in ranges
    )
