import os

import cv2
import matplotlib.pyplot as plt
import mplcursors
import numpy as np
from centroid_counterpoint_module import centroid_counterpoint
from localization_map_generator import map_generator
from localization_module import localizator

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"


def centroid_detector(image, model_path, params):
    """
    Detection of centroids of objects in an image.

    Parameters
    ----------
    image : numpy.ndarray
        colored or grayscale image represented as numpy.ndarray
    model_path : str
    params : dict
        dictionary of centroid counterpoint parameters: filter_size, min_distance, threshold,
        specified for all classes

    Returns
    -------
    centroids : numpy.ndarray
        numpy.ndarray of all centroids detected in an image,
        each element of the array is an array of three values (x, y, class)
    """

    localization_map = map_generator(image, model_path=model_path)
    centroid_map = centroid_counterpoint(localization_map, params)
    centroids = localizator(centroid_map)

    return centroids


def print_results(centroids):
    """
    Prints the number and positions of detected objects for each class.

    Parameters
    ----------
    centroids : numpy.ndarray
        A numpy array of all centroids detected in an image, where each element of the array
        is an array of three values (x, y, class).
    """

    classes = np.unique(centroids[:, 2])
    for c in classes:
        class_objs = centroids[centroids[:, 2] == c]
        print(f"\nNumber of class_{c} objects: {len(class_objs)}\n")
        for obj in class_objs:
            print(f"\tx: {obj[0]}, y: {obj[1]}")

    return None


def draw_results(image, centroids):
    """
    Visualisation of detected centroids on the original image.

    Parameters
    ----------
    image : numpy.ndarray
        colored or grayscale image represented as numpy.ndarray
    centroids : numpy.ndarray
        numpy.ndarray of all centroids detected in an image,
        each element of the array is an array of three values (x, y, class)
    """

    if image.ndim == 2:
        colored = False
    else:
        colored = True

    _, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) if colored else ax.imshow(
        image, cmap="gray"
    )

    color_bank = ["red", "lime", "blue", "black", "yellow", "purple"]
    scatter_plots = []
    for centroid in centroids:
        y = centroid[0]
        x = centroid[1]
        c = centroid[2]

        color = color_bank[c]

        scatter_plot = ax.scatter(x, y, s=100, c=color, marker="+")
        scatter_plots.append(scatter_plot)

        class_objs = centroids[centroids[:, 2] == c]
        num_objs = len(class_objs)
        class_str = f"class_{c}: {num_objs}"

        ax.text(
            0.90,
            0.95 - 0.04 * c,
            class_str,
            transform=ax.transAxes,
            ha="center",
            va="center",
            bbox=dict(fc="silver", ec=color, alpha=1),
        )

    ax.set_xlabel("y")
    ax.set_ylabel("x")

    ax.xaxis.set_label_position("top")
    ax.xaxis.set_tick_params(labeltop=True)
    ax.xaxis.set_tick_params(labelbottom=False)

    cursor = mplcursors.cursor(scatter_plots, hover=True)

    @cursor.connect("add")
    def on_add(sel):
        x, y = sel.target
        sel.annotation.set_text(f"x: {y:.0f}\ny: {x:.0f}")
        sel.annotation.get_bbox_patch().set(fc="white", ec="none", alpha=0.7)

    plt.show()

    return None


if __name__ == "__main__":

    # set image and model paths
    image_path = ""
    model_path = ""

    # set counterpoint params
    counterpoint_params = {}

    # read the image
    image = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)

    # centroid detector
    centroids = centroid_detector(image, model_path, counterpoint_params)

    # print results
    print_results(centroids)

    # draw results into the original image
    draw_results(image, centroids)
