import logging

from cobra.exceptions import OptimizationError

from cameobrs.core.manipulation import swap_cofactors
from cameobrs.strain_design.heuristic.evolutionary.decoders import SetDecoder
from cameobrs.strain_design.heuristic.evolutionary.objective_functions import (
    ObjectiveFunction,
)
from cameobrs.util import ProblemCache, memoize

logger = logging.getLogger(__name__)

__all__ = ["KnockoutEvaluator", "SwapEvaluator"]


class Evaluator(object):
    """
    Any evaluator that takes a population and returns a fitness.
    """

    def __call__(self, population):
        raise NotImplementedError

    def evaluate_individual(self, individual):
        raise NotImplementedError

    def reset(self):
        raise NotImplementedError


class TargetEvaluator(Evaluator):
    """
    Evaluator for targets in a model. It gets a representation, decodes into targets and simulates
    the model with a given method.

    Attributes
    ----------
    model : cobra.Model
        A constraint-based model
    decoder : SetDecoder
        A decoder to convert the representation into knockouts
    objective_function : objective_function or list(objective_function)
        The objectives of the algorithm
    simulation_method : see flux_analysis.simulation
        The method use to simulate the knockouts
    simulation_kwargs : dict
        The extra parameters used by the simulation method

    See Also
    --------
    cameo.strain_design.heuristic.objective_function

    Methods
    -------
    __call__(population)
        calling the object will evaluate a population (see inspyred)

    """

    def __init__(
        self, model, decoder, objective_function, simulation_method, simulation_kwargs
    ):
        self.model = model
        if not isinstance(decoder, SetDecoder):
            raise ValueError("Invalid decoder %s" % decoder)
        self.decoder = decoder

        if not isinstance(objective_function, ObjectiveFunction):
            raise ValueError(
                "'objective_function' must be instance of ObjectiveFunction (%s)"
                % objective_function
            )

        self.objective_function = objective_function
        self.simulation_method = simulation_method
        self.simulation_kwargs = simulation_kwargs
        self.cache = ProblemCache(model)

    def __call__(self, population):
        return [self.evaluate_individual(tuple(i)) for i in population]

    def reset(self):
        self.cache.reset()


class KnockoutEvaluator(TargetEvaluator):
    """
    Knockout evaluator for genes or reactions.
    """

    @memoize
    def evaluate_individual(self, individual):
        """
        Evaluates a single individual.

        Parameters
        ----------

        individual: set
            The encoded representation of a single individual.


        Returns
        -------
        fitness
            A single real value or a Pareto, depending on the number of objectives.
        """
        targets = self.decoder(individual)[0]
        with self.model:
            for target in targets:
                target.knock_out()
            try:
                solution = self.simulation_method(
                    self.model,
                    cache=self.cache,
                    volatile=False,
                    raw=True,
                    reactions=self.objective_function.reactions,
                    **self.simulation_kwargs
                )
                fitness = self.objective_function(self.model, solution, targets)
            except OptimizationError as e:
                logger.debug(e)
                fitness = self.objective_function.worst_fitness()
            return fitness


class SwapEvaluator(TargetEvaluator):
    """Evaluate reaction swaps where we knock one reaction in favor of another"""

    def __init__(self, swap_pair=None, *args, **kwargs):
        super(SwapEvaluator, self).__init__(*args, **kwargs)
        self.swap_pair = swap_pair

    @memoize
    def evaluate_individual(self, individual):
        swap_reactions = self.decoder(individual)[0]
        with self.model:
            for reaction in swap_reactions:
                swap_cofactors(reaction, self.model, self.swap_pair, inplace=True)
            try:
                solution = self.simulation_method(
                    self.model,
                    cache=self.cache,
                    volatile=False,
                    raw=True,
                    reactions=self.objective_function.reactions,
                    **self.simulation_kwargs
                )
                fitness = self.objective_function(self.model, solution, swap_reactions)
            except OptimizationError as e:
                logger.debug(e)
                fitness = self.objective_function.worst_fitness()

            return fitness


class KnockinKnockoutEvaluator(KnockoutEvaluator):
    def __init__(self, database, *args, **kwargs):
        super(KnockinKnockoutEvaluator, self).__init__(*args, **kwargs)
        self._database = database
