import math

import cv2
import numpy as np
from scipy.ndimage import maximum_filter


def _neighbor_exists(x, y, min_distance, centroids):
    """
    Detects if there is previosly detected centroid within defined distance
    from currently detected centroid.

    Parameters
    ----------
    x : int
    y : int
    min_distance : float
    centroids : list
        previously detected centroids

    Returns
    -------
    exists : boolean
    """

    exists = False

    for centroid in centroids:
        distance = math.dist([x, y], centroid)
        if distance < min_distance:
            return True
    return exists


def centroid_counterpoint(localization_map, params):
    """
    Centroid counterpoint module of the centroid detector.

    Parameters
    ----------
    localization_map : numpy.ndarray
    params : dict
        dictionary of centroid counterpoint parameters: filter_size, min_distance, threshold,
        specified for all classes

    Returns
    -------
    centroid_map : numpy.ndarray
    """

    height = localization_map.shape[0]
    width = localization_map.shape[1]
    n_channels = len(localization_map.shape) > 2 and localization_map.shape[2] or 1

    # denormalize
    centroid_map = np.zeros((height, width, n_channels), np.uint8)

    # do for every class (every channel)
    for i in range(n_channels):

        localization_map_channel = localization_map[:, :, i]

        # set parameters according to class
        filter_size = params[f"class_{i}"]["filter_size"]
        min_distance = params[f"class_{i}"]["min_distance"]
        threshold = params[f"class_{i}"]["threshold"]

        # M
        M = np.zeros((height, width), np.uint8)
        M = maximum_filter(localization_map_channel, size=filter_size)

        # M1
        M1 = np.zeros((height, width), np.uint8)
        M1[np.equal(M, localization_map_channel)] = 1

        # OMEGA
        OMEGA = np.zeros((height, width), np.uint8)
        OMEGA[localization_map_channel == 0] = 1

        # OMEGA_ERODED
        kernel = np.ones((filter_size, filter_size), np.uint8)
        OMEGA_ERODED = cv2.erode(OMEGA, kernel, iterations=1)

        # M_OMEGA
        M_OMEGA = np.zeros((height, width), np.uint8)
        M_OMEGA[np.not_equal(M1, OMEGA_ERODED)] = 1

        # XI
        XI = np.zeros((height, width), np.uint8)
        valid_pixels = np.logical_and(
            M_OMEGA == 1, localization_map_channel >= threshold
        )
        valid_pixels_indexes = np.where(valid_pixels)
        valid_pixels_indexes = np.transpose(valid_pixels_indexes)

        # check for min_distance between two centroids
        centroids = []
        for x, y in valid_pixels_indexes:
            if not _neighbor_exists(x, y, min_distance, centroids):
                XI[x][y] = 1
            centroids.append([x, y])

        centroid_map[:, :, i] = XI

    return centroid_map
