import pytest
import numpy as np
from weighted_mcc.multiclass import (
    extended_corr_coef,
    mpc_trace_ratio,
    mpc_mean_corr
)

class TestExtendedCorrCoef:
    def test_extended_corr_coef_perfect(self, multiclass_data):
        y_true, _, weights = multiclass_data
        # Perfect prediction: y_pred = y_true
        result = extended_corr_coef(y_true, y_true, weights)
        assert np.isclose(result, 1.0)

    def test_extended_corr_coef_range(self, multiclass_data):
        y_true, y_pred, weights = multiclass_data
        result = extended_corr_coef(y_true, y_pred, weights)
        assert -1.0 <= result <= 1.0

class TestMPCTraceRatio:
    def test_mpc_trace_ratio_perfect(self, multiclass_data):
        y_true, _, weights = multiclass_data
        result = mpc_trace_ratio(y_true, y_true, weights)
        assert np.isclose(result, 1.0)

    def test_mpc_trace_ratio_uniform(self, multiclass_data):
        """Consistency check with uniform weights."""
        y_true, y_pred, _ = multiclass_data
        w_uni = np.ones(len(y_true))
        res = mpc_trace_ratio(y_true, y_pred, w_uni)
        assert not np.isnan(res)

class TestMPCMeanCorr:
    def test_mpc_mean_corr_perfect(self, multiclass_data):
        y_true, _, weights = multiclass_data
        result = mpc_mean_corr(y_true, y_true, weights)
        assert np.isclose(result, 1.0)

    def test_mpc_mean_corr_zero_variance_handling(self):
        """Test MPC2 behavior when class variance is zero (should be 0 by continuity)."""
        # Class 1 never appears in Truth or Pred
        y_true = np.array([[1, 0], [1, 0]])
        y_pred = np.array([[1, 0], [1, 0]])
        weights = np.array([1.0, 1.0])
        
        # For class 2 (index 1), R_tt and R_cc will be 0.
        # Implementation should handle divide by zero.
        result = mpc_mean_corr(y_true, y_pred, weights)
        # (1.0 for class 1 + 0.0 for class 2) / 2 = 0.5
        assert np.isclose(result, 0.5)