"""
Tests for label generation and potential leakage.

Ref: revisions.txt Step 1-A, 1.4
"""
import pytest
import pandas as pd
import numpy as np
import sys, os

# --- Add path for src imports --- #
# Assuming tests is one level down from the package root
script_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(script_dir) # Go up one level
src_path = os.path.join(project_root, 'src')
if src_path not in sys.path:
    sys.path.insert(0, src_path)
# --- End Add path --- #

# Import the function to test
from trading_pipeline import _generate_direction_labels

# --- Fixtures --- #
@pytest.fixture
def sample_close_data() -> pd.DataFrame:
    """Creates a sample DataFrame with close prices and DatetimeIndex."""
    # Generate data with some variation
    np.random.seed(42)
    prices = 100 + np.cumsum(np.random.randn(200) * 0.5)
    data = {'close': prices}
    index = pd.date_range(start='2023-01-01', periods=len(data['close']), freq='min', tz='UTC')
    df = pd.DataFrame(data, index=index)
    return df

@pytest.fixture
def sample_config() -> dict:
    """Provides a basic config dictionary."""
    return {
        'gru': {
            'prediction_horizon': 5,
            'use_ternary': False,
            'flat_sigma_multiplier': 0.25
        },
        'data': {
            'label_smoothing': 0.0
        }
    }

# --- Tests --- #

def test_lookahead_bias(sample_close_data, sample_config):
    """
    Test 1.4.a: Verify labels don't depend on information *beyond* the prediction horizon.
    Strategy: Modify future close prices (beyond horizon) and check if labels change.
    """
    df = sample_close_data
    config = sample_config
    horizon = config['gru']['prediction_horizon']

    # Generate baseline labels (binary)
    df_labeled_base, label_col_base = _generate_direction_labels(df.copy(), config)

    # Modify close prices far into the future (beyond the horizon needed for any label)
    df_modified = df.copy()
    future_index = len(df) - 1 # Index of the last point
    modify_point = future_index - horizon - 5 # Index well beyond the last needed future price
    if modify_point > 0:
        df_modified.iloc[modify_point:, df_modified.columns.get_loc('close')] *= 1.5 # Modify future prices

    # Generate labels with modified future data
    df_labeled_mod, label_col_mod = _generate_direction_labels(df_modified.copy(), config)

    # Align based on index (label function drops NaNs at the end)
    common_index = df_labeled_base.index.intersection(df_labeled_mod.index)
    labels_base_aligned = df_labeled_base.loc[common_index, label_col_base]
    labels_mod_aligned = df_labeled_mod.loc[common_index, label_col_mod]

    # Assert: Labels should be identical, as modification was beyond the horizon
    pd.testing.assert_series_equal(labels_base_aligned, labels_mod_aligned, check_names=False)

    # --- Repeat for Ternary --- #
    config['gru']['use_ternary'] = True
    df_labeled_base_t, label_col_base_t = _generate_direction_labels(df.copy(), config)
    df_labeled_mod_t, label_col_mod_t = _generate_direction_labels(df_modified.copy(), config)

    common_index_t = df_labeled_base_t.index.intersection(df_labeled_mod_t.index)
    labels_base_aligned_t = df_labeled_base_t.loc[common_index_t, label_col_base_t]
    labels_mod_aligned_t = df_labeled_mod_t.loc[common_index_t, label_col_mod_t]

    # Assert: Ternary labels should also be identical
    # Need careful comparison for list/array column
    assert labels_base_aligned_t.equals(labels_mod_aligned_t)

def test_binary_label_distribution(sample_close_data, sample_config):
    """
    Test 1.4.b: Check binary label distribution has >= 5% in each class.
    """
    df = sample_close_data
    config = sample_config
    config['gru']['use_ternary'] = False
    config['data']['label_smoothing'] = 0.0 # Ensure hard binary for this test

    df_labeled, label_col = _generate_direction_labels(df.copy(), config)

    assert not df_labeled.empty, "Label generation resulted in empty DataFrame"
    assert label_col in df_labeled.columns, f"Label column '{label_col}' not found"

    labels = df_labeled[label_col]
    counts = labels.value_counts(normalize=True)

    assert len(counts) == 2, f"Expected 2 binary classes, found {len(counts)}"
    assert counts.min() >= 0.05, f"Minimum binary class proportion ({counts.min():.2%}) is less than 5%"
    print(f"\nBinary Dist: {counts.to_dict()}") # Print for info

def test_soft_binary_label_distribution(sample_close_data, sample_config):
    """
    Test 1.4.b: Check soft binary label distribution has >= 5% in each effective class.
    """
    df = sample_close_data
    config = sample_config
    config['gru']['use_ternary'] = False
    config['data']['label_smoothing'] = 0.2 # Example smoothing
    smoothing = config['data']['label_smoothing']
    low_label = smoothing / 2.0
    high_label = 1.0 - smoothing / 2.0

    df_labeled, label_col = _generate_direction_labels(df.copy(), config)

    assert not df_labeled.empty, "Label generation resulted in empty DataFrame"
    assert label_col in df_labeled.columns, f"Label column '{label_col}' not found"

    labels = df_labeled[label_col]
    counts = labels.value_counts(normalize=True)

    assert len(counts) == 2, f"Expected 2 soft binary classes, found {len(counts)}"
    assert counts.min() >= 0.05, f"Minimum soft binary class proportion ({counts.min():.2%}) is less than 5%"
    assert low_label in counts.index, f"Low label {low_label} not found in counts"
    assert high_label in counts.index, f"High label {high_label} not found in counts"
    print(f"\nSoft Binary Dist: {counts.to_dict()}")

def test_ternary_label_distribution(sample_close_data, sample_config):
    """
    Test 1.4.b: Check ternary label distribution (flat=[0.15, 0.45], others >= 0.10).
    Uses default k=0.25.
    """
    df = sample_close_data
    config = sample_config
    config['gru']['use_ternary'] = True
    k = config['gru']['flat_sigma_multiplier'] # Should be 0.25 from fixture

    df_labeled, label_col = _generate_direction_labels(df.copy(), config)

    assert not df_labeled.empty, "Label generation resulted in empty DataFrame"
    assert label_col in df_labeled.columns, f"Label column '{label_col}' not found"

    # Decode one-hot labels back to ordinal for distribution check
    labels_one_hot = np.stack(df_labeled[label_col].values)
    assert labels_one_hot.shape[1] == 3, "Ternary labels should have 3 columns"
    ordinal_labels = np.argmax(labels_one_hot, axis=1)

    counts = np.bincount(ordinal_labels, minlength=3)
    total = len(ordinal_labels)
    dist_pct = counts / total * 100

    print(f"\nTernary Dist (k={k}): Down={dist_pct[0]:.1f}%, Flat={dist_pct[1]:.1f}%, Up={dist_pct[2]:.1f}%")

    # Check constraints based on design doc / implementation
    assert 15.0 <= dist_pct[1] <= 45.0, f"Flat class ({dist_pct[1]:.1f}%) out of expected range [15%, 45%] for k={k}"
    assert dist_pct[0] >= 10.0, f"Down class ({dist_pct[0]:.1f}%) is less than 10% (check impl threshold)"
    assert dist_pct[2] >= 10.0, f"Up class ({dist_pct[2]:.1f}%) is less than 10% (check impl threshold)"

# --- Old Tests (Keep or Remove?) ---
# The original tests checked 'future_close', which is related but not the final label.
# We can keep test_future_close_shift as it verifies the shift logic used internally.
# The NaN test is less relevant now as the main function handles NaN dropping.

def test_future_close_shift(sample_close_data):
    """Verify that 'future_close' is correctly shifted and has NaNs at the end."""
    df = sample_close_data
    horizon = 5 # Example horizon

    # Apply the logic directly for testing the shift itself
    df['future_close'] = df['close'].shift(-horizon)
    df['fwd_log_ret'] = np.log(df['future_close'] / df['close'])

    # Assertions
    # 1. Check for correct shift in fwd_log_ret
    # The first valid fwd_log_ret depends on close[0] and close[horizon]
    assert pd.notna(df['fwd_log_ret'].iloc[0])
    # The last valid fwd_log_ret depends on close[end-horizon-1] and close[end-1]
    assert pd.notna(df['fwd_log_ret'].iloc[len(df) - horizon - 1])

    # 2. Check for NaNs at the end due to shift
    assert pd.isna(df['fwd_log_ret'].iloc[-horizon:]).all()
    assert pd.notna(df['fwd_log_ret'].iloc[:-horizon]).all()

# def test_no_nan_in_future_close_output():
#     """Unit test to ensure no unexpected NaNs in the output of label creation (specific to the function)."""
#     # Setup similar to above, potentially call the actual DataLoader/label function
#     # Assert pd.notna(output_df['future_close'][:-horizon]).all()
#     pytest.skip("Test covered by NaN dropping in _generate_direction_labels and its tests.") 