"""
Unit tests for leakage measurement.

Tests the core compute_data_leakage function with different scenarios.
"""

import pytest
import pandas as pd
from mleakdetect.core.measure import compute_data_leakage


def test_random_split_high_leakage():
    """
    Random split should have high leakage.

    When we randomly shuffle and split, train and test overlap in time,
    so leakage should be close to 1.
    """
    # Create simple time series data
    dates = pd.date_range('2026-01-01', periods=100, freq='D')
    df = pd.DataFrame({
        'date': dates,
        'value': range(100)
    })

    # Random split (shuffle then split)
    shuffled = df.sample(frac=1, random_state=42).reset_index(drop=True)
    train = shuffled.iloc[:70]
    test = shuffled.iloc[70:]

    # Measure leakage
    leakage = compute_data_leakage(train, test, 'date', alpha=1.0)

    # Random split should have high leakage
    assert leakage > 0.5, f"Random split should have high leakage, got {leakage}"
    print(f"Random split leakage: {leakage:.4f}")

def test_temporal_split_low_leakage():
    """
    Temporal split should have low leakage.

    When train comes before test chronologically,
    leakage should be close to 0.
    """
    dates = pd.date_range('2026-01-01', periods=100, freq='D')
    df = pd.DataFrame({
        'date': dates,
        'value': range(100)
    })

    # Temporal split (first 70 for train, last 30 for test)
    train = df.iloc[:70].copy()
    test = df.iloc[70:].copy()

    # Measure leakage
    leakage = compute_data_leakage(train, test, 'date', alpha=1.0)

    # Temporal split should have low leakage
    assert leakage < 0.3, f"Temporal split should have low leakage, got {leakage}"
    print(f"Temporal split leakage: {leakage:.4f}")

def test_same_data_max_leakage():
    """
    Using same data for train and test should give leakage = 1.

    This is the worst case scenario.
    """
    dates = pd.date_range('2026-01-01', periods=50, freq='D')
    df = pd.DataFrame({
        'date': dates,
        'value': range(50)
    })

    # Use same data for both train and test
    train = df.copy()
    test = df.copy()

    leakage = compute_data_leakage(train, test, 'date', alpha=1.0)

    # Should be exactly 1
    assert abs(leakage - 1.0) < 0.01, f"Same data should give leakage=1, got {leakage}"
    print(f"Same data leakage: {leakage:.4f}")

def test_alpha_parameter():
    """
    Test that alpha parameter changes leakage value.

    alpha controls how much we penalize backward temporal leakage.
    """
    dates = pd.date_range('2026-01-01', periods=100, freq='D')
    df = pd.DataFrame({
        'date': dates,
        'value': range(100)
    })

    # Random split
    shuffled = df.sample(frac=1, random_state=42).reset_index(drop=True)
    train = shuffled.iloc[:70]
    test = shuffled.iloc[70:]

    # Test with different alpha values
    leakage_alpha_0 = compute_data_leakage(train, test, 'date', alpha=0.0)
    leakage_alpha_1 = compute_data_leakage(train, test, 'date', alpha=1.0)

    # Different alpha should give different results
    assert leakage_alpha_0 != leakage_alpha_1
    print(f"Alpha=0: {leakage_alpha_0:.4f}, Alpha=1: {leakage_alpha_1:.4f}")

def test_empty_test_raises_error():
    """
    Empty test set should raise an error.

    We can't normalize if there's no test data.
    """
    dates = pd.date_range('2026-01-01', periods=50, freq='D')
    train = pd.DataFrame({
        'date': dates,
        'value': range(50)
    })

    test = pd.DataFrame({
        'date': pd.Series([], dtype='datetime64[ns]'),
        'value': pd.Series([], dtype='int64')
    })

    with pytest.raises(ValueError):
        compute_data_leakage(train, test, 'date', alpha=1.0)


def test_invalid_time_column():
    """
    Missing time column should raise error.
    """
    df = pd.DataFrame({
        'date': pd.date_range('2026-01-01', periods=50),
        'value': range(50)
    })

    train = df.iloc[:30].copy()
    test = df.iloc[30:].copy()

    with pytest.raises(ValueError):
        compute_data_leakage(train, test, 'wrong_column', alpha=1.0)


if __name__ == "__main__":
    # Run tests manually
    print("\n=== Running Leakage Measurement Tests ===\n")

    test_random_split_high_leakage()
    test_temporal_split_low_leakage()
    test_same_data_max_leakage()
    test_alpha_parameter()
    test_empty_test_raises_error()
    test_invalid_time_column()