"""
Unit tests for the PPMI library.
"""

import pytest
import math
import numpy as np
from ppmi import PMI, calculate_pmi_from_probabilities, calculate_ppmi_from_probabilities
from ppmi.pmi import calculate_pmi_from_counts, calculate_ppmi_from_counts


class TestPMIFunctions:
    """Test standalone PMI calculation functions."""
    
    def test_calculate_pmi_basic(self):
        """Test basic PMI calculation."""
        pmi = calculate_pmi_from_probabilities(0.1, 0.3, 0.4, base=2)
        expected = math.log2(0.1 / (0.3 * 0.4))
        assert abs(pmi - expected) < 1e-10
    
    def test_calculate_pmi_base_10(self):
        """Test PMI with base 10."""
        pmi = calculate_pmi_from_probabilities(0.1, 0.3, 0.4, base=10)
        expected = math.log10(0.1 / (0.3 * 0.4))
        assert abs(pmi - expected) < 1e-10
    
    def test_calculate_pmi_base_e(self):
        """Test PMI with natural logarithm."""
        pmi = calculate_pmi_from_probabilities(0.1, 0.3, 0.4, base='e')
        expected = math.log(0.1 / (0.3 * 0.4))
        assert abs(pmi - expected) < 1e-10
    
    def test_calculate_pmi_custom_base(self):
        """Test PMI with custom base."""
        pmi = calculate_pmi_from_probabilities(0.1, 0.3, 0.4, base=5)
        expected = math.log(0.1 / (0.3 * 0.4)) / math.log(5)
        assert abs(pmi - expected) < 1e-10
    
    def test_calculate_pmi_positive_association(self):
        """Test PMI for positive association."""
        # When P(x,y) > P(x) * P(y), PMI should be positive
        pmi = calculate_pmi_from_probabilities(0.5, 0.6, 0.7, base=2)
        assert pmi > 0
    
    def test_calculate_pmi_negative_association(self):
        """Test PMI for negative association."""
        # When P(x,y) < P(x) * P(y), PMI should be negative
        pmi = calculate_pmi_from_probabilities(0.01, 0.3, 0.4, base=2)
        assert pmi < 0
    
    def test_calculate_ppmi_positive(self):
        """Test PPMI when PMI is positive."""
        ppmi = calculate_ppmi_from_probabilities(0.5, 0.6, 0.7, base=2)
        pmi = calculate_pmi_from_probabilities(0.5, 0.6, 0.7, base=2)
        assert ppmi == pmi
        assert ppmi > 0
    
    def test_calculate_ppmi_negative(self):
        """Test PPMI when PMI is negative."""
        ppmi = calculate_ppmi_from_probabilities(0.01, 0.3, 0.4, base=2)
        assert ppmi == 0
    
    def test_calculate_pmi_invalid_probability(self):
        """Test that invalid probabilities raise ValueError."""
        with pytest.raises(ValueError):
            calculate_pmi_from_probabilities(0, 0.3, 0.4)
        with pytest.raises(ValueError):
            calculate_pmi_from_probabilities(0.1, 0, 0.4)
        with pytest.raises(ValueError):
            calculate_pmi_from_probabilities(0.1, 0.3, 0)
        with pytest.raises(ValueError):
            calculate_pmi_from_probabilities(1.5, 0.3, 0.4)


class TestPMIFromCounts:
    """Test PMI calculation from counts."""
    
    def test_calculate_pmi_from_counts_basic(self):
        """Test PMI calculation from counts."""
        pmi = calculate_pmi_from_counts(10, 50, 40, 1000, base=2)
        
        # Manually calculate expected value
        p_xy = 10 / 1000
        p_x = 50 / 1000
        p_y = 40 / 1000
        expected = math.log2(p_xy / (p_x * p_y))
        
        assert abs(pmi - expected) < 1e-10
    
    def test_calculate_ppmi_from_counts(self):
        """Test PPMI calculation from counts."""
        ppmi = calculate_ppmi_from_counts(10, 50, 40, 1000, base=2)
        pmi = calculate_pmi_from_counts(10, 50, 40, 1000, base=2)
        assert ppmi == max(0, pmi)
    
    def test_calculate_pmi_from_counts_invalid(self):
        """Test that invalid counts raise ValueError."""
        with pytest.raises(ValueError):
            calculate_pmi_from_counts(0, 50, 40, 1000)
        with pytest.raises(ValueError):
            calculate_pmi_from_counts(10, 0, 40, 1000)
        with pytest.raises(ValueError):
            calculate_pmi_from_counts(10, 50, 0, 1000)
        with pytest.raises(ValueError):
            calculate_pmi_from_counts(10, 50, 40, 0)


class TestPMIClass:
    """Test PMI class functionality."""
    
    def test_initialization(self):
        """Test PMI class initialization."""
        pmi_calc = PMI()
        assert pmi_calc.total == 0
        assert len(pmi_calc.x_counts) == 0
        assert len(pmi_calc.y_counts) == 0
        assert len(pmi_calc.co_occurrence) == 0
    
    def test_add_observation(self):
        """Test adding a single observation."""
        pmi_calc = PMI()
        pmi_calc.add_observation("word1", "word2")
        
        assert pmi_calc.total == 1
        assert pmi_calc.x_counts["word1"] == 1
        assert pmi_calc.y_counts["word2"] == 1
        assert pmi_calc.co_occurrence["word1"]["word2"] == 1
    
    def test_add_observation_with_count(self):
        """Test adding observation with custom count."""
        pmi_calc = PMI()
        pmi_calc.add_observation("word1", "word2", count=5)
        
        assert pmi_calc.total == 5
        assert pmi_calc.x_counts["word1"] == 5
        assert pmi_calc.y_counts["word2"] == 5
        assert pmi_calc.co_occurrence["word1"]["word2"] == 5
    
    def test_add_multiple_observations(self):
        """Test adding multiple observations."""
        pmi_calc = PMI()
        pmi_calc.add_observation("word1", "word2")
        pmi_calc.add_observation("word1", "word3")
        pmi_calc.add_observation("word2", "word3")
        
        assert pmi_calc.total == 3
        assert pmi_calc.x_counts["word1"] == 2
        assert pmi_calc.x_counts["word2"] == 1
    
    def test_add_observations_batch(self):
        """Test batch observation addition."""
        pmi_calc = PMI()
        pairs = [
            ("word1", "word2"),
            ("word1", "word3"),
            ("word2", "word3", 5),
        ]
        pmi_calc.add_observations_batch(pairs)
        
        assert pmi_calc.total == 7
        assert pmi_calc.co_occurrence["word2"]["word3"] == 5
    
    def test_calculate_pmi_from_class(self):
        """Test PMI calculation from PMI class."""
        pmi_calc = PMI()
        pmi_calc.add_observation("word1", "word2", count=10)
        pmi_calc.add_observation("word1", "word3", count=40)
        pmi_calc.add_observation("word2", "word3", count=50)
        
        pmi = pmi_calc.calculate_pmi("word1", "word2", base=2)
        
        # Manually verify
        p_xy = 10 / 100
        p_x = 50 / 100
        p_y = 10 / 100
        expected = math.log2(p_xy / (p_x * p_y))
        
        assert abs(pmi - expected) < 1e-10
    
    def test_calculate_pmi_nonexistent_pair(self):
        """Test PMI for pair that was never observed."""
        pmi_calc = PMI()
        pmi_calc.add_observation("word1", "word2")
        
        pmi = pmi_calc.calculate_pmi("word1", "word3")
        assert pmi is None
    
    def test_calculate_ppmi_from_class(self):
        """Test PPMI calculation from PMI class."""
        pmi_calc = PMI()
        pmi_calc.add_observation("word1", "word2", count=100)
        pmi_calc.add_observation("word3", "word4", count=1)
        
        ppmi1 = pmi_calc.calculate_ppmi("word1", "word2")
        ppmi2 = pmi_calc.calculate_ppmi("word3", "word4")
        
        assert ppmi1 >= 0
        assert ppmi2 >= 0
    
    def test_calculate_npmi(self):
        """Test NPMI calculation."""
        pmi_calc = PMI()
        pmi_calc.add_observation("word1", "word2", count=50)
        pmi_calc.add_observation("word1", "word3", count=25)
        pmi_calc.add_observation("word3", "word2", count=25)
        
        npmi = pmi_calc.calculate_npmi("word1", "word2")
        
        # NPMI should be between -1 and 1
        assert -1 <= npmi <= 1
    
    def test_get_all_pmi(self):
        """Test getting all PMI scores."""
        pmi_calc = PMI()
        pmi_calc.add_observation("word1", "word2")
        pmi_calc.add_observation("word1", "word3")
        pmi_calc.add_observation("word2", "word3")
        
        all_pmi = pmi_calc.get_all_pmi()
        
        assert len(all_pmi) == 3
        assert ("word1", "word2") in all_pmi
        assert ("word1", "word3") in all_pmi
        assert ("word2", "word3") in all_pmi
    
    def test_get_all_ppmi(self):
        """Test getting all PPMI scores."""
        pmi_calc = PMI()
        pmi_calc.add_observation("word1", "word2")
        pmi_calc.add_observation("word1", "word3")
        
        all_ppmi = pmi_calc.get_all_ppmi()
        
        assert len(all_ppmi) == 2
        for value in all_ppmi.values():
            assert value >= 0  # PPMI is always non-negative
    
    def test_to_matrix_ppmi(self):
        """Test matrix export with PPMI."""
        pmi_calc = PMI()
        pmi_calc.add_observation("word1", "context1")
        pmi_calc.add_observation("word1", "context2")
        pmi_calc.add_observation("word2", "context1")
        
        matrix, x_labels, y_labels = pmi_calc.to_matrix(metric='ppmi')
        
        assert isinstance(matrix, np.ndarray)
        assert matrix.shape == (len(x_labels), len(y_labels))
        assert "word1" in x_labels
        assert "word2" in x_labels
        assert "context1" in y_labels
        assert "context2" in y_labels
        
        # All PPMI values should be non-negative
        assert np.all(matrix >= 0)
    
    def test_to_matrix_pmi(self):
        """Test matrix export with PMI."""
        pmi_calc = PMI()
        # Add observations with some that co-occur more/less than expected
        pmi_calc.add_observation("word1", "context1", count=100)
        pmi_calc.add_observation("word1", "context2", count=1)
        pmi_calc.add_observation("word2", "context1", count=50)
        pmi_calc.add_observation("word2", "context2", count=50)
        
        matrix, x_labels, y_labels = pmi_calc.to_matrix(metric='pmi')
        
        assert isinstance(matrix, np.ndarray)
        # With varied observations, we should get non-zero PMI values
        assert np.any(matrix != 0)
    
    def test_to_matrix_npmi(self):
        """Test matrix export with NPMI."""
        pmi_calc = PMI()
        pmi_calc.add_observation("word1", "context1", count=50)
        pmi_calc.add_observation("word2", "context2", count=30)
        
        matrix, x_labels, y_labels = pmi_calc.to_matrix(metric='npmi')
        
        assert isinstance(matrix, np.ndarray)
        # NPMI should be between -1 and 1 where defined
        non_zero = matrix[matrix != 0]
        if len(non_zero) > 0:
            assert np.all(non_zero >= -1)
            assert np.all(non_zero <= 1)


class TestEdgeCases:
    """Test edge cases and special scenarios."""
    
    def test_symmetric_probabilities(self):
        """Test PMI when P(x) == P(y)."""
        pmi = calculate_pmi_from_probabilities(0.25, 0.5, 0.5, base=2)
        expected = math.log2(0.25 / 0.25)
        assert abs(pmi - expected) < 1e-10
    
    def test_perfect_correlation(self):
        """Test PMI for strong positive association."""
        pmi_calc = PMI()
        # Create a scenario where x and y co-occur more than expected
        # x and y together
        pmi_calc.add_observation("x", "y", count=50)
        # x with other contexts (less frequent)
        pmi_calc.add_observation("x", "z", count=5)
        # y with other words (less frequent)
        pmi_calc.add_observation("w", "y", count=5)
        # Other unrelated pairs
        pmi_calc.add_observation("w", "z", count=40)
        
        pmi = pmi_calc.calculate_pmi("x", "y")
        # Should be positive since they co-occur more than expected by chance
        assert pmi > 0
    
    def test_independence(self):
        """Test PMI for independent events."""
        # When P(x,y) = P(x) * P(y), PMI should be close to 0
        p_x = 0.3
        p_y = 0.4
        p_xy = p_x * p_y  # Independent events
        
        pmi = calculate_pmi_from_probabilities(p_xy, p_x, p_y, base=2)
        assert abs(pmi) < 1e-10
    
    def test_single_occurrence(self):
        """Test with single occurrence."""
        pmi_calc = PMI()
        pmi_calc.add_observation("rare", "word")
        
        pmi = pmi_calc.calculate_pmi("rare", "word")
        assert pmi is not None
    
    def test_large_counts(self):
        """Test with large counts."""
        pmi_calc = PMI()
        pmi_calc.add_observation("common", "word", count=1000000)
        
        assert pmi_calc.total == 1000000
        pmi = pmi_calc.calculate_pmi("common", "word")
        assert pmi is not None


class TestWeightedPPMI:
    """Test weighted PPMI with alpha parameter."""
    
    def test_ppmi_alpha_default(self):
        """Test that alpha=1.0 gives same result as standard PPMI."""
        pmi_calc = PMI()
        pmi_calc.add_observation("word1", "context1", count=10)
        pmi_calc.add_observation("word1", "context2", count=5)
        pmi_calc.add_observation("word2", "context1", count=3)
        
        ppmi_standard = pmi_calc.calculate_ppmi("word1", "context1", alpha=1.0)
        ppmi_no_alpha = pmi_calc.calculate_ppmi("word1", "context1")
        
        assert abs(ppmi_standard - ppmi_no_alpha) < 1e-10
    
    def test_ppmi_alpha_075(self):
        """Test PPMI with alpha=0.75 (recommended value)."""
        pmi_calc = PMI()
        # Create scenario with common and rare contexts
        pmi_calc.add_observation("word", "common_context", count=100)
        pmi_calc.add_observation("word", "rare_context", count=2)
        pmi_calc.add_observation("other", "common_context", count=50)
        
        ppmi_standard = pmi_calc.calculate_ppmi("word", "rare_context", alpha=1.0)
        ppmi_smoothed = pmi_calc.calculate_ppmi("word", "rare_context", alpha=0.75)
        
        # Alpha < 1 should decrease PPMI for rare contexts
        assert ppmi_smoothed < ppmi_standard
    
    def test_ppmi_alpha_rare_vs_common(self):
        """Test that alpha affects rare and common contexts differently."""
        pmi_calc = PMI()
        # Create data with clear common and rare contexts
        pmi_calc.add_observation("target", "common", count=100)
        pmi_calc.add_observation("target", "rare", count=5)
        # Add other words to create distribution
        for i in range(10):
            pmi_calc.add_observation(f"word{i}", "common", count=10)
        
        # Standard PPMI
        ppmi_common_1 = pmi_calc.calculate_ppmi("target", "common", alpha=1.0)
        ppmi_rare_1 = pmi_calc.calculate_ppmi("target", "rare", alpha=1.0)
        
        # Smoothed PPMI with alpha=0.75
        ppmi_common_075 = pmi_calc.calculate_ppmi("target", "common", alpha=0.75)
        ppmi_rare_075 = pmi_calc.calculate_ppmi("target", "rare", alpha=0.75)
        
        # Smoothing should change the values
        assert ppmi_common_075 != ppmi_common_1
        assert ppmi_rare_075 != ppmi_rare_1
        
        # With alpha < 1, rare contexts get higher probability,
        # which typically reduces their PPMI
        assert ppmi_rare_075 < ppmi_rare_1
    
    def test_get_all_ppmi_with_alpha(self):
        """Test get_all_ppmi with alpha parameter."""
        pmi_calc = PMI()
        pmi_calc.add_observation("w1", "c1", count=10)
        pmi_calc.add_observation("w1", "c2", count=2)
        pmi_calc.add_observation("w2", "c1", count=5)
        
        all_ppmi_standard = pmi_calc.get_all_ppmi(alpha=1.0)
        all_ppmi_smoothed = pmi_calc.get_all_ppmi(alpha=0.75)
        
        assert len(all_ppmi_standard) == len(all_ppmi_smoothed)
        # Values should be different
        assert all_ppmi_standard[("w1", "c2")] != all_ppmi_smoothed[("w1", "c2")]
    
    def test_to_matrix_with_alpha(self):
        """Test matrix export with alpha parameter."""
        pmi_calc = PMI()
        pmi_calc.add_observation("w1", "c1", count=10)
        pmi_calc.add_observation("w2", "c2", count=5)
        
        matrix_standard, _, _ = pmi_calc.to_matrix(metric='ppmi', alpha=1.0)
        matrix_smoothed, _, _ = pmi_calc.to_matrix(metric='ppmi', alpha=0.75)
        
        # Matrices should have same shape
        assert matrix_standard.shape == matrix_smoothed.shape
        # But different values
        assert not np.allclose(matrix_standard, matrix_smoothed)
    
    def test_smoothed_context_prob(self):
        """Test the smoothed context probability calculation."""
        pmi_calc = PMI()
        pmi_calc.add_observation("w1", "common", count=100)
        pmi_calc.add_observation("w2", "rare", count=1)
        
        # Standard probability
        p_common_1 = pmi_calc._calculate_smoothed_context_prob("common", alpha=1.0)
        p_rare_1 = pmi_calc._calculate_smoothed_context_prob("rare", alpha=1.0)
        
        # Smoothed probability
        p_common_075 = pmi_calc._calculate_smoothed_context_prob("common", alpha=0.75)
        p_rare_075 = pmi_calc._calculate_smoothed_context_prob("rare", alpha=0.75)
        
        # Rare context should get higher probability with alpha < 1
        assert p_rare_075 > p_rare_1
        # Common context should get lower probability
        assert p_common_075 < p_common_1


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