from abc import abstractmethod
from typing import List

import numpy as np
import torch
from pymoo.algorithms.soo.nonconvex.ga import GA
from pymoo.core.algorithm import Algorithm
from pymoo.core.problem import ElementwiseProblem
from pymoo.factory import get_crossover, get_mutation, get_sampling
from pymoo.optimize import minimize
from rt_search_based.datasets.datasets import DatasetItem
from rt_search_based.fitness_functions.fitness_functions import FitnessFunction
from rt_search_based.models.classifiers import Classifier
from rt_search_based.strategies.strategies import Strategy
from rt_search_based.transformations import stickers
from rt_search_based.transformations.stickers import Color


class PyMooSingleObjectiveProblem(ElementwiseProblem):
    """
    pymoo problem definition. This is the core of the pymoo optimization.
    The problem is defined by the following attributes:
    - n_var: number of variables in the search space
    - n_obj: number of objectives
    - n_constr: number of constraints
    - xl: lower bound of the variables
    - xu: upper bound of the variables
    - type_var: type of the variables
    """

    def __init__(
        self,
        classifier: Classifier,
        fitness_function: FitnessFunction,
        dataset_item: DatasetItem,
        sticker_colors: List[Color],
        sticker_count: int = 1,
        use_constraints: bool = False,
    ):
        self.classifier = classifier
        self.fitness_function = fitness_function
        self.dataset_item = dataset_item
        self.sticker_colors = sticker_colors
        self.use_constraints = use_constraints

        _, self.img_height, self.img_width = self.dataset_item.image.size()

        super().__init__(
            n_var=5 * sticker_count,
            n_obj=1,
            n_constr=(2 * sticker_count + 1) * int(self.use_constraints),
            xl=np.zeros(sticker_count * 5),
            xu=np.array(
                [
                    self.img_width,
                    self.img_width,
                    self.img_height,
                    self.img_height,
                    len(self.sticker_colors) - 1,
                ]
                * sticker_count
            ),
            type_var=int,
        )

    def _evaluate(self, x, out, *args, **kwargs):
        properties = stickers.create_multi_sticker_properties(x, self.sticker_colors)
        img = stickers.add_multi_sticker_to_image(properties, self.dataset_item.image)

        # set fitness value as minimization target
        out["F"] = self.fitness_function(properties, self.dataset_item)
        # if specified also make use of constraints
        if self.use_constraints:
            # add the constraints for valid rectangles to the pymoo constraints list
            out["G"] = []
            for i in range(1, len(properties), 7):
                x1, y1, x2, y2, _, _, _ = properties[i : i + 7]
                out["G"].extend([x1 - x2, y1 - y2])
            # if specified also make the fooling of the model a constraint
            out["G"].append(self.passes_test(img))

    def passes_test(self, img):
        """
        Function to check if the model is fooled by the given stickered image.
        Used as a constraint in the pymoo optimization to only consider images
        that fool the model.
        """

        # Add extra dimension used for classifier (equivalent to batch_size 1)
        img = torch.unsqueeze(img, dim=0)
        predicted_class = self.classifier.get_predicted_class(img)

        correct_prediction = predicted_class == self.dataset_item.label
        return int(correct_prediction)


class PyMooStrategy(Strategy):
    def __init__(
        self,
        fitness_function: FitnessFunction,
        classifier: Classifier,
        sticker_colors: List[Color] = None,
        sticker_count: int = 1,
        need_constraints=False,
        **kwargs,
    ):

        if self.__class__.__name__ == "PyMooStrategy":
            raise TypeError

        super().__init__(
            fitness_function, classifier, sticker_colors, sticker_count, **kwargs
        )

        self.algorithm = self.get_new_algorithm()
        self.need_constraints = need_constraints

    @abstractmethod
    def get_new_algorithm(self) -> Algorithm:
        """return the algorithm that should be used"""

    def search_for_sticker(self, dataset_item: DatasetItem) -> torch.Tensor:
        """
        Function that searches for an optimal sticker in the given image.
        It creates a pymoo problem and optimizes it using the chosen pymoo algorithm,
        and then it returns the optimal sticker properties (solution).
        """

        problem = PyMooSingleObjectiveProblem(
            self.classifier,
            self.fitness_function,
            dataset_item,
            self.sticker_colors,
            self.sticker_count,
            self.need_constraints,
        )

        result = minimize(
            problem, self.algorithm, seed=1, save_history=True, verbose=True
        )
        solution = stickers.create_multi_sticker_properties(
            result.X, self.sticker_colors
        )

        # the code below checks if the solution fools the classifier
        img = stickers.add_multi_sticker_to_image(solution, dataset_item.image)

        # Add extra dimension used for classifier (equivalent to batch_size 1)
        img = torch.unsqueeze(img, dim=0)
        if self.classifier.get_predicted_class(img) == dataset_item.label:
            return torch.zeros_like(solution) - 1

        return solution


class PyMooGeneticStrategy(PyMooStrategy):
    def get_new_algorithm(self) -> Algorithm:
        return GA(
            sampling=get_sampling("int_random"),
            crossover=get_crossover("int_sbx"),
            mutation=get_mutation("int_pm"),
            # pop_size=100,
            # n_offsprings=2,
        )