import numpy as np


def _relative_distance(width, height, point1, point2):
    """
    Function to calculate relative distance between two points.

    Parameters
    ----------
    width : int
        width of an image
    height : int
        height of an image
    point1 : numpy.ndarray
    point2 : numpy.ndarray

    Returns
    -------
    distance : float
        if points are equal distance is 0, while the maximum distance is 1
    """

    return min(width, abs(point1[0] - point2[0])) / (2 * width) + min(
        height, abs(point1[1] - point2[1])
    ) / (2 * height)


def localization_error(img_size, true, pred, n_classes, mode):
    """
    Function to calculate total or mean localization error for each class and
    combined. Total localization error is defined as mean localization error of
    one image. Mean localization error is defined as mean localization error of
    one object.

    Parameters
    ----------
    img_size : tuple
        (width, height) of an image
    true : list
        list of numpy numpy.ndarrays of centroid predictions
    pred : list
        list of numpy numpy.ndarrays of centroid predictions
    n_classes : int
    mode : str
        "total" / "mean"

    Returns
    -------
    error_total / error_mean : dict
        dictionary with keys as classes and values as errors
    """

    N = len(true)

    width = img_size[0]
    height = img_size[1]

    classes = list()
    no_per_class = dict()
    error_per_class = dict()

    for i in range(n_classes):
        classes.append(f"class_{i}")

    for c in classes:
        no_per_class[c] = 0
        error_per_class[c] = 0

    for (t, p) in zip(true, pred):
        for layer, c in enumerate(classes):
            true_class = t[np.where(t[:, -1] == layer)]
            pred_class = p[np.where(p[:, -1] == layer)]

            diff = abs(len(true_class) - len(pred_class))

            for _ in range(diff):
                if len(true_class) < len(pred_class):
                    true_class = np.concatenate((true_class, [[np.inf, np.inf, layer]]))
                else:
                    pred_class = np.concatenate((pred_class, [[np.inf, np.inf, layer]]))

            no_per_class[c] += len(true_class)

            for i in true_class:
                smallest_dist = np.inf
                for idx_j, j in enumerate(pred_class):
                    current_dist = _relative_distance(width, height, i, j)
                    if current_dist < smallest_dist:
                        smallest_dist = current_dist
                        idx_j_del = idx_j

                pred_class[idx_j_del] = [
                    np.iinfo(np.int64).max,
                    np.iinfo(np.int64).max,
                    layer,
                ]

                error_per_class[c] += smallest_dist

    if mode == "total":
        error_total = dict()
        for c in classes:
            error_total[c] = error_per_class[c] / N

        error_total["combined"] = sum(error_per_class.values()) / N
        return error_total

    elif mode == "mean":
        error_mean = dict()
        for c in classes:
            error_mean[c] = error_per_class[c] / no_per_class[c]
        error_mean["combined"] = sum(error_per_class.values()) / sum(
            no_per_class.values()
        )
        return error_mean


def count_error(true, pred, n_classes):
    """
    Function to calculate count error for each class and combined.

    Parameters
    ----------
    true : numpy.ndarray
        numpy array of true centroids (labels)
    pred : numpy.ndarray
        numpy array of detected centroids (predictions)
    n_classes : int

    Returns
    -------
    count_err : dict
    """

    classes = list()
    count_err = dict()

    for i in range(n_classes):
        classes.append(f"class_{i}")
        count_err[classes[i]] = dict()

    count_err["combined"] = dict()

    for (t, p) in zip(true, pred):
        diff_temp = 0
        for layer, c in enumerate(classes):
            true_class = t[np.where(t[:, -1] == layer)]
            pred_class = p[np.where(p[:, -1] == layer)]

            diff = abs(len(true_class) - len(pred_class))
            diff_temp += diff

            if diff in count_err[c]:
                count_err[c][diff] += 1
            else:
                count_err[c][diff] = 1

        if diff_temp in count_err["combined"]:
            count_err["combined"][diff_temp] += 1
        else:
            count_err["combined"][diff_temp] = 1

    # sort dictionaries on keys for printing
    for key, value in count_err.items():
        count_err[key] = dict(sorted(value.items(), key=lambda x: x[0]))

    return count_err
