# weighted_mcc/utils.py
import numpy as np
from numpy.typing import NDArray
from typing import Tuple

def compute_weighted_means(
    y_one_hot: NDArray[np.int_],
    sample_weights: NDArray[np.float64]
) -> NDArray[np.float64]:
    """
    Computes the weighted mean vector (t_bar or c_bar) for the multiclass case.
    
    Args:
        y_one_hot: 2D Array [n_samples, n_classes].
        sample_weights: 1D Array [n_samples].

    Returns:
        1D Array [n_classes] representing the weighted means.
    """
    # Ensure inputs are correct type/shape
    y = np.asarray(y_one_hot)
    w = np.asarray(sample_weights)
    
    tr_S = np.sum(w)
    if tr_S == 0:
        return np.zeros(y.shape[1])
        
    # Weighted sum: w^T * Y
    weighted_sum = np.dot(w, y)
    
    return weighted_sum / tr_S

def compute_covariance_matrices(
    y_true_one_hot: NDArray[np.int_],
    y_pred_one_hot: NDArray[np.int_],
    sample_weights: NDArray[np.float64]
) -> Tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]:
    """
    Computes the weighted covariance matrices R_tc, R_tt, and R_cc.

    Args:
        y_true_one_hot: 2D Array [n_samples, n_classes].
        y_pred_one_hot: 2D Array [n_samples, n_classes].
        sample_weights: 1D Array [n_samples].

    Returns:
        Tuple containing (R_tc, R_tt, R_cc).
    """
    t = np.asarray(y_true_one_hot)
    c = np.asarray(y_pred_one_hot)
    w = np.asarray(sample_weights)
    
    tr_S = np.sum(w)
    if tr_S == 0:
        K = t.shape[1]
        return np.zeros((K,K)), np.zeros((K,K)), np.zeros((K,K))
    
    # Compute means
    t_bar = compute_weighted_means(t, w)
    c_bar = compute_weighted_means(c, w)
    
    # Center data: (t_n - t_bar)
    # t is [N, K], t_bar is [K]. Broadcasting handles subtraction row-wise.
    t_centered = t - t_bar
    c_centered = c - c_bar
    
    # R = 1/tr(S) * (T_centered^T * S * C_centered)
    # Efficiently: T_centered^T * (w[:, None] * C_centered)
    
    # Pre-weight one of the matrices
    w_c_centered = c_centered * w[:, np.newaxis]
    w_t_centered = t_centered * w[:, np.newaxis]
    
    r_tc = np.dot(t_centered.T, w_c_centered) / tr_S
    r_tt = np.dot(t_centered.T, w_t_centered) / tr_S
    r_cc = np.dot(c_centered.T, w_c_centered) / tr_S
    
    return r_tc, r_tt, r_cc