import pytest
import numpy as np
from weighted_mcc.utils import compute_weighted_means, compute_covariance_matrices

class TestComputeWeightedMeans:
    def test_compute_weighted_means_simple(self):
        """Test mean calculation with a simple 2-class example."""
        # T (N=2, K=2)
        t = np.array([[1, 0], [0, 1]])
        weights = np.array([3.0, 1.0]) # Trace = 4
        # Mean = (1/4) * (t.T @ weights)
        # Class 0: (1*3 + 0*1)/4 = 0.75
        # Class 1: (0*3 + 1*1)/4 = 0.25
        means = compute_weighted_means(t, weights)
        expected = np.array([0.75, 0.25])
        assert np.allclose(means, expected)

    def test_compute_weighted_means_shape(self):
        """Verify output shape is (n_classes,)."""
        t = np.eye(3)
        w = np.ones(3)
        means = compute_weighted_means(t, w)
        assert means.shape == (3,)

class TestComputeCovarianceMatrices:
    def test_compute_covariance_matrices_simple(self):
        t = np.array([[1], [0]])
        c = np.array([[1], [0]])
        w = np.array([1.0, 1.0])
        # Function computes means internally; unpack the resulting tuple
        R_tc, _, _ = compute_covariance_matrices(t, c, w)
        assert np.isclose(R_tc[0,0], 0.25)

    def test_compute_covariance_matrices_symmetry(self):
        t = np.random.randint(0, 2, (10, 3))
        w = np.random.rand(10)
        _, R_tt, _ = compute_covariance_matrices(t, t, w)
        assert np.allclose(R_tt, R_tt.T)

    def test_compute_covariance_matrices_output_shape(self):
        t = np.zeros((5, 4))
        w = np.ones(5)
        R_tc, _, _ = compute_covariance_matrices(t, t, w)
        assert R_tc.shape == (4, 4)