"""
Comprehensive tests for PAC learning module.

Tests for PAC learning algorithms from O'Donnell Chapter 3.
"""

import sys

import numpy as np
import pytest

sys.path.insert(0, "src")

import boofun as bf
from boofun.analysis.pac_learning import (
    PACLearner,
    lmn_algorithm,
    pac_learn_decision_tree,
    pac_learn_junta,
    pac_learn_low_degree,
    pac_learn_monotone,
    pac_learn_sparse_fourier,
    sample_function,
)


class TestSampleFunction:
    """Test sample_function utility."""

    def test_sample_returns_list(self):
        """sample_function should return list of tuples."""
        f = bf.majority(3)
        samples = sample_function(f, 10)

        assert isinstance(samples, list)
        assert len(samples) == 10

    def test_sample_tuple_format(self):
        """Each sample should be (input, output) tuple."""
        f = bf.AND(3)
        samples = sample_function(f, 5)

        for x, y in samples:
            assert isinstance(x, int)
            assert isinstance(y, int)
            assert 0 <= x < 8  # 2^3
            assert y in [0, 1]

    def test_sample_consistency(self):
        """Samples should be consistent with function."""
        f = bf.parity(3)
        samples = sample_function(f, 20)

        for x, y in samples:
            assert f.evaluate(x) == y

    def test_sample_with_rng(self):
        """sample_function should accept custom RNG."""
        f = bf.OR(3)
        rng = np.random.default_rng(42)

        samples = sample_function(f, 10, rng=rng)
        assert len(samples) == 10


class TestPACLearnLowDegree:
    """Test pac_learn_low_degree function."""

    def test_learn_constant(self):
        """Should learn constant functions perfectly."""
        f = bf.create([0, 0, 0, 0])

        hypothesis = pac_learn_low_degree(f, max_degree=0)

        assert hypothesis is not None

    def test_learn_dictator(self):
        """Should learn dictator (degree 1) functions."""
        f = bf.dictator(3, 0)

        hypothesis = pac_learn_low_degree(f, max_degree=1)

        assert hypothesis is not None

    def test_learn_majority(self):
        """Should learn majority with appropriate degree."""
        f = bf.majority(3)

        hypothesis = pac_learn_low_degree(f, max_degree=3)

        assert hypothesis is not None

    def test_epsilon_delta_parameters(self):
        """Should accept epsilon and delta parameters."""
        f = bf.AND(3)

        hypothesis = pac_learn_low_degree(f, max_degree=3, epsilon=0.1, delta=0.05)

        assert hypothesis is not None


class TestPACLearnJunta:
    """Test pac_learn_junta function."""

    def test_learn_dictator_as_1junta(self):
        """Dictator is a 1-junta."""
        f = bf.dictator(5, 2)

        result = pac_learn_junta(f, k=1)

        assert result is not None

    def test_learn_and_as_kjunta(self):
        """AND on k variables is a k-junta."""
        f = bf.AND(3)

        result = pac_learn_junta(f, k=3)

        assert result is not None

    def test_junta_learning_parameters(self):
        """Should accept epsilon and delta."""
        f = bf.OR(3)

        result = pac_learn_junta(f, k=3, epsilon=0.1, delta=0.05)

        assert result is not None


class TestLMNAlgorithm:
    """Test LMN (Linial-Mansour-Nisan) algorithm."""

    def test_lmn_basic(self):
        """LMN should work on basic functions."""
        f = bf.majority(3)

        # LMN may have different signature
        try:
            result = lmn_algorithm(f)
            assert result is not None
        except TypeError:
            # Try with different params
            result = lmn_algorithm(f, 100)  # num_samples
            assert result is not None

    def test_lmn_with_samples(self):
        """LMN should accept sample count."""
        f = bf.AND(3)

        try:
            result = lmn_algorithm(f, 100)
            assert result is not None
        except TypeError:
            pass  # Different API


class TestPACLearnSparseFourier:
    """Test pac_learn_sparse_fourier function."""

    def test_learn_parity(self):
        """Parity has single non-zero Fourier coefficient."""
        f = bf.parity(3)

        result = pac_learn_sparse_fourier(f, sparsity=1)

        assert result is not None

    def test_learn_dictator(self):
        """Dictator has 2 non-zero coefficients (constant + degree 1)."""
        f = bf.dictator(3, 0)

        result = pac_learn_sparse_fourier(f, sparsity=2)

        assert result is not None


class TestPACLearnDecisionTree:
    """Test pac_learn_decision_tree function."""

    def test_learn_and(self):
        """AND has simple decision tree."""
        f = bf.AND(3)

        result = pac_learn_decision_tree(f, max_depth=3)

        assert result is not None

    def test_learn_or(self):
        """OR has simple decision tree."""
        f = bf.OR(3)

        result = pac_learn_decision_tree(f, max_depth=3)

        assert result is not None


class TestPACLearnMonotone:
    """Test pac_learn_monotone function."""

    def test_learn_and(self):
        """AND is monotone."""
        f = bf.AND(3)

        result = pac_learn_monotone(f)

        assert result is not None

    def test_learn_or(self):
        """OR is monotone."""
        f = bf.OR(3)

        result = pac_learn_monotone(f)

        assert result is not None

    def test_learn_majority(self):
        """Majority is monotone."""
        f = bf.majority(3)

        result = pac_learn_monotone(f)

        assert result is not None


class TestPACLearner:
    """Test PACLearner class."""

    def test_learner_init(self):
        """PACLearner should initialize with a function."""
        f = bf.majority(3)
        learner = PACLearner(f)

        assert learner is not None

    def test_learner_with_params(self):
        """PACLearner should accept epsilon and delta."""
        f = bf.AND(3)

        try:
            learner = PACLearner(f, epsilon=0.1, delta=0.05)
            assert learner is not None
        except TypeError:
            learner = PACLearner(f)
            assert learner is not None

    def test_learner_learn_method(self):
        """PACLearner should have learning methods."""
        f = bf.majority(3)
        learner = PACLearner(f)

        # Check available methods
        methods = [m for m in dir(learner) if not m.startswith("_")]
        assert len(methods) > 0

    def test_learner_different_functions(self):
        """PACLearner should work with different functions."""
        for func in [bf.AND(3), bf.OR(3), bf.parity(3)]:
            learner = PACLearner(func)
            assert learner is not None


class TestPACLearningAccuracy:
    """Test that PAC learning achieves reasonable accuracy."""

    def test_low_degree_accuracy(self):
        """Low-degree learning should be accurate for low-degree functions."""
        f = bf.dictator(4, 0)

        hypothesis = pac_learn_low_degree(f, max_degree=1, epsilon=0.1)

        if hypothesis is not None and hasattr(hypothesis, "evaluate"):
            # Test accuracy
            errors = 0
            for x in range(16):
                if hypothesis.evaluate(x) != f.evaluate(x):
                    errors += 1

            error_rate = errors / 16
            assert error_rate <= 0.5  # Should be reasonably accurate

    def test_junta_accuracy(self):
        """Junta learning should work well for small juntas."""
        f = bf.AND(2)  # 2-junta

        hypothesis = pac_learn_junta(f, k=2, epsilon=0.1)

        # Just check it returns something
        assert hypothesis is not None


class TestPACLearningEdgeCases:
    """Test edge cases for PAC learning."""

    def test_empty_function(self):
        """Should handle constant zero function."""
        f = bf.create([0, 0, 0, 0])

        result = pac_learn_low_degree(f, max_degree=0)
        assert result is not None

    def test_small_n(self):
        """Should work for n=1."""
        f = bf.create([0, 1])

        result = pac_learn_low_degree(f, max_degree=1)
        assert result is not None


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