import numpy as np

from artemis.general.mymath import cosine_distance, softmax


def get_evaluation_function(name):
    return {
        'mse': mean_squared_error,
        'mean_cosine_distance': lambda a, b: cosine_distance(a, b, axis=1).mean(),
        'mean_squared_error': mean_squared_error,
        'mean_l1_error': mean_l1_error,
        'percent_argmax_correct': percent_argmax_correct,
        'percent_argmax_incorrect': percent_argmax_incorrect,
        'percent_correct': percent_correct,
        'softmax_categorical_xe': softmax_categorical_xe
        }[name]


def mean_l1_error(actual, target):
    return np.mean(np.sum(np.abs(actual-target), axis=-1), axis=-1)


def mean_squared_error(actual, target):
    return np.mean(np.sum((actual-target)**2, axis = -1), axis = -1)


def softmax_categorical_xe(actual, target):
    """
    :param actual: An (n_samples, n_dims) array identifying pre-logistic output
    :param target: An (n_samples, ) integer array identifying labels
    :return:
    """
    if target.ndim==1:
        assert target.dtype==int and np.max(target) < actual.shape[1]
    elif target.ndim==2:
        assert np.all(target.sum(axis=1)==1) and np.all(np.max(target, axis=1)==1)
        target = np.argmax(target, axis=1)
    else:
        raise Exception("Don't know how to interpret a {}-D target".format(target))

    return np.mean(softmax(actual, axis=1)[np.arange(actual.shape[0]), target], axis=0)


def fraction_correct(actual, target):
    return np.mean(actual == target)


def percent_correct(actual, target):
    return 100*fraction_correct(actual, target)


def percent_argmax_correct(actual, target):
    """
    :param actual: An (n_samples, n_dims) array
    :param target: An (n_samples, ) array of indices OR an (n_samples, n_dims) array
    :return:
    """
    actual = collapse_onehot_if_necessary(actual)
    target = collapse_onehot_if_necessary(target)
    return 100*fraction_correct(actual, target)


def percent_binary_incorrect(actual, target):
    return 100.-percent_binary_correct(actual, target)


def percent_binary_correct(actual, target):
    """
    :param actual:  A (n_samples, ) array of floats between 0 and 1
    :param target: A (n_samples, ) array of True/False
    :return: The percent of times the "actual" was closes to the correct.
    """
    assert len(actual) == len(target)
    if actual.ndim==2:
        assert actual.shape[1]==1
        actual = actual[:, 0]
    if target.ndim==2:
        assert target.shape[1]==1
        target = target[:, 0]
    assert target.ndim==1
    if actual.ndim>1:
        assert actual.shape == (len(target), 1)
        actual = actual[:, 0]
    if np.array_equal(np.unique(target), (0, 1)):
        assert np.all(actual)>=0 and np.all(actual)<=1
        assert np.all((target==0)|(target==1))
        return 100*np.mean(np.round(actual) == target)
    elif np.array_equal(np.unique(target), (-1, 1)):
        assert np.all((target==-1)|(target==1))
        return 100*np.mean((actual>0)*2-1 == target)
    else:
        raise Exception("Go away I'm tired.")


def percent_argmax_incorrect(actual, target):
    return 100 - percent_argmax_correct(actual, target)


def collapse_onehot_if_necessary(output_data):
    """
    Given an input that could either be in onehot encoding or not, return it in onehot encoding.

    :param output_data: Either an (n_samples, n_dims) array, or an (n_samples, ) array of labels.
    :return: An (n_samples, ) array.
    """

    # output_data = np.squeeze(output_data)

    if output_data.ndim == 2:
        return np.argmax(output_data, axis = 1)
    else:
        assert output_data.ndim == 1 and output_data.dtype in (int, 'int32', bool)
        return output_data