import numpy as np
from numpy.typing import NDArray
from .binary import weighted_mcc
from .multiclass import extended_corr_coef, mpc_trace_ratio, mpc_mean_corr
from .utils import compute_covariance_matrices, compute_weighted_means

def calculate_binary_stability_bound(
    y_true: NDArray[np.int_],
    y_pred: NDArray[np.int_],
    sample_weights: NDArray[np.float64],
    epsilon: float
) -> float:
    """
    Calculates the theoretical upper bound for the change in Weighted MCC 
    given a perturbation epsilon to the weights.

    Args:
        y_true: 1D Array [n_samples] of binary truth labels.
        y_pred: 1D Array [n_samples] of binary prediction labels.
        sample_weights: 1D Array [n_samples].
        epsilon: Scalar float representing max change in weights (|S_ii - W_ii|).
                 Must be < min(m/2, 1/N).

    Returns:
        float: The maximum expected change in MCC value, or None if preconditions
               (m > 0 and epsilon check) are not met.
    """
    t = np.asarray(y_true)
    c = np.asarray(y_pred)
    w = np.asarray(sample_weights)
    
    # Calculate inner products used in the stability bound (The margins m and M)
    # These represent the weighted counts of TP, Truth Pos, Pred Pos, etc.
    ip_tc = np.dot(w, t * c)           # Weighted True Positives
    ip_t1 = np.dot(w, t)               # Weighted Total Positives (Truth)
    ip_1c = np.dot(w, c)               # Weighted Total Predicted Positives
    ip_1mt_1 = np.dot(w, 1 - t)        # Weighted Total Negatives (Truth)
    ip_1_1mc = np.dot(w, 1 - c)        # Weighted Total Predicted Negatives
    
    A = np.array([ip_tc, ip_t1, ip_1c, ip_1mt_1, ip_1_1mc])
    M = np.max(A)
    m = np.min(A)
    N = len(t)
    
    # Return None for precondition violations as expected by property tests
    if m <= 0 or epsilon >= min(m/2, 1/N):
        return None
    
    mcc_val = weighted_mcc(y_true, y_pred, sample_weights)
    tr_s = np.sum(w)
    
    # Stability factor logic from Theorem 3.4
    k1 = (2**5 * tr_s**2 * N) / (m**2)
    inner_term = 1 + M**2 * (1 + k1 * epsilon)
    
    bound = abs(mcc_val) * k1 * inner_term * epsilon
    return float(bound)

def calculate_multiclass_stability_bound(
    y_true_one_hot: NDArray[np.int_],
    y_pred_one_hot: NDArray[np.int_],
    sample_weights: NDArray[np.float64],
    epsilon: float,
    metric_type: str = 'ECC'
) -> float:
    """
    Calculates the theoretical upper bound for the change in multiclass metrics 
    (ECC, MPC1, MPC2) given a perturbation epsilon to the weights.
    """
    r_tc, r_tt, r_cc = compute_covariance_matrices(y_true_one_hot, y_pred_one_hot, sample_weights)
    s = np.sum(sample_weights)
    
    if s == 0:
        return 0.0
    if epsilon >= s/2:
        raise ValueError(f"Epsilon {epsilon} must be < s/2 = {s/2}")
    
    N = y_true_one_hot.shape[0]
    
    t_bar = compute_weighted_means(y_true_one_hot, sample_weights)
    c_bar = compute_weighted_means(y_pred_one_hot, sample_weights)
    
    t_centered = y_true_one_hot - t_bar
    c_centered = y_pred_one_hot - c_bar
    
    M_t = np.max(np.abs(t_centered))
    M_c = np.max(np.abs(c_centered))
    
    diag_tt = np.diagonal(r_tt)
    diag_cc = np.diagonal(r_cc)
    C_k = diag_tt * (M_c**2) + diag_cc * (M_t**2) + (4 * N / s**2)
    
    if metric_type == 'ECC':
        val = extended_corr_coef(y_true_one_hot, y_pred_one_hot, sample_weights)
        y = np.trace(r_tt) * np.trace(r_cc)
        if y == 0: return 0.0
        
        C_global = (M_t**2) * np.trace(r_cc) + (M_c**2) * np.trace(r_tt) + (4 * N**2 / s**2)
        
        term1 = (16 * N**2 * C_global) / (y * np.sqrt(2) * s**2)
        term2 = (4 * N**2 * M_t * M_c) / (s**2)
        term3 = (16 * N * C_global) / (y * np.sqrt(2) * s**2)
        
        factor = term1 + term2 * (1 + term3 * epsilon**2)
        return abs(val) * factor * epsilon**2
        
    elif metric_type == 'MPC1':
        val = mpc_trace_ratio(y_true_one_hot, y_pred_one_hot, sample_weights)
        sqrt_diags = np.sqrt(diag_tt * diag_cc)
        y = np.sum(sqrt_diags)
        if y == 0: return 0.0
        
        denom_k = np.sqrt(2 * diag_tt * diag_cc)
        mask = denom_k > 0
        C_global = np.sum(C_k[mask] / denom_k[mask])
        
        term1 = (16 * N * C_global) / (y * s**2)
        term2 = (4 * N**2 * M_t * M_c) / (s**2)
        term3 = (16 * N * C_global) / (y * s**2)
        
        factor = term1 + term2 * (1 + term3 * epsilon**2)
        return abs(val) * factor * epsilon**2
        
    elif metric_type == 'MPC2':
        C_max = np.max(C_k)
        diag_tc = np.diagonal(r_tc)
        sqrt_diags = np.sqrt(diag_tt * diag_cc)
        K = len(diag_tc)
        
        mask = sqrt_diags > 1e-15
        ratios = np.zeros_like(diag_tc)
        ratios[mask] = np.abs(diag_tc[mask]) / (K * sqrt_diags[mask])
        sum_ratios = np.sum(ratios)
        
        term1 = (16 * N * C_max) / (s**2)
        term2 = (4 * N * C_max) / (s**2)
        term3 = (16 * N * C_max) / (s**2)
        
        factor = term1 + term2 * (1 + term3 * epsilon**2)
        return sum_ratios * factor * epsilon**2
    else:
        raise ValueError("Invalid metric_type. Choose 'ECC', 'MPC1', or 'MPC2'.")