import numpy as np
from scipy.optimize import minimize


M = 5  
p = 10  


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


def loss_function(params, local_grads, global_grad):
    B_squared, H_squared = params[0], params[1]
    

    local_grad_square_sum = sum([np.linalg.norm(g)**2 for g in local_grads]) / M
    

    global_grad_square = np.linalg.norm(global_grad)**2
    right_side = B_squared * global_grad_square + H_squared
    

    loss = (local_grad_square_sum - right_side) ** 2
    return loss


initial_params = [1.0, 1.0]  

result = minimize(loss_function, initial_params, args=(local_gradients, global_gradient), method='L-BFGS-B', bounds=[(0, None), (0, None)])


B_squared_opt, H_squared_opt = result.x
B_opt = np.sqrt(B_squared_opt)
H_opt = np.sqrt(H_squared_opt)


print(f"Optimized B: {B_opt}, Optimized H: {H_opt}")
