"""
Unit tests for parallel resampling functionality.

Tests the n_jobs parameter in resample() function.
"""

import pytest
import numpy as np
import pandas as pd
import time

from mlpy.tasks import TaskClassif, TaskRegr
from mlpy.learners.sklearn import LearnerRandomForestClassifier, LearnerRandomForestRegressor
from mlpy.measures import MeasureClassifAccuracy, MeasureRegrRMSE
from mlpy.resamplings import ResamplingCV
from mlpy.resample import resample


# Mark all tests as requiring joblib
pytest.importorskip("joblib")
pytest.importorskip("sklearn")


@pytest.fixture
def simple_classif_task():
    """Create a simple classification task."""
    from sklearn.datasets import make_classification

    X, y = make_classification(
        n_samples=200,
        n_features=10,
        n_informative=5,
        n_redundant=2,
        n_classes=3,
        random_state=42
    )

    df = pd.DataFrame(X, columns=[f'feature_{i}' for i in range(X.shape[1])])
    df['target'] = y

    task = TaskClassif(data=df, target='target')
    return task


@pytest.fixture
def simple_regr_task():
    """Create a simple regression task."""
    from sklearn.datasets import make_regression

    X, y = make_regression(
        n_samples=200,
        n_features=10,
        n_informative=5,
        random_state=42
    )

    df = pd.DataFrame(X, columns=[f'feature_{i}' for i in range(X.shape[1])])
    df['target'] = y

    task = TaskRegr(data=df, target='target')
    return task


def test_sequential_vs_parallel_classification(simple_classif_task):
    """Test that n_jobs=1 (sequential) and n_jobs=2 give same results."""
    task = simple_classif_task
    learner = LearnerRandomForestClassifier(n_estimators=10, random_state=42)
    measure = MeasureClassifAccuracy()
    resampling = ResamplingCV(folds=5)

    # Sequential
    result_seq = resample(
        task=task,
        learner=learner,
        resampling=resampling,
        measures=[measure],
        n_jobs=1
    )

    # Parallel with 2 jobs
    result_par = resample(
        task=task,
        learner=learner.clone(),
        resampling=resampling.clone(),
        measures=[measure],
        n_jobs=2
    )

    # Check that both completed successfully
    assert result_seq.n_iters == 5
    assert result_par.n_iters == 5
    assert result_seq.n_errors == 0
    assert result_par.n_errors == 0

    # Scores should be identical (same random seed)
    score_seq = result_seq.score('classif.acc', average='mean')
    score_par = result_par.score('classif.acc', average='mean')

    # Allow small numerical differences due to floating point
    assert abs(score_seq - score_par) < 0.01, f"Sequential: {score_seq}, Parallel: {score_par}"


def test_n_jobs_minus_one(simple_classif_task):
    """Test that n_jobs=-1 (all cores) works."""
    task = simple_classif_task
    learner = LearnerRandomForestClassifier(n_estimators=10, random_state=42)
    measure = MeasureClassifAccuracy()

    result = resample(
        task=task,
        learner=learner,
        resampling=ResamplingCV(folds=5),
        measures=[measure],
        n_jobs=-1
    )

    assert result.n_iters == 5
    assert result.n_errors == 0
    score = result.score('classif.acc', average='mean')
    assert 0 <= score <= 1


def test_parallel_regression(simple_regr_task):
    """Test parallel execution with regression."""
    task = simple_regr_task
    learner = LearnerRandomForestRegressor(n_estimators=10, random_state=42)
    measure = MeasureRegrRMSE()

    result = resample(
        task=task,
        learner=learner,
        resampling=ResamplingCV(folds=5),
        measures=[measure],
        n_jobs=2
    )

    assert result.n_iters == 5
    assert result.n_errors == 0
    score = result.score('regr.rmse', average='mean')
    assert score > 0


def test_parallel_speedup_exists(simple_classif_task):
    """Test that parallel execution is faster than sequential (roughly)."""
    task = simple_classif_task
    # Use larger model to make training time more significant
    learner = LearnerRandomForestClassifier(n_estimators=50, random_state=42)
    measure = MeasureClassifAccuracy()
    resampling = ResamplingCV(folds=5)

    # Sequential
    start = time.time()
    result_seq = resample(
        task=task,
        learner=learner,
        resampling=resampling,
        measures=[measure],
        n_jobs=1
    )
    time_seq = time.time() - start

    # Parallel
    start = time.time()
    result_par = resample(
        task=task,
        learner=learner.clone(),
        resampling=resampling.clone(),
        measures=[measure],
        n_jobs=2
    )
    time_par = time.time() - start

    print(f"\nSequential time: {time_seq:.2f}s")
    print(f"Parallel time: {time_par:.2f}s")
    print(f"Speedup: {time_seq/time_par:.2f}x")

    # Parallel should be faster (or at least not much slower due to overhead)
    # We don't enforce strict speedup because it depends on system load
    # Just check that both work
    assert result_seq.n_errors == 0
    assert result_par.n_errors == 0


def test_n_jobs_compatibility_with_backend():
    """Test that explicit backend takes precedence over n_jobs."""
    from sklearn.datasets import make_classification
    from mlpy.parallel.joblib import BackendJoblib

    X, y = make_classification(n_samples=100, n_features=5, random_state=42)
    df = pd.DataFrame(X, columns=[f'f{i}' for i in range(X.shape[1])])
    df['target'] = y

    task = TaskClassif(data=df, target='target')
    learner = LearnerRandomForestClassifier(n_estimators=10, random_state=42)
    measure = MeasureClassifAccuracy()

    # Create explicit backend
    backend = BackendJoblib(n_jobs=3)

    # Pass both n_jobs and backend
    # Backend should take precedence
    result = resample(
        task=task,
        learner=learner,
        resampling=ResamplingCV(folds=5),
        measures=[measure],
        n_jobs=1,  # This should be ignored
        backend=backend  # This should be used
    )

    assert result.n_iters == 5
    assert result.n_errors == 0


def test_multiple_measures_parallel(simple_classif_task):
    """Test parallel execution with multiple measures."""
    from mlpy.measures import MeasureClassifF1

    task = simple_classif_task
    learner = LearnerRandomForestClassifier(n_estimators=10, random_state=42)
    measures = [MeasureClassifAccuracy(), MeasureClassifF1()]

    result = resample(
        task=task,
        learner=learner,
        resampling=ResamplingCV(folds=5),
        measures=measures,
        n_jobs=2
    )

    assert result.n_iters == 5
    assert result.n_errors == 0

    # Check both measures computed
    acc = result.score('classif.acc', average='mean')
    f1 = result.score('classif.f1', average='mean')

    assert 0 <= acc <= 1
    assert 0 <= f1 <= 1


if __name__ == '__main__':
    pytest.main([__file__, '-v', '-s'])
