# SPDX-FileCopyrightText: 2025 Contributors to the OpenSTEF project <openstef@lfenergy.org>
#
# SPDX-License-Identifier: MPL-2.0

"""Unit tests for the VersionedLagsAdder."""

from datetime import timedelta

import numpy as np
import pandas as pd
import pytest

from openstef_core.datasets.versioned_timeseries_dataset import VersionedTimeSeriesDataset
from openstef_core.exceptions import MissingColumnsError
from openstef_core.types import LeadTime
from openstef_models.transforms.time_domain.versioned_lags_adder import VersionedLagsAdder


def test_multiple_lag_features():
    """Test that multiple lags create separate features with correctly shifted timestamps.

    This test covers both single and multiple lag functionality:
    - Single lag: How -1h lag shifts timestamps forward by 1 hour
    - Multiple lags: How -1h and -2h lags create different shift patterns
    - Timestamp filtering: How lag features stay within original timestamp range
    """
    # Arrange
    simple_versioned_dataset = VersionedTimeSeriesDataset.from_dataframe(
        data=pd.DataFrame(
            {
                "available_at": pd.to_datetime([
                    "2025-01-01T10:00:00",
                    "2025-01-01T11:00:00",
                    "2025-01-01T12:00:00",
                    "2025-01-01T13:00:00",
                ]),
                "load": [100.0, 110.0, 120.0, 130.0],
            },
            index=pd.to_datetime([
                "2025-01-01T10:00:00",
                "2025-01-01T11:00:00",
                "2025-01-01T12:00:00",
                "2025-01-01T13:00:00",
            ]),
        ),
        sample_interval=timedelta(hours=1),
    )
    transform = VersionedLagsAdder(feature="load", lags=[timedelta(hours=-1), timedelta(hours=-2)])

    # Act
    result = transform.transform(simple_versioned_dataset)
    snapshot = result.select_version()

    # Assert
    lag_1h_col = "load_lag_-PT1H"
    lag_2h_col = "load_lag_-PT2H"

    # How multiple lags work with timestamp filtering:
    # Original data: 10:00->100, 11:00->110, 12:00->120, 13:00->130
    # 1h lag part: 11:00->100, 12:00->110, 13:00->120, 14:00->130 (shifted +1h)
    #              After filtering: 11:00->100, 12:00->110, 13:00->120 (14:00 filtered out)
    # 2h lag part: 12:00->100, 13:00->110, 14:00->120, 15:00->130 (shifted +2h)
    #              After filtering: 12:00->100, 13:00->110 (14:00-15:00 filtered out)
    # Combined index stays original: [10:00, 11:00, 12:00, 13:00]

    # Expected results (within original timestamp range):
    # 1h lag: [NaN, 100, 110, 120] - available at timestamps 11:00-13:00
    # 2h lag: [NaN, NaN, 100, 110] - available at timestamps 12:00-13:00
    expected_1h = pd.Series(
        [np.nan, 100.0, 110.0, 120.0],
        index=pd.DatetimeIndex(
            [
                "2025-01-01T10:00:00",
                "2025-01-01T11:00:00",
                "2025-01-01T12:00:00",
                "2025-01-01T13:00:00",
            ],
            name="timestamp",
        ),
        name=lag_1h_col,
    )

    expected_2h = pd.Series(
        [np.nan, np.nan, 100.0, 110.0],
        index=pd.DatetimeIndex(
            [
                "2025-01-01T10:00:00",
                "2025-01-01T11:00:00",
                "2025-01-01T12:00:00",
                "2025-01-01T13:00:00",
            ],
            name="timestamp",
        ),
        name=lag_2h_col,
    )

    actual_1h = snapshot.data[lag_1h_col]
    actual_2h = snapshot.data[lag_2h_col]

    pd.testing.assert_series_equal(actual_1h, expected_1h, check_freq=False)
    pd.testing.assert_series_equal(actual_2h, expected_2h, check_freq=False)


def test_lag_transform_with_horizon_filtering():
    """Test how lag values determine which data version quality is accessed based on lead time constraints.

    This test demonstrates the core insight of versioned lag transforms: longer lags require more advance notice,
    which forces the system to use higher quality (but later available) data versions.

    Scenario Setup:
    - Signal has 3 versions: "bad" (3d delay), "okay" (5d delay), "best" (10d delay)
    - Forecasting with 2-day lead time
    - Different lag values require different amounts of advance notice

    The Mathematics:
    For a forecast at time T with lag L and lead time H:
    - Data must be available by: T - L - H
    - Available data versions depend on this availability deadline

    Expected Behavior:
    - 5-day lag: needs data available by T-(-5d)-2d = T+3d → "bad" version (3d delay) just qualifies
    - 7-day lag: needs data available by T-(-7d)-2d = T+5d → "okay" version (5d delay) qualifies
    - 12-day lag: needs data available by T-(-12d)-2d = T+10d → "best" version (10d delay) qualifies

    Note: With timestamp filtering, only values within the original forecasting range are preserved.
    """
    # Arrange
    timestamps = pd.date_range("2025-01-01T10:00:00", periods=5, freq="D")

    multi_version_signal_dataset = VersionedTimeSeriesDataset.from_dataframe(
        data=pd.concat([
            pd.DataFrame(
                {
                    "available_at": timestamps + timedelta(days=3),
                    "signal": ["bad"] * len(timestamps),
                },
                index=timestamps,
            ),
            pd.DataFrame(
                {
                    "available_at": timestamps + timedelta(days=5),
                    "signal": ["okay"] * len(timestamps),
                },
                index=timestamps,
            ),
            pd.DataFrame(
                {
                    "available_at": timestamps + timedelta(days=10),
                    "signal": ["best"] * len(timestamps),
                },
                index=timestamps,
            ),
        ]),
        sample_interval=timedelta(days=1),
    )

    lead_time = LeadTime.from_string("P2D")  # 2-day forecast horizon
    transform = VersionedLagsAdder(feature="signal", lags=[timedelta(days=-5), timedelta(days=-7), timedelta(days=-12)])

    # Act
    result = transform.transform(multi_version_signal_dataset)

    # Assert
    snapshot = result.filter_by_lead_time(lead_time).select_version()

    # With timestamp filtering, lag features are constrained to original timestamp range
    # The test verifies that the correct version quality is accessed when data is available

    # Check if we have data for each lag and verify the version quality when available
    lag_5d_data = snapshot.data["signal_lag_-P5D"].dropna()
    lag_7d_data = snapshot.data["signal_lag_-P7D"].dropna()
    lag_12d_data = snapshot.data["signal_lag_-P12D"].dropna()

    # Test that when data is available, it's from the correct quality version
    if len(lag_5d_data) > 0:
        assert lag_5d_data.eq("bad").all(), "5-day lag should access 'bad' version when available"
    if len(lag_7d_data) > 0:
        assert lag_7d_data.eq("okay").all(), "7-day lag should access 'okay' version when available"
    if len(lag_12d_data) > 0:
        assert lag_12d_data.eq("best").all(), "12-day lag should access 'best' version when available"

    # Verify that lag features exist (even if empty due to filtering)
    assert "signal_lag_-P5D" in snapshot.feature_names, "5-day lag feature should exist"
    assert "signal_lag_-P7D" in snapshot.feature_names, "7-day lag feature should exist"
    assert "signal_lag_-P12D" in snapshot.feature_names, "12-day lag feature should exist"

    # Verify that the snapshot maintains the original timestamp range
    original_range = pd.date_range("2025-01-01T10:00:00", periods=5, freq="D")
    assert len(snapshot.data.index) == len(original_range), "Snapshot should maintain original timestamp range"


def test_missing_feature_error():
    """Test that lag transform raises MissingColumnsError when trying to create lag of non-existent feature.

    This ensures proper error handling when configuration specifies a feature that doesn't exist
    in the dataset, preventing silent failures in production pipelines.
    """
    # Arrange
    simple_versioned_dataset = VersionedTimeSeriesDataset.from_dataframe(
        data=pd.DataFrame(
            {
                "available_at": pd.to_datetime([
                    "2025-01-01T10:00:00",
                    "2025-01-01T11:00:00",
                    "2025-01-01T12:00:00",
                    "2025-01-01T13:00:00",
                ]),
                "load": [100.0, 110.0, 120.0, 130.0],
            },
            index=pd.to_datetime([
                "2025-01-01T10:00:00",
                "2025-01-01T11:00:00",
                "2025-01-01T12:00:00",
                "2025-01-01T13:00:00",
            ]),
        ),
        sample_interval=timedelta(hours=1),
    )
    transform = VersionedLagsAdder(feature="non_existent", lags=[timedelta(hours=-1)])

    # Act & Assert
    # Verify that attempting to create lag features for a non-existent feature fails appropriately
    with pytest.raises(MissingColumnsError) as exc_info:
        transform.transform(simple_versioned_dataset)
    assert "non_existent" in str(exc_info.value)
