__all__ = ['geomedian_eval']
import numpy as np
import numpy.linalg as la


def EvaluateFunction(x, component, B):
    """
    Evaluates the sum of squares-variant chwirut function
    """
    m = B.shape[0]
    assert B.shape[1] == len(x)
    assert 0 <= component <= m-1
    i = component
    b_i = B[i]

    f_i = 1.0/m * la.norm(x-b_i)
    return f_i


def EvaluateJacobian(x, component, B):
    """
    Evaluates the sum of squares-variant chwirut Jacobian
    """
    m = B.shape[0]
    assert B.shape[1] == len(x)
    assert 0 <= component <= m-1
    i = component
    b_i = B[i]

    df_i = 1.0/m * (x-b_i)/la.norm(x-b_i)
    return df_i


def geomedian_eval(H, persis_info, sim_specs, _):

    B = persis_info['params']['B']

    num_xs = len(H['x'])  # b==1 always?
    H_o = np.zeros(num_xs, dtype=sim_specs['out'])

    for k, x in enumerate(H['x']):
        i = H[k]['obj_component']   # f_i

        if H[k]['get_grad']:
            H_o['gradf_i'][k] = EvaluateJacobian(x, i, B)
        else:
            H_o['f_i'][k] = EvaluateFunction(x, i, B)

    return H_o, persis_info
