"""
Tests for time encoding, specifically DST transitions.
"""
import pytest
import pandas as pd
import numpy as np
import pytz # For timezone handling

@pytest.fixture(scope="module")
def generate_dst_timeseries():
    """
    Generate a minute-frequency timestamp series crossing DST transitions
    for a specific timezone (e.g., US/Eastern).
    """
    # Example: US/Eastern DST Start (e.g., March 10, 2024 2:00 AM -> 3:00 AM)
    # Example: US/Eastern DST End (e.g., Nov 3, 2024 2:00 AM -> 1:00 AM)
    tz = pytz.timezone('US/Eastern')

    # Create timestamps around DST start
    dst_start_range = pd.date_range(
        start='2024-03-10 01:00:00', end='2024-03-10 04:00:00', freq='T', tz=tz
    )
    # Create timestamps around DST end
    dst_end_range = pd.date_range(
        start='2024-11-03 00:00:00', end='2024-11-03 03:00:00', freq='T', tz=tz
    )

    # Combine and ensure uniqueness/order (though disjoint here)
    timestamps = dst_start_range.union(dst_end_range)
    df = pd.DataFrame(index=timestamps)
    df.index.name = 'timestamp'
    return df

def calculate_cyclical_features(df):
    """Helper to calculate sin/cos features from a datetime index."""
    if not isinstance(df.index, pd.DatetimeIndex):
        raise TypeError("Input DataFrame must have a DatetimeIndex.")

    # Ensure timezone is present (fixture provides it)
    if df.index.tz is None:
         print("Warning: Index timezone is None, assuming UTC for calculation.")
         timestamp_source = df.index.tz_localize('utc')
    else:
         timestamp_source = df.index

    # Use UTC hour for consistent calculation if timezone handling upstream is complex
    # Or use localized hour if pipeline guarantees consistent local TZ
    # Here, let's use the localized hour provided by the fixture
    hour_of_day = timestamp_source.hour
    # minute_of_day = timestamp_source.hour * 60 + timestamp_source.minute # Alternative

    df['hour_sin'] = np.sin(2 * np.pi * hour_of_day / 24)
    df['hour_cos'] = np.cos(2 * np.pi * hour_of_day / 24)
    return df


def test_cyclical_features_continuity(generate_dst_timeseries):
    """
    Check if hour_sin and hour_cos features are continuous (no large jumps)
    across DST transitions, assuming calculation uses localized time.
    If using UTC hour, continuity is guaranteed, but might not capture
    local market patterns intended.
    """
    df = generate_dst_timeseries
    df = calculate_cyclical_features(df)

    # Check differences between consecutive values
    sin_diff = df['hour_sin'].diff().abs()
    cos_diff = df['hour_cos'].diff().abs()

    # Define a reasonable threshold for a jump (e.g., difference > value for 15 mins)
    # Max change in sin(2*pi*h/24) over 1 minute is small.
    # A jump of 1 hour means h changes by 1, argument changes by pi/12.
    # Max diff sin(x+pi/12) - sin(x) is approx pi/12 ~ 0.26
    max_allowed_diff = 0.3 # Allow slightly more than 1 hour jump equivalent

    print(f"\nMax Sin Diff: {sin_diff.max():.4f}")
    print(f"Max Cos Diff: {cos_diff.max():.4f}")

    assert sin_diff.max() < max_allowed_diff, \
        f"Large jump detected in hour_sin ({sin_diff.max():.4f}) around DST. Check time source/calculation."
    assert cos_diff.max() < max_allowed_diff, \
        f"Large jump detected in hour_cos ({cos_diff.max():.4f}) around DST. Check time source/calculation."

    # Optional: Plot to visually inspect
    # import matplotlib.pyplot as plt
    # plt.figure()
    # plt.plot(df.index, df['hour_sin'], '.-.', label='sin')
    # plt.plot(df.index, df['hour_cos'], '.-.', label='cos')
    # plt.title('Cyclical Features Across DST')
    # plt.legend()
    # plt.xticks(rotation=45)
    # plt.tight_layout()
    # plt.show() 