from ceco.benchmark import Benchmark
import numpy as np


class Schwefel1_2(Benchmark):
    """
    A class representing the Schwefel 1.2 function, which is a benchmark function
    for optimization.

    Inherits from the Benchmark class and applies rotation and shift to the
    input vector before evaluating the Schwefel 1.2 function.

    Attributes:
        dimension (int): The number of dimensions for the input space. Must be a positive integer.
        rotation (np.ndarray): A rotation matrix for transforming the input vector.
        shift (np.ndarray): A shift vector for adjusting the input vector.
        f_bias (float): A bias term added to the benchmark function's output. Defaults to 0.
    """

    def __init__(self, dimension: int, rotation: np.ndarray, shift: np.ndarray, f_bias: float = 0) -> None:
        """
        Initializes the Schwefel1_2 class with a rotation matrix and a shift vector.

        Parameters:
            dimension (int): The number of dimensions for the input space. Must be a positive integer.
            rotation (np.ndarray): The rotation matrix for transforming the input vector.
            shift (np.ndarray): The shift vector for adjusting the input vector.
            f_bias (float): A bias term added to the benchmark function's output. Defaults to 0.
        """
        super().__init__(dimension)
        self.cec_init(rotation, shift, f_bias)

    def evaluate(self, input_vector: np.ndarray) -> float:
        """
        Evaluates the Schwefel 1.2 function at a given input vector after applying
        rotation and shift.

        Parameters:
            input_vector (np.ndarray): The input vector [x1, x2, ..., xD].

        Returns:
            float: The result of the Schwefel 1.2 function after applying rotation
            and shift.
        """
        # Apply shift and rotation
        shifted_rotated_vector = np.matmul(
            self.rotation, input_vector - self.shift)

        # Compute the Schwefel 1.2 function
        cumulative_sums = np.cumsum(shifted_rotated_vector)
        total_sum = np.sum(cumulative_sums ** 2)

        return total_sum + self.f_bias
