import numpy as np
from .common import Benchmark, safe_import

with safe_import():
    from scipy.optimize import linear_sum_assignment
    from scipy.spatial.distance import cdist


def random_uniform(shape):
    return np.random.uniform(-20, 20, shape)


def random_logarithmic(shape):
    return 10**np.random.uniform(-20, 20, shape)


def random_integer(shape):
    return np.random.randint(-1000, 1000, shape)


def random_binary(shape):
    return np.random.randint(0, 2, shape)


def random_spatial(shape):
    P = np.random.uniform(-1, 1, size=(shape[0], 2))
    Q = np.random.uniform(-1, 1, size=(shape[1], 2))
    return cdist(P, Q, 'sqeuclidean')


class LinearAssignment(Benchmark):

    sizes = range(100, 401, 100)
    shapes = [(i, i) for i in sizes]
    shapes.extend([(i, 2 * i) for i in sizes])
    shapes.extend([(2 * i, i) for i in sizes])
    cost_types = ['uniform', 'spatial', 'logarithmic', 'integer', 'binary']
    param_names = ['shape', 'cost_type']
    params = [shapes, cost_types]

    def setup(self, shape, cost_type):

        cost_func = {'uniform': random_uniform,
                     'spatial': random_spatial,
                     'logarithmic': random_logarithmic,
                     'integer': random_integer,
                     'binary': random_binary}[cost_type]

        self.cost_matrix = cost_func(shape)

    def time_evaluation(self, *args):
        linear_sum_assignment(self.cost_matrix)