import torch


def pos_def(ws, alpha=0.001, eps=1e-20):
    """Diagonal modification.

    This method takes a complex Hermitian matrix represented by its upper
    triangular part and adds the value of its trace multiplied by alpha
    to the real part of its diagonal. The output will have the format:
    (*,2,C+P)

    Arguments
    ---------
    ws : tensor
        An input matrix. The tensor must have the following format:
        (*,2,C+P)

    alpha : float
        A coefficient to multiply the trace. The default value is 0.001.

    eps : float
        A small value to increase the real part of the diagonal. The
        default value is 1e-20.
    """

    # Extracting data
    D = ws.dim()
    P = ws.shape[D - 1]
    C = int(round(((1 + 8 * P) ** 0.5 - 1) / 2))

    # Finding the indices of the diagonal
    ids_triu = torch.triu_indices(C, C)
    ids_diag = torch.eq(ids_triu[0, :], ids_triu[1, :])

    # Computing the trace
    trace = torch.sum(ws[..., 0, ids_diag], D - 2)
    trace = trace.view(trace.shape + (1,))
    trace = trace.repeat((1,) * (D - 2) + (C,))

    # Adding the trace multiplied by alpha to the diagonal
    ws_pf = ws.clone()
    ws_pf[..., 0, ids_diag] += alpha * trace + eps

    return ws_pf