import torch
import numpy as np

def U_proj(x, p, beta):
    k = -beta
    handle_obj = lambda a: np.linalg.norm(x - a, 2) ** 2
    Ux = np.zeros_like(x)
    m = np.linalg.norm(x[:p], 2) ** 2
    n = np.linalg.norm(x[p:], 2) ** 2
    a = 1 / 4 * k
    b = 0
    c = m - n - 2 * k
    d = -4 * m - 4 * n
    e = 4 * m - 4 * n + 4 * k
    coeffs = [a, b, c, d, e]
    roots = np.roots(coeffs)

    real_roots = roots[np.isreal(roots)].real
    fobj_list = np.zeros(real_roots.size)

    for i in range(real_roots.size):
        Ux_temp = np.zeros_like(x)
        Ux_temp[:p] = 2 * x[:p] / (2 + real_roots[i])
        Ux_temp[p:] = 2 * x[p:] / (2 - real_roots[i])
        fobj_list[i] = handle_obj(Ux_temp)

    index = np.argmin(fobj_list)
    lambda_val = real_roots[index]

    Ux[:p] = 2 * x[:p] / (2 + lambda_val)
    Ux[p:] = 2 * x[p:] / (2 - lambda_val)

    return torch.tensor(Ux)

