from ambiegen.mutation.abstract_mutation import AbstractMutation
from ambiegen.generators.abstract_generator import AbstractGenerator
import numpy as np
import random
import logging

log = logging.getLogger(__name__)


class ObstacleMutation(AbstractMutation):
    """
    ObstacleMutation applies mutation operations to obstacle-based test configurations.
    This class provides mutation strategies for modifying, adding, or removing obstacles in a test configuration.
    It is designed to work with a generator that supplies obstacle parameters and normalization/denormalization
    methods for the test representation.

    Arguments:
        mut_prob (float, optional): Probability of mutation being applied. Defaults to 0.4.
        generator (AbstractGenerator, optional): Generator instance for obstacle parameter generation and test normalization.

    Methods:
        do_mutation(x):
            Applies a random mutation (modification, addition, or removal of an obstacle) to the input solution `x`.
            Returns the mutated solution if valid, otherwise returns the original solution.
        _random_modification(test):
            Randomly modifies three parameters of a randomly selected obstacle in the test configuration.
        _add_obstacle(test):
            Adds a new randomly generated obstacle to the test configuration, if the maximum number is not exceeded.
        _remove_obstacle(test):
            Removes a randomly selected obstacle from the test configuration, if more than one obstacle exists.
    """

    def __init__(self, mut_prob: float = 0.4, generator: AbstractGenerator = None):
        self.generator = (
            generator  # this allows to reuse mutation with differet frameworks
        )
        super().__init__(mut_prob)

    def do_mutation(self, x) -> np.ndarray:
        """
        Applies a random mutation to the given solution `x` and returns the mutated solution if it is valid.
        The mutation is selected randomly from the following options:
            - Random modification of the solution.
            - Addition of an obstacle.
            - Removal of an obstacle.
        If the mutated solution is not valid according to the validator, the original solution is returned.
        Args:
            x: The solution to be mutated.
        Returns:
            np.ndarray: The mutated solution if valid, otherwise the original solution.
        """
        if not self.generator:
            raise ValueError("Generator not set.")

        possible_mutations = [
            self._random_modification,
            self._add_obstacle,
            self._remove_obstacle,
        ]

        x = self.generator.denormalize_flattened_test(x)

        mutator = np.random.choice(possible_mutations)

        mutated_x = mutator(x)
        mutated_x = self.generator.normalize_flattened_test(mutated_x)

        return mutated_x

    def _random_modification(self, test):
        """
        Randomly modifies parameters of a single obstacle in the given test configuration.
        Selects one obstacle at random, then randomly chooses three of its six parameters to replace
        with new values generated by the obstacle generator. The test configuration is modified in place.
        Args:
            test (list): A list representing the test configuration, where the first element is the number
                of obstacles, followed by groups of six values per obstacle.
        Returns:
            list: The modified test configuration with updated obstacle parameters.
        """

        num_obstacles = test[0]

        idx = random.randint(0, num_obstacles - 1)
        param_idx = random.sample(range(0, 6), 3)
        start = 1 + idx * 6

        new_obstacle = self.generator.get_random_box_vals()

        for p in param_idx:
            test[start + p] = new_obstacle[p]

        return test

    def _add_obstacle(self, test):

        num_obst = test[0]
        max_num_obst = self.generator.max_box_num

        if num_obst >= max_num_obst:
            return test  # Already at max capacity

        new_obstacle = self.generator.get_random_box_vals()
        start = 1 + num_obst * 6
        test[start : start + 6] = new_obstacle
        test[0] += 1

        return test

    def _remove_obstacle(self, test):
        num_obst = test[0]
        if num_obst == 1:
            return test  # Nothing to remove

        idx = random.randint(0, num_obst - 1)
        start = 1 + idx * 6
        end = start + 6

        # Shift remaining elements left
        for i in range(end, 1 + 6 * num_obst):
            test[i - 6] = test[i]

        # Fill the last obstacle slot with a new random obstacle
        new_obstacle = self.generator.get_random_box_vals()
        fill_start = 1 + 6 * (num_obst - 1)
        test[fill_start : fill_start + 6] = new_obstacle

        test[0] -= 1

        return test
