import numpy as np


def gradient_square_sum(gradients, M):
    return sum([np.linalg.norm(g)**2 for g in gradients]) / M

def global_gradient_square(global_grad):
    return np.linalg.norm(global_grad)**2

def calculate_H2(local_grads, global_grad, B_squared, M):
    H_squared = 0
    global_grad_norm_square = np.linalg.norm(global_grad)**2


    if global_grad_norm_square < 1e-10:
        print("Warning: Global gradient norm squared is too small, setting H^2 to 0")
        return 0

    for grad in local_grads:
        grad_norm_square = np.linalg.norm(grad)**2
        

        diff = grad_norm_square - B_squared * global_grad_norm_square
        
        if np.abs(diff) < 1e-10:  
            diff = 0

        H_squared += diff

    H_squared /= M

    if not np.isfinite(H_squared) or H_squared < 0:
        print("Warning: H_squared is non-finite (nan or inf) or negative, setting H_squared to 0")
        H_squared = 0  
    
    return H_squared

def calculate_B_and_H(local_grads, global_grad, M):

    local_grad_square_sum = gradient_square_sum(local_grads, M)


    global_grad_square_sum = global_gradient_square(global_grad)

    if global_grad_square_sum == 0:
        print("Warning: Global gradient square sum is zero, returning B as 0")
        return 0, 0  

    B_squared = local_grad_square_sum / global_grad_square_sum


    H_squared = calculate_H2(local_grads, global_grad, B_squared, M)


    return np.sqrt(B_squared), np.sqrt(H_squared)


local_gradients = [np.random.rand(10) for _ in range(5)]  
global_gradient = np.mean(local_gradients, axis=0)  

M = len(local_gradients)
B, H = calculate_B_and_H(local_gradients, global_gradient, M)

print(f"B: {B}, H: {H}")
