"""
Unit tests for data splitting functions.

Tests temporal_split and group_split to ensure they work correctly.
"""

import pytest
import pandas as pd
from mleakdetect.core.split import temporal_split, group_split


def test_temporal_split_basic():
    """
    Test basic temporal split with 70/30 ratio.

    Should split data chronologically.
    """
    # Create 100 days of data
    dates = pd.date_range('2026-01-01', periods=100, freq='D')
    df = pd.DataFrame({
        'date': dates,
        'value': range(100)
    })

    train, test = temporal_split(df, 'date', train_ratio=0.7, gap='0D')

    # Check sizes
    assert len(train) == 70, f"Expected 70 train samples, got {len(train)}"
    assert len(test) == 30, f"Expected 30 test samples, got {len(test)}"

    # Check temporal order: all train should be before test
    assert train['date'].max() <= test['date'].min(), "Train should come before test"

    print(f"Temporal split: train={len(train)}, test={len(test)}")


def test_temporal_split_with_gap():
    """
    Test temporal split with a 7-day gap.

    Gap should remove data between train and test.
    """
    # Create daily data
    dates = pd.date_range('2026-01-01', periods=100, freq='D')
    df = pd.DataFrame({
        'date': dates,
        'value': range(100)
    })

    train, test = temporal_split(df, 'date', train_ratio=0.7, gap='7D')

    # Check that gap exists
    train_end = train['date'].max()
    test_start = test['date'].min()

    gap_days = (test_start - train_end).days

    assert gap_days >= 7, f"Gap should be at least 7 days, got {gap_days}"

    # Test set should be smaller due to gap
    assert len(test) < 30, f"Test set should be smaller with gap, got {len(test)}"

    print(f"Temporal split with gap: train={len(train)}, test={len(test)}, gap={gap_days} days")


def test_temporal_split_preserves_order():
    """
    Make sure temporal split doesn't shuffle data.

    Train should contain first 70% of chronological data.
    """
    dates = pd.date_range('2026-01-01', periods=100, freq='D')
    df = pd.DataFrame({
        'date': dates,
        'value': range(100)
    })

    train, test = temporal_split(df, 'date', train_ratio=0.7, gap='0D')

    # Train should have values 0-69
    assert train['value'].min() == 0
    assert train['value'].max() == 69

    # Test should have values 70-99
    assert test['value'].min() == 70
    assert test['value'].max() == 99

    print(f"Order preserved: train=[0-69], test=[70-99]")


def test_group_split_no_overlap():
    """
    Test that group split keeps groups separate.

    No group should appear in both train and test.
    """
    # Create data with 3 groups
    df = pd.DataFrame({
        'group_id': ['A'] * 20 + ['B'] * 20 + ['C'] * 20,
        'value': range(60)
    })

    train, test = group_split(df, 'group_id', train_ratio=0.67, random_state=42)

    # Get unique groups in each set
    train_groups = set(train['group_id'].unique())
    test_groups = set(test['group_id'].unique())

    # Check no overlap
    overlap = train_groups & test_groups
    assert len(overlap) == 0, f"Groups should not overlap, found: {overlap}"

    print(f"Group split: train_groups={train_groups}, test_groups={test_groups}")


def test_group_split_all_samples():
    """
    Test that group split includes all samples from each group.

    If group A goes to train, ALL rows of group A should be in train.
    """
    df = pd.DataFrame({
        'group_id': ['A'] * 10 + ['B'] * 10 + ['C'] * 10,
        'value': range(30)
    })

    train, test = group_split(df, 'group_id', train_ratio=0.67, random_state=42)

    # Check total samples preserved
    assert len(train) + len(test) == 30, "All samples should be preserved"

    # Check each group is complete in one split
    for group in df['group_id'].unique():
        group_in_train = (train['group_id'] == group).sum()
        group_in_test = (test['group_id'] == group).sum()

        # Each group should be entirely in one split (10 samples)
        assert (group_in_train == 10 and group_in_test == 0) or \
               (group_in_train == 0 and group_in_test == 10), \
            f"Group {group} should be entirely in train or test"

    print(f"All samples preserved and grouped correctly")

def test_group_split_preserves_time_order():
    """
    Group split should preserve chronological order
    within each group if time_column is provided.
    """
    # Create shuffled time-series data
    df = pd.DataFrame({
        'group_id': ['A'] * 5 + ['B'] * 5,
        'date': list(pd.date_range('2026-01-01', periods=5)) * 2
    })

    # Shuffle rows to break order
    df = df.sample(frac=1, random_state=42).reset_index(drop=True)

    train, test = group_split(
        df,
        group_column='group_id',
        train_ratio=0.5,
        random_state=42,
        time_column='date'
    )

    # Check monotonicity inside each split
    for split_df in [train, test]:
        for group in split_df['group_id'].unique():
            group_df = split_df[split_df['group_id'] == group]
            assert group_df['date'].is_monotonic_increasing, \
                f"Group {group} is not sorted by date"

def test_temporal_split_different_ratios():
    """
    Test temporal split with different train ratios.
    """
    dates = pd.date_range('2026-01-01', periods=100, freq='D')
    df = pd.DataFrame({
        'date': dates,
        'value': range(100)
    })

    # Test 50/50 split
    train_50, test_50 = temporal_split(df, 'date', train_ratio=0.5, gap='0D')
    assert len(train_50) == 50
    assert len(test_50) == 50

    # Test 80/20 split
    train_80, test_80 = temporal_split(df, 'date', train_ratio=0.8, gap='0D')
    assert len(train_80) == 80
    assert len(test_80) == 20

    print(f"Different ratios: 50/50 and 80/20 work correctly")


def test_invalid_ratio_raises_error():
    """
    Invalid train_ratio should raise error.

    Ratio must be between 0 and 1.
    """
    df = pd.DataFrame({
        'date': pd.date_range('2026-01-01', periods=50),
        'value': range(50)
    })

    # Test ratio > 1
    with pytest.raises(ValueError):
        temporal_split(df, 'date', train_ratio=1.5, gap='0D')

    # Test ratio <= 0
    with pytest.raises(ValueError):
        temporal_split(df, 'date', train_ratio=0.0, gap='0D')

    print("Invalid ratios raise errors correctly")


def test_invalid_gap_format():
    """
    Invalid gap format should raise error.
    """
    df = pd.DataFrame({
        'date': pd.date_range('2026-01-01', periods=50),
        'value': range(50)
    })

    # Test invalid gap format
    with pytest.raises(ValueError):
        temporal_split(df, 'date', train_ratio=0.7, gap='invalid')

    print("✓ Invalid gap format raises error correctly")


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

    test_temporal_split_basic()
    test_temporal_split_with_gap()
    test_temporal_split_preserves_order()
    test_group_split_no_overlap()
    test_group_split_all_samples()
    test_temporal_split_different_ratios()
    test_invalid_ratio_raises_error()
    test_invalid_gap_format()
