import numpy as np

n_dof = 3
n_gains = 1

def number_of_gains(n_dof, n_gains, sym=True, diag=False):
    
    if diag is True:
        mat = np.eye(n_dof) * n_gains
    else:
        if sym is True:
            mat = np.triu(n_gains * np.ones((n_dof,n_dof)))
        elif sym is False:
            mat = np.ones((n_dof,n_dof)) * n_gains
        
    print(mat)
    return int(np.sum(mat))


def make_gain_matrix(n_dof, gains=None, sym=True, diag=False):
    
    n_tot = number_of_gains(n_dof=n_dof,n_gains=1,sym=sym,diag=diag)

    if gains is None:
        gains = np.arange(n_tot)+1
    else:
        if len(gains) != n_tot:
            raise ValueError(f"Length of gains is {len(gains)}, should be {n_tot}")

    if diag is True:
        whole_sym = np.diag(gains)
    else:
        if sym is True:
            inds = np.triu_indices(n_dof)
            gain_mat = np.zeros((n_dof,n_dof))
            gain_mat[inds] = gains
            
            diag = np.diag(np.diag(gain_mat))
            upper_tri = np.triu(gain_mat,1)
            lower_tri = np.tril(np.transpose(gain_mat),-1)
            
            whole_sym = diag + upper_tri + lower_tri
        elif sym is False:
            whole_sym = np.reshape(gains,(n_dof,n_dof))
    
    return whole_sym

make_gain_matrix(n_dof=n_dof, 
                #  gains=[1,2,3,4,5,6],
                 sym=True,
                 diag=False)