import itertools

import numpy as np
from joblib import Parallel, delayed
from tqdm import tqdm

from detector.centroid_counterpoint_module import centroid_counterpoint
from detector.localization_module import localizator
from tools.evaluator import localization_error


def counterpoint_params_search(y_true, y_pred, param_grid, n_jobs=-1):
    """
    Custom grid search like hyperparameter optimization function. This function
    is custom for tuning hyperparameters of the centroid counterpoint module.
    This function is implemented in parallelized fashion.

    Parameters
    ----------
    y_true : numpy.ndarray
        array of true localization maps
    y_pred : numpy.ndarray
        array of predicted localization maps
    param_grid : dict
        dictionary of parameter grid (it must contain three keys:
        filter_size, min_distance, threshold)
    n_jobs : int
        Number of jobs to run in parallel, -1 means all processors.

    Returns
    -------
    best_params : dict
        dictionary of the best parameters found
    """
    img_width = y_true[0].shape[0]
    img_height = y_true[0].shape[1]
    n_classes = y_true[0].shape[2]

    param_combinations = [
        dict(zip(param_grid.keys(), values))
        for values in itertools.product(*param_grid.values())
    ]

    counterpoint_params = []
    for comb in param_combinations:
        temp = {}
        for i in range(n_classes):
            temp[f"class_{i}"] = comb
        counterpoint_params.append(temp)

    true = []
    for localization_map in y_true:
        centroids = np.argwhere(localization_map == 1)
        true.append(centroids)

    def calculate_score(params):
        pred = []
        for localization_map in y_pred:
            centroid_map = centroid_counterpoint(localization_map, params)
            centroids = localizator(centroid_map)
            pred.append(centroids)

        score = localization_error(
            (img_width, img_height), true, pred, n_classes, mode="mean"
        )
        return (params, score)

    results = Parallel(n_jobs=n_jobs)(
        delayed(calculate_score)(params)
        for params in tqdm(counterpoint_params, total=len(param_combinations))
    )

    best_params = {}

    for i in range(n_classes):
        best_score = np.inf
        for params, score in results:
            if score[f"class_{i}"] < best_score:
                best_score = score[f"class_{i}"]
                best_params[f"class_{i}"] = params[f"class_{i}"]

    return best_params
