import pytest
import numpy as np
from weighted_mcc.binary import weighted_mcc

class TestWeightedMCC:
    def test_weighted_mcc_perfect_prediction(self):
        """Test that MCC is 1.0 for perfect prediction with uniform weights."""
        y_true = np.array([1, 0, 1, 0])
        y_pred = np.array([1, 0, 1, 0])
        weights = np.ones_like(y_true)
        result = weighted_mcc(y_true, y_pred, weights)
        assert np.isclose(result, 1.0)

    def test_weighted_mcc_total_disagreement(self):
        """Test that MCC is -1.0 for total disagreement."""
        y_true = np.array([1, 1, 0, 0])
        y_pred = np.array([0, 0, 1, 1])
        weights = np.ones_like(y_true)
        result = weighted_mcc(y_true, y_pred, weights)
        assert np.isclose(result, -1.0)

    def test_weighted_mcc_random_prediction(self):
        """Test that MCC is approx 0.0 for random prediction (uncorrelated)."""
        # Create a balanced scenario where TP*TN approx FP*FN
        y_true = np.array([1, 0, 1, 0])
        y_pred = np.array([1, 1, 0, 0]) # Half right, half wrong
        weights = np.ones_like(y_true)
        # TP=1, TN=1, FP=1, FN=1 -> Num=1*1 - 1*1 = 0
        result = weighted_mcc(y_true, y_pred, weights)
        assert np.isclose(result, 0.0)

    def test_weighted_mcc_unweighted_equivalence(self):
        """Test weighted MCC with uniform weights equals manual standard calculation."""
        y_true = np.array([1, 1, 0, 1, 0, 0])
        y_pred = np.array([1, 0, 0, 1, 1, 0])
        # TP=2, TN=2, FP=1, FN=1
        # MCC = (4 - 1) / sqrt(3*3*3*3) = 3 / 9 = 0.3333
        weights = np.ones_like(y_true)
        result = weighted_mcc(y_true, y_pred, weights)
        assert np.isclose(result, 1.0/3.0)

    def test_weighted_mcc_high_weight_importance(self):
        """Test that high weights on correct samples drive the score higher."""
        y_true = np.array([1, 0])
        y_pred = np.array([1, 1]) # 1 Correct (Pos), 1 Wrong (Neg)
        
        # Case 1: Weight the correct prediction heavily
        weights_good = np.array([100.0, 1.0])
        # TP=100, TN=0, FP=1, FN=0. 
        # Note: MCC undefined if denominator is 0.
        # Let's add dummy samples to prevent 0 denominator
        y_true = np.array([1, 0, 1, 0])
        y_pred = np.array([1, 1, 0, 0]) # Index 0 correct, Index 1 wrong, Index 2 wrong, Index 3 correct
        
        # Heavy weight on correct indices (0 and 3)
        w_good = np.array([10.0, 1.0, 1.0, 10.0])
        score_good = weighted_mcc(y_true, y_pred, w_good)
        
        # Heavy weight on wrong indices (1 and 2)
        w_bad = np.array([1.0, 10.0, 10.0, 1.0])
        score_bad = weighted_mcc(y_true, y_pred, w_bad)
        
        assert score_good > score_bad

    def test_weighted_mcc_invalid_input_shape(self):
        """Test that ValueError is raised when input arrays have different lengths."""
        y_true = np.array([1, 0])
        y_pred = np.array([1, 0, 1])
        with pytest.raises(ValueError):
            weighted_mcc(y_true, y_pred, np.ones(2))

    def test_weighted_mcc_single_class_input(self):
        """Test behavior when y_true contains only one class (denominator 0 check)."""
        y_true = np.array([1, 1, 1])
        y_pred = np.array([1, 1, 1])
        weights = np.ones(3)
        # Standard implementations often return 0.0 or raise warning when valid MCC cannot be computed
        # Assuming implementation returns 0.0 for undefined cases as per paper Remark 1 continuity
        assert weighted_mcc(y_true, y_pred, weights) == 0.0