import itertools as iterator
import numpy as np
from tqdm import tqdm
from concurrent.futures import ProcessPoolExecutor, as_completed


class GridSearch:
    def __init__(self, *, cls_regression, executor=None, max_parallel_workers=1):
        self.cls_regression = cls_regression
        self.executor = executor or ProcessPoolExecutor(
            max_workers=max_parallel_workers
        )

    @staticmethod
    def _train_and_test(
        model: object,
        n_classes: int,
        xtrain: np.ndarray,
        ytrain: np.ndarray,
        xtest: np.ndarray,
        ytest: np.ndarray,
        params: dict,
    ):
        regression = model(n_classes=n_classes, **params)
        regression.fit(xtrain=xtrain, ytrain=ytrain)
        return regression.precision(x=xtest, y=ytest), params

    def search(
        self,
        n_classes: int,
        xtrain: np.ndarray,
        ytrain: np.ndarray,
        xtest: np.ndarray,
        ytest: np.ndarray,
        params: dict,
    ):
        names = sorted(params.keys())
        combinations = iterator.product(*(params[param_name] for param_name in names))

        def __map_names_to_values(values):
            return {k: v for k, v in zip(names, values)}

        tasks = [
            self.executor.submit(
                self._train_and_test,
                model=self.cls_regression,
                n_classes=n_classes,
                xtrain=xtrain,
                ytrain=ytrain,
                xtest=xtest,
                ytest=ytest,
                params=__map_names_to_values(combination),
            )
            for combination in combinations
        ]

        max_precision = -float("inf")
        optimized_parameters = None
        for task in tqdm(as_completed(tasks), desc="Search by grid", total=len(tasks)):
            precision, parameters = task.result()
            if precision > max_precision:
                max_precision = precision
                optimized_parameters = parameters
        return max_precision, optimized_parameters
