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

def extended_corr_coef(
    y_true_one_hot: NDArray[np.int_],
    y_pred_one_hot: NDArray[np.int_],
    sample_weights: NDArray[np.float64]
) -> float:
    """
    Calculates Extended Correlation Coefficient (ECC).
    """
    r_tc, r_tt, r_cc = compute_covariance_matrices(y_true_one_hot, y_pred_one_hot, sample_weights)
    
    numerator = np.trace(r_tc)
    denominator = np.sqrt(np.trace(r_tt) * np.trace(r_cc))
    
    if denominator == 0:
        return 0.0
        
    return float(numerator / denominator)

def mpc_trace_ratio(
    y_true_one_hot: NDArray[np.int_],
    y_pred_one_hot: NDArray[np.int_],
    sample_weights: NDArray[np.float64]
) -> float:
    """
    Calculates MPC1 (Trace Ratio).
    """
    r_tc, r_tt, r_cc = compute_covariance_matrices(y_true_one_hot, y_pred_one_hot, sample_weights)
    
    diag_tc = np.diagonal(r_tc)
    diag_tt = np.diagonal(r_tt)
    diag_cc = np.diagonal(r_cc)
    
    numerator = np.sum(diag_tc)
    denominator = np.sum(np.sqrt(diag_tt * diag_cc))
    
    if denominator == 0:
        return 0.0
        
    return float(numerator / denominator)

def mpc_mean_corr(
    y_true_one_hot: NDArray[np.int_],
    y_pred_one_hot: NDArray[np.int_],
    sample_weights: NDArray[np.float64]
) -> float:
    """
    Calculates MPC2 (Mean Correlation) with continuity handling for zero-variance classes.
    """
    y_t = np.asarray(y_true_one_hot)
    y_p = np.asarray(y_pred_one_hot)
    
    r_tc, r_tt, r_cc = compute_covariance_matrices(y_t, y_p, sample_weights)
    
    diag_tc = np.diagonal(r_tc)
    diag_tt = np.diagonal(r_tt)
    diag_cc = np.diagonal(r_cc)
    
    K = len(diag_tc)
    if K == 0:
        return 0.0
        
    denom_terms = np.sqrt(diag_tt * diag_cc)
    mask = denom_terms > 1e-15
    
    # Continuity logic to satisfy Remark 1 and zero-variance test cases:
    # If a class matches perfectly and is 'present', correlation is 1.0.
    # Otherwise (e.g., absent or mismatch on constant), it is 0.0.
    matches = np.all(y_t == y_p, axis=0)
    presence = np.sum(y_t, axis=0) > 0
    continuity_ones = (~mask) & matches & presence
    
    terms = np.zeros_like(diag_tc, dtype=np.float64)
    terms[mask] = diag_tc[mask] / denom_terms[mask]
    terms[continuity_ones] = 1.0
    
    return float(np.sum(terms) / K)