# SPDX-FileCopyrightText: 2026 bartzbeielstein
# SPDX-FileCopyrightText: Max Kuhn
# SPDX-License-Identifier: GPL-2.0-or-later
#
# The R 'desirability' package is maintained and developed by Max Kuhn.
# It is available on CRAN: https://CRAN.R-project.org/package=desirability
# DOI: https://doi.org/10.32614/CRAN.package.desirability
#
# This Python implementation (spotdesirability) is a port by bartzbeielstein.

import numpy as np
import matplotlib.pyplot as plt
from typing import Any, Dict


class DesirabilityBase:
    """Base class for all desirability functions.

    Provides a method to print class attributes and extend the range of values.

    Methods:
        print_class_attributes(indent=0):
            Prints the attributes of the class object in a generic and recursive manner.

        extend_range(values, factor=0.05):
            Extends the range of values by a given factor.

    References:
        Many thanks to Max Kuhn for his implementation of the 'desirability' package in R.
        This class is based on the 'desirability' package in R, see:
        https://cran.r-project.org/package=desirability
    """

    def __init__(self):
        """Initializes the DesirabilityBase object."""
        pass

    def print_class_attributes(self, indent=0):
        """
        Prints the attributes of the class object in a generic and recursive manner.

        Args:
            indent (int): The indentation level for nested objects.
        """
        # Print the class name of the current object
        print("\n" + " " * indent + f"Class: {type(self).__name__}")

        # Get the attributes of the object as a dictionary
        attributes = vars(self)
        for attr, value in attributes.items():
            if isinstance(value, DesirabilityBase):  # Check if the attribute is another desirability object
                print(" " * indent + f"{attr}:")
                value.print_class_attributes(indent=indent + 2)  # Recursive call with increased indentation
            elif isinstance(value, (list, tuple)) and all(isinstance(v, DesirabilityBase) for v in value):
                print(" " * indent + f"{attr}: [")
                for v in value:
                    v.print_class_attributes(indent=indent + 2)  # Recursive call for each object in the list/tuple
                print(" " * indent + "]")
            else:
                print(" " * indent + f"{attr}: {value}")

    def extend_range(self, values, factor=0.05):
        """Extend the range of values by a given factor."""
        range_span = max(values) - min(values)
        return [min(values) - factor * range_span, max(values) + factor * range_span]

    def plot(self, add: bool = False, non_inform: bool = True, xlabel="Input", ylabel="Desirability", figsize: tuple = (10, 8), grid=True, **kwargs: Dict[str, Any]) -> None:
        """
        Plots the continuous desirability functions, i.e., `DMax`, `DMin`, and `DTarget`.

        Args:
            add (bool, optional):
                Whether to add the plot to an existing figure. Defaults to False.
            non_inform (bool, optional):
                Whether to display the non-informative value as a dashed line. Defaults to True.
            xlabel (str, optional):
                Label for the x-axis. Defaults to "Input".
            ylabel (str, optional):
                Label for the y-axis. Defaults to "Desirability".
            figsize (tuple, optional):
                Size of the figure (width, height) in inches. Defaults to (10, 8).
            grid (bool, optional):
                Whether to display a grid. Defaults to True.
            **kwargs (Dict[str, Any]):
                Additional keyword arguments for the plot.

        Examples:
            ```{python}
            from spotdesirability import DMax
            dmax = DMax(low=0, high=10, scale=1)
            dmax.plot(figsize=(12, 6))  # Custom figure size
            ```
            ```{python}
            from spotdesirability import DMin
            dmin = DMin(low=0, high=10, scale=1)
            dmin.plot()  # Default figure size
            ```
            ```{python}
            from spotdesirability import DTarget
            dtarget = DTarget(low=0, target=5, high=10, low_scale=1, high_scale=1)
            dtarget.plot(figsize=(8, 4))
            ```
            ```{python}
            from spotdesirability import DArb
            x = [-5, 0, 5, 10]
            d = [0, 0.5, 1, 0.2]
            darb = DArb(x, d)
            darb.plot(figsize=(15, 10))
            ```
        """
        x_range = self.extend_range([self.low, self.high])
        if not add:
            plt.figure(figsize=figsize)  # Set the figure size
            plt.plot([], [])  # Create an empty plot
            plt.xlim(x_range)
            plt.ylim(0, 1)
            plt.grid(grid)
            plt.xlabel(xlabel)
            plt.ylabel(ylabel)
        plt.hlines(0, x_range[0], self.low, **kwargs)
        plt.hlines(1, self.high, x_range[1], **kwargs)
        input_values = np.linspace(self.low, self.high, 100)
        output_values = self.predict(input_values)
        plt.plot(input_values, output_values, **kwargs)
        if non_inform:
            plt.axhline(y=self.missing, linestyle="--", **kwargs)
        plt.show()


class DMax(DesirabilityBase):
    """
    Implements a desirability function for maximization.

    The desirability function assigns a value of 0 for inputs below the `low` threshold,
    a value of 1 for inputs above the `high` threshold, and scales the desirability
    between 0 and 1 for inputs within the range `[low, high]` using a specified scale factor.

    Attributes:
        low (float): The lower threshold for the desirability function.
        high (float): The upper threshold for the desirability function.
        scale (float): The scaling factor for the desirability function. Must be greater than 0.
        tol (float, optional): A tolerance value to replace desirability values of 0. Defaults to None.
        missing (float, optional): The value to use for missing inputs. Defaults to a non-informative value.

    Methods:
        predict(newdata=None, missing=None):
            Predicts the desirability values for the given input data.

        plot(add=False, non_inform=True, **kwargs):
            Plots the desirability function.

    References:
        Many thanks to Max Kuhn for his implementation of the 'desirability' package in R.
        This class is based on the 'desirability' package in R, see:
        https://cran.r-project.org/package=desirability

    Examples:
        ```{python}
        from spotdesirability import DMax
        import numpy as np
        import matplotlib.pyplot as plt
        # Create a DMax object
        dmax = DMax(low=0, high=10, scale=1)
        # Predict desirability for a range of inputs
        inputs = np.array([-5, 0, 5, 10, 15])
        desirability = dmax.predict(inputs)
        print(desirability)
        # [0. 0. 0.5 1. 1.]
        # Plot the desirability function
        dmax.plot()
        ```
    """

    def __init__(self, low, high, scale=1, tol=None, missing=None):
        """
        Initializes the DMax object.

        Args:
            low (float): The lower threshold for the desirability function.
            high (float): The upper threshold for the desirability function.
            scale (float): The scaling factor for the desirability function. Must be greater than 0.
            tol (float, optional): A tolerance value to replace desirability values of 0. Defaults to None.
            missing (float, optional): The value to use for missing inputs. Defaults to a non-informative value.

        Raises:
            ValueError: If `low` is greater than or equal to `high`.
            ValueError: If `scale` is less than or equal to 0.
        """
        super().__init__()
        if low >= high:
            raise ValueError("The low value must be less than the high value.")
        if scale <= 0:
            raise ValueError("The scale parameter must be greater than zero.")

        self.low = low
        self.high = high
        self.scale = scale
        self.tol = tol
        self.missing = missing
        if self.missing is None:
            self.missing = self._calculate_non_informative_value()

    def _calculate_non_informative_value(self) -> float:
        """
        Calculates a non-informative value for missing inputs.

        Returns:
            float: The mean desirability value over the range `[low, high]`.
        """
        test_seq = np.linspace(self.low, self.high, 100)
        return np.mean(self.predict(test_seq))

    def predict(self, newdata=None, missing=None) -> np.ndarray:
        """
        Predicts the desirability values for the given input data.

        Args:
            newdata (array-like, optional): The input data for which to compute desirability values.
                If None, an empty array is used. Defaults to None.
            missing (float, optional): The value to use for missing inputs. Defaults to the object's `missing` attribute.

        Returns:
            numpy.ndarray: The desirability values for the input data.

        Examples:
            ```{python}
            from spotdesirability import DMax
            import numpy as np
            dmax = DMax(low=0, high=10, scale=1)
            inputs = np.array([-5, 0, 5, 10, 15])
            desirability = dmax.predict(inputs)
            print(desirability)
            # [0. 0. 0.5 1. 1.]
            ```
        """
        if newdata is None:
            newdata = np.array([])
        else:
            newdata = np.atleast_1d(newdata)
        if missing is None:
            missing = self.missing

        out = np.full(len(newdata), np.nan)
        out[newdata < self.low] = 0
        out[newdata > self.high] = 1
        mask = (newdata <= self.high) & (newdata >= self.low)
        out[mask] = ((newdata[mask] - self.low) / (self.high - self.low)) ** self.scale
        out[np.isnan(out)] = missing
        if self.tol is not None:
            out[out == 0] = self.tol
        return out


class DMin(DesirabilityBase):
    """
    Implements a desirability function for minimization.

    The desirability function assigns a value of 1 for inputs below the `low` threshold,
    a value of 0 for inputs above the `high` threshold, and scales the desirability
    between 1 and 0 for inputs within the range `[low, high]` using a specified scale factor.

    Attributes:
        low (float): The lower threshold for the desirability function.
        high (float): The upper threshold for the desirability function.
        scale (float): The scaling factor for the desirability function. Must be greater than 0.
        tol (float, optional): A tolerance value to replace desirability values of 0. Defaults to None.
        missing (float, optional): The value to use for missing inputs. Defaults to a non-informative value.

    Methods:
        predict(newdata=None, missing=None):
            Predicts the desirability values for the given input data.

        plot(add=False, non_inform=True, **kwargs):
            Plots the desirability function.

    References:
        Many thanks to Max Kuhn for his implementation of the 'desirability' package in R.
        This class is based on the 'desirability' package in R, see:
        https://cran.r-project.org/package=desirability

    Examples:
        ```{python}
        from spotdesirability import DMin
        import numpy as np
        import matplotlib.pyplot as plt
        # Create a DMin object
        dmin = DMin(low=0, high=10, scale=1)
        # Predict desirability for a range of inputs
        inputs = np.array([-5, 0, 5, 10, 15])
        desirability = dmin.predict(inputs)
        print(desirability)
        # [1. 1. 0.5 0. 0.]
        # Plot the desirability function
        dmin.plot()
        ```
    """

    def __init__(self, low, high, scale=1, tol=None, missing=None):
        """
        Initializes the DMin object.

        Args:
            low (float): The lower threshold for the desirability function.
            high (float): The upper threshold for the desirability function.
            scale (float): The scaling factor for the desirability function. Must be greater than 0.
            tol (float, optional): A tolerance value to replace desirability values of 0. Defaults to None.
            missing (float, optional): The value to use for missing inputs. Defaults to a non-informative value.

        Raises:
            ValueError: If `low` is greater than or equal to `high`.
            ValueError: If `scale` is less than or equal to 0.
        """
        super().__init__()
        if low >= high:
            raise ValueError("The low value must be less than the high value.")
        if scale <= 0:
            raise ValueError("The scale parameter must be greater than zero.")

        self.low = low
        self.high = high
        self.scale = scale
        self.tol = tol
        self.missing = missing
        if self.missing is None:
            self.missing = self._calculate_non_informative_value()

    def _calculate_non_informative_value(self):
        """
        Calculates a non-informative value for missing inputs.

        Returns:
            float: The mean desirability value over the range `[low, high]`.
        """
        test_seq = np.linspace(self.low, self.high, 100)
        return np.mean(self.predict(test_seq))

    def predict(self, newdata=None, missing=None):
        """
        Predicts the desirability values for the given input data.

        Args:
            newdata (array-like, optional): The input data for which to compute desirability values.
                If None, an empty array is used. Defaults to None.
            missing (float, optional): The value to use for missing inputs. Defaults to the object's `missing` attribute.

        Returns:
            (numpy.ndarray): The desirability values for the input data.

        Examples:
            ```{python}
            from spotdesirability import DMin
            import numpy as np
            dmin = DMin(low=0, high=10, scale=1)
            inputs = np.array([-5, 0, 5, 10, 15])
            desirability = dmin.predict(inputs)
            print(desirability)
            # [1. 1. 0.5 0. 0.]
            ```
        """
        if newdata is None:
            newdata = np.array([])
        else:
            newdata = np.atleast_1d(newdata)
        if missing is None:
            missing = self.missing

        out = np.full(len(newdata), np.nan)
        out[newdata < self.low] = 1
        out[newdata > self.high] = 0
        mask = (newdata <= self.high) & (newdata >= self.low)
        out[mask] = ((newdata[mask] - self.high) / (self.low - self.high)) ** self.scale
        out[np.isnan(out)] = missing
        if self.tol is not None:
            out[out == 0] = self.tol
        return out


class DTarget(DesirabilityBase):
    """
    Implements a desirability function for target optimization.

    The desirability function assigns a value of 0 for inputs outside the range `[low, high]`,
    scales the desirability between 0 and 1 for inputs within `[low, target]` using `low_scale`,
    and scales the desirability between 1 and 0 for inputs within `[target, high]` using `high_scale`.

    Attributes:
        low (float): The lower threshold for the desirability function.
        target (float): The target value for the desirability function.
        high (float): The upper threshold for the desirability function.
        low_scale (float): The scaling factor for the desirability function below the target. Must be greater than 0.
        high_scale (float): The scaling factor for the desirability function above the target. Must be greater than 0.
        tol (float, optional): A tolerance value to replace desirability values of 0. Defaults to None.
        missing (float, optional): The value to use for missing inputs. Defaults to a non-informative value.

    Methods:
        predict(newdata=None, missing=None):
            Predicts the desirability values for the given input data.

        plot(add=False, non_inform=True, **kwargs):
            Plots the desirability function.

    References:
        Many thanks to Max Kuhn for his implementation of the 'desirability' package in R.
        This class is based on the 'desirability' package in R, see:
        https://cran.r-project.org/package=desirability

    Examples:
        ```{python}
        from spotdesirability import DTarget
        import numpy as np
        import matplotlib.pyplot as plt
        # Create a DTarget object
        dtarget = DTarget(low=0, target=5, high=10, low_scale=1, high_scale=1)
        # Predict desirability for a range of inputs
        inputs = np.array([-5, 0, 2.5, 5, 7.5, 10, 15])
        desirability = dtarget.predict(inputs)
        print(desirability)
        # [0.   0.   0.5  1.   0.5  0.   0.  ]
        # Plot the desirability function
        dtarget.plot()
        ```
    """

    def __init__(self, low, target, high, low_scale=1, high_scale=1, tol=None, missing=None):
        """
        Initializes the DTarget object.

        Args:
            low (float): The lower threshold for the desirability function.
            target (float): The target value for the desirability function.
            high (float): The upper threshold for the desirability function.
            low_scale (float): The scaling factor for the desirability function below the target. Must be greater than 0.
            high_scale (float): The scaling factor for the desirability function above the target. Must be greater than 0.
            tol (float, optional): A tolerance value to replace desirability values of 0. Defaults to None.
            missing (float, optional): The value to use for missing inputs. Defaults to a non-informative value.

        Raises:
            ValueError: If `low` is greater than or equal to `high`.
            ValueError: If `low` is greater than or equal to `target`.
            ValueError: If `target` is greater than or equal to `high`.
            ValueError: If `low_scale` or `high_scale` is less than or equal to 0.
        """
        super().__init__()
        if low >= high:
            raise ValueError("The low value must be less than the high value.")
        if low >= target:
            raise ValueError("The low value must be less than the target.")
        if target >= high:
            raise ValueError("The target value must be less than the high value.")
        if low_scale <= 0 or high_scale <= 0:
            raise ValueError("The scale parameters must be greater than zero.")

        self.low = low
        self.target = target
        self.high = high
        self.low_scale = low_scale
        self.high_scale = high_scale
        self.tol = tol
        self.missing = missing
        if self.missing is None:
            self.missing = self._calculate_non_informative_value()

    def _calculate_non_informative_value(self):
        """
        Calculates a non-informative value for missing inputs.

        Returns:
            float: The mean desirability value over the range `[low, high]`.
        """
        test_seq = np.linspace(self.low, self.high, 100)
        return np.mean(self.predict(test_seq))

    def predict(self, newdata=None, missing=None):
        """
        Predicts the desirability values for the given input data.

        Args:
            newdata (array-like, optional): The input data for which to compute desirability values.
                If None, an empty array is used. Defaults to None.
            missing (float, optional): The value to use for missing inputs. Defaults to the object's `missing` attribute.

        Returns:
            (numpy.ndarray): The desirability values for the input data.

        Examples:
            ```{python}
            from spotdesirability import DTarget
            import numpy as np
            # Create a DTarget object
            dtarget = DTarget(low=0, target=5, high=10, low_scale=1, high_scale=1)
            inputs = np.array([-5, 0, 2.5, 5, 7.5, 10, 15])
            desirability = dtarget.predict(inputs)
            print(desirability)
            # [0.   0.   0.5  1.   0.5  0.   0.  ]
            ```
        """
        if newdata is None:
            newdata = np.array([])
        else:
            newdata = np.atleast_1d(newdata)
        if missing is None:
            missing = self.missing

        out = np.full(len(newdata), np.nan)
        out[(newdata < self.low) | (newdata > self.high)] = 0
        mask_low = (newdata <= self.target) & (newdata >= self.low)
        out[mask_low] = ((newdata[mask_low] - self.low) / (self.target - self.low)) ** self.low_scale
        mask_high = (newdata <= self.high) & (newdata >= self.target)
        out[mask_high] = ((newdata[mask_high] - self.high) / (self.target - self.high)) ** self.high_scale
        out[np.isnan(out)] = missing
        if self.tol is not None:
            out[out == 0] = self.tol
        return out


class DArb(DesirabilityBase):
    """
    Implements an arbitrary desirability function.

    This class allows users to define a custom desirability function by specifying
    input values (`x`) and their corresponding desirability values (`d`).

    Attributes:
        x (numpy.ndarray): The input values for the desirability function.
        d (numpy.ndarray): The desirability values corresponding to the input values.
        tol (float, optional): A tolerance value to replace desirability values of 0. Defaults to None.
        missing (float, optional): The value to use for missing inputs. Defaults to a non-informative value.

    Methods:
        predict(newdata=None, missing=None):
            Predicts the desirability values for the given input data.

        plot(add=False, non_inform=True, **kwargs):
            Plots the desirability function.

    References:
        Many thanks to Max Kuhn for his implementation of the 'desirability' package in R.
        This class is based on the 'desirability' package in R, see:
        https://cran.r-project.org/package=desirability

    Examples:
        ```{python}
        from spotdesirability import DArb
        import numpy as np
        import matplotlib.pyplot as plt
        # Define input and desirability values
        x = [-5, 0, 5, 10]
        d = [0, 0.5, 1, 0.2]
        # Create a DArb object
        darb = DArb(x, d)
        # Predict desirability for a range of inputs
        inputs = np.array([-10, -5, 0, 5, 10, 15])
        desirability = darb.predict(inputs)
        print(desirability)
        # [0.  0.  0.5 1.  0.2 0.2]
        # Plot the desirability function
        darb.plot()
        ```
    """

    def __init__(self, x, d, tol=None, missing=None):
        """
        Initializes the DArb object.

        Args:
            x (list or numpy.ndarray): The input values for the desirability function.
            d (list or numpy.ndarray): The desirability values corresponding to the input values.
            tol (float, optional): A tolerance value to replace desirability values of 0. Defaults to None.
            missing (float, optional): The value to use for missing inputs. Defaults to a non-informative value.

        Raises:
            ValueError: If any desirability value is not in the range [0, 1].
            ValueError: If `x` and `d` do not have the same length.
            ValueError: If `x` or `d` has fewer than two values.
        """
        super().__init__()
        self.x = np.atleast_1d(x).astype(float)
        self.d = np.atleast_1d(d).astype(float)

        if np.any(self.d > 1) or np.any(self.d < 0):
            raise ValueError("The desirability values must be 0 <= d <= 1.")
        if len(self.x) != len(self.d):
            raise ValueError("x and d must have the same length.")
        if len(self.x) < 2 or len(self.d) < 2:
            raise ValueError("x and d must have at least two values.")
        self.tol = tol
        self.missing = missing
        if self.missing is None:
            self.missing = self._calculate_non_informative_value()

    def _calculate_non_informative_value(self):
        """
        Calculates a non-informative value for missing inputs.

        Returns:
            float: The mean desirability value over the range `[min(x), max(x)]`.
        """
        test_seq = np.linspace(min(self.x), max(self.x), 100)
        return np.mean(self.predict(test_seq))

    def predict(self, newdata=None, missing=None):
        """
        Predicts the desirability values for the given input data.

        Args:
            newdata (array-like, optional): The input data for which to compute desirability values.
                If None, an empty array is used. Defaults to None.
            missing (float, optional): The value to use for missing inputs. Defaults to the object's `missing` attribute.

        Returns:
            (numpy.ndarray): The desirability values for the input data.

        Examples:
            ```{python}
            from spotdesirability import DArb
            import numpy as np
            # Define input and desirability values
            x = [-5, 0, 5, 10]
            d = [0, 0.5, 1, 0.2]
            darb = DArb(x, d)
            inputs = np.array([-10, -5, 0, 5, 10, 15])
            desirability = darb.predict(inputs)
            print(desirability)
            # [0.  0.  0.5 1.  0.2 0.2]
            ```
        """
        if newdata is None:
            newdata = np.array([])
        else:
            newdata = np.atleast_1d(newdata)
        if missing is None:
            missing = self.missing

        out = np.full(len(newdata), np.nan)
        out[newdata < min(self.x)] = self.d[0]
        out[newdata > max(self.x)] = self.d[-1]

        in_between = (newdata >= min(self.x)) & (newdata <= max(self.x))
        if np.any(in_between):
            out[in_between] = np.interp(newdata[in_between], self.x, self.d)

        out[np.isnan(out)] = missing
        if self.tol is not None:
            out[out == 0] = self.tol
        return out

    def plot(self, add: bool = False, non_inform: bool = True, xlabel: str = "Input", ylabel: str = "Desirability", figsize: tuple = (10, 8), grid: bool = True, **kwargs: Dict[str, Any]) -> None:
        """
        Plots the arbitrary desirability function.

        Args:
            add (bool, optional):
                Whether to add the plot to an existing figure. Defaults to False.
            non_inform (bool, optional):
                Whether to display the non-informative value as a dashed line. Defaults to True.
            xlabel (str, optional):
                Label for the x-axis. Defaults to "Input".
            ylabel (str, optional):
                Label for the y-axis. Defaults to "Desirability".
            figsize (tuple, optional):
                Size of the figure (width, height) in inches. Defaults to (10, 8).
            grid (bool, optional):
                Whether to display a grid. Defaults to True.
            **kwargs (Dict[str, Any]):
                Additional keyword arguments for the plot.

        Examples:
            ```{python}
            from spotdesirability import DArb
            x = [-5, 0, 5, 10]
            d = [0, 0.5, 1, 0.2]
            darb = DArb(x, d)
            darb.plot()  # Default figure size
            darb.plot(figsize=(12, 6))  # Custom figure size
            ```
        """
        x_range = self.extend_range(self.x)
        if not add:
            plt.figure(figsize=figsize)  # Set the figure size
            plt.plot([], [])  # Create an empty plot
            plt.xlim(x_range)
            plt.ylim(0, 1)
            plt.grid(grid)
            plt.xlabel(xlabel)
            plt.ylabel(ylabel)
        input_values = np.linspace(x_range[0], x_range[1], 100)
        output_values = self.predict(input_values)
        plt.plot(input_values, output_values, **kwargs)
        if non_inform:
            plt.axhline(y=self.missing, linestyle="--", **kwargs)
        plt.show()


class DBox(DesirabilityBase):
    """
    Implements a box-like desirability function.

    The desirability function assigns a value of 1 for inputs within the range `[low, high]`
    and a value of 0 for inputs outside this range.

    Attributes:
        low (float): The lower threshold for the desirability function.
        high (float): The upper threshold for the desirability function.
        tol (float, optional): A tolerance value to replace desirability values of 0. Defaults to None.
        missing (float, optional): The value to use for missing inputs. Defaults to a non-informative value.

    Methods:
        predict(newdata=None, missing=None):
            Predicts the desirability values for the given input data.

        plot(add=False, non_inform=True, **kwargs):
            Plots the desirability function.

    References:
        Many thanks to Max Kuhn for his implementation of the 'desirability' package in R.
        This class is based on the 'desirability' package in R, see:
        https://cran.r-project.org/package=desirability

    Examples:
        ```{python}
        from spotdesirability import DBox
        import numpy as np
        import matplotlib.pyplot as plt
        # Create a DBox object
        dbox = DBox(low=-1.682, high=1.682)
        # Predict desirability for a range of inputs
        inputs = np.array([-3, -1.682, 0, 1.682, 3])
        desirability = dbox.predict(inputs)
        print(desirability)
        # [0. 1. 1. 1. 0.]
        # Plot the desirability function
        dbox.plot()
        ```
    """

    def __init__(self, low, high, tol=None, missing=None):
        """
        Initializes the DBox object.

        Args:
            low (float): The lower threshold for the desirability function.
            high (float): The upper threshold for the desirability function.
            tol (float, optional): A tolerance value to replace desirability values of 0. Defaults to None.
            missing (float, optional): The value to use for missing inputs. Defaults to a non-informative value.

        Raises:
            ValueError: If `low` is greater than or equal to `high`.
        """
        super().__init__()
        if low >= high:
            raise ValueError("The low value must be less than the high value.")

        self.low = low
        self.high = high
        self.tol = tol
        self.missing = missing
        if self.missing is None:
            self.missing = self._calculate_non_informative_value()

    def _calculate_non_informative_value(self):
        """
        Calculates a non-informative value for missing inputs.

        Returns:
            (float): The mean desirability value over the range `[low, high]`.
        """
        test_seq = np.linspace(self.low, self.high, 100)
        return np.mean(self.predict(test_seq))

    def predict(self, newdata=None, missing=None):
        """
        Predicts the desirability values for the given input data.

        Args:
            newdata (array-like, optional): The input data for which to compute desirability values.
                If None, an empty array is used. Defaults to None.
            missing (float, optional): The value to use for missing inputs. Defaults to the object's `missing` attribute.

        Returns:
            (numpy.ndarray): The desirability values for the input data.

        Examples:
            ```{python}
            from spotdesirability import DBox
            import numpy as np
            # Create a DBox object
            dbox = DBox(low=-1.682, high=1.682)
            inputs = np.array([-3, -1.682, 0, 1.682, 3])
            desirability = dbox.predict(inputs)
            print(desirability)
            # [0. 1. 1. 1. 0.]
            ```
        """
        if newdata is None:
            newdata = np.array([])
        else:
            newdata = np.atleast_1d(newdata)
        if missing is None:
            missing = self.missing

        out = np.full(len(newdata), np.nan)
        out[(newdata < self.low) | (newdata > self.high)] = 0
        out[(newdata >= self.low) & (newdata <= self.high)] = 1
        out[np.isnan(out)] = missing
        if self.tol is not None:
            out[out == 0] = self.tol
        return out

    def plot(self, add: bool = False, non_inform: bool = True, xlabel: str = "Input", ylabel: str = "Desirability", figsize: tuple = (10, 8), **kwargs: Dict[str, Any]) -> None:
        """
        Plots the box-like desirability function.

        Args:
            add (bool, optional):
                Whether to add the plot to an existing figure. Defaults to False.
            non_inform (bool, optional):
                Whether to display the non-informative value as a dashed line. Defaults to True.
            xlabel (str, optional):
                Label for the x-axis. Defaults to "Input".
            ylabel (str, optional):
                Label for the y-axis. Defaults to "Desirability".
            figsize (tuple, optional):
                Size of the figure (width, height) in inches. Defaults to (10, 8).
            **kwargs (Dict[str, Any]):
                Additional keyword arguments for the plot.

        Examples:
            ```{python}
            from spotdesirability import DBox
            dbox = DBox(low=-1.682, high=1.682)
            dbox.plot()  # Default figure size
            dbox.plot(figsize=(12, 6))  # Custom figure size
            ```
        """
        x_range = self.extend_range([self.low, self.high])
        if not add:
            plt.figure(figsize=figsize)  # Set the figure size
            plt.plot([], [])  # Create an empty plot
            plt.xlim(x_range)
            plt.ylim(0, 1)
            plt.xlabel(xlabel)
            plt.ylabel(ylabel)
        plt.hlines(0, x_range[0], self.low, **kwargs)
        plt.hlines(0, self.high, x_range[1], **kwargs)
        plt.vlines(self.low, 0, 1, **kwargs)
        plt.vlines(self.high, 0, 1, **kwargs)
        plt.hlines(1, self.low, self.high, **kwargs)
        if non_inform:
            plt.axhline(y=self.missing, linestyle="--", **kwargs)
        plt.show()


class DCategorical(DesirabilityBase):
    """
    Implements a desirability function for categorical inputs.

    This class allows users to define desirability values for specific categorical inputs.

    Attributes:
        values (dict): A dictionary where keys are category names (strings) and values are desirability scores (floats).
        tol (float, optional): A tolerance value to replace desirability values of 0. Defaults to None.
        missing (float, optional): The value to use for missing inputs. Defaults to a non-informative value.

    Methods:
        predict(newdata=None, missing=None):
            Predicts the desirability values for the given categorical input data.

        plot(non_inform=True, **kwargs):
            Plots the desirability function for the categorical inputs.

    References:
        Many thanks to Max Kuhn for his implementation of the 'desirability' package in R.
        This class is based on the 'desirability' package in R, see:
        https://cran.r-project.org/package=desirability

    Examples:
        ```{python}
        from spotdesirability import DCategorical
        import matplotlib.pyplot as plt
        # Define desirability values for categories
        values = {"A": 0.1, "B": 0.9, "C": 0.5}
        # Create a DCategorical object
        dcat = DCategorical(values)
        # Predict desirability for a list of categories
        inputs = ["A", "B", "C"]
        desirability = dcat.predict(inputs)
        print(desirability)
        # [0.1 0.9 0.5]
        # Plot the desirability function
        dcat.plot()
        ```
    """

    def __init__(self, values, tol=None, missing=None):
        """
        Initializes the DCategorical object.

        Args:
            values (dict): A dictionary where keys are category names (strings) and values are desirability scores (floats).
            tol (float, optional): A tolerance value to replace desirability values of 0. Defaults to None.
            missing (float, optional): The value to use for missing inputs. Defaults to a non-informative value.

        Raises:
            ValueError: If `values` has fewer than two entries.
            ValueError: If keys in `values` are not strings.
        """
        super().__init__()
        if len(values) < 2:
            raise ValueError("'values' should have at least two values.")
        if not all(isinstance(k, str) for k in values.keys()):
            raise ValueError("'values' should be a named dictionary.")

        self.values = values
        self.tol = tol
        self.missing = missing
        if self.missing is None:
            self.missing = self._calculate_non_informative_value()

    def _calculate_non_informative_value(self):
        """
        Calculates a non-informative value for missing inputs.

        Returns:
            (float): The mean desirability value across all categories.
        """
        return np.mean(list(self.values.values()))

    def predict(self, newdata=None, missing=None):
        """
        Predicts the desirability values for the given categorical input data.

        Args:
            newdata (list or array-like, optional): A list or array of categorical inputs.
                If None, an empty array is used. Defaults to None.
            missing (float, optional): The value to use for missing inputs. Defaults to the object's `missing` attribute.

        Returns:
            (numpy.ndarray): The desirability values for the input data.

        Raises:
            ValueError: If a category in `newdata` is not in the allowed categories.

        Examples:
            ```{python}
            from spotdesirability import DCategorical
            values = {"A": 0.1, "B": 0.9, "C": 0.5}
            dcat = DCategorical(values)
            inputs = ["A", "B", "C"]
            desirability = dcat.predict(inputs)
            print(desirability)
            # [0.1 0.9 0.5]
            ```
        """
        if newdata is None:
            newdata = np.array([])
        else:
            newdata = np.atleast_1d(newdata)
        if missing is None:
            missing = self.missing

        out = np.full(len(newdata), np.nan)
        for i, val in enumerate(newdata):
            if val in self.values:
                out[i] = self.values[val]
            else:
                raise ValueError(f"Value '{val}' not in allowed values: {list(self.values.keys())}")

        out[np.isnan(out)] = missing
        if self.tol is not None:
            out[out == 0] = self.tol
        return out

    def plot(self, non_inform: bool = True, xlabel: str = "Category", ylabel: str = "Desirability", figsize: tuple = (10, 8), **kwargs: Dict[str, Any]) -> None:
        """
        Plots the desirability function for the categorical inputs.

        Args:
            non_inform (bool, optional):
                Whether to display the non-informative value as a dashed line. Defaults to True.
            xlabel (str, optional):
                Label for the x-axis. Defaults to "Category".
            ylabel (str, optional):
                Label for the y-axis. Defaults to "Desirability".
            figsize (tuple, optional):
                Size of the figure (width, height) in inches. Defaults to (10, 8).
            **kwargs (Dict[str, Any]):
                Additional keyword arguments for the plot.

        Examples:
            ```{python}
            from spotdesirability import DCategorical
            values = {"A": 0.1, "B": 0.9, "C": 0.5}
            dcat = DCategorical(values)
            dcat.plot()  # Default figure size
            dcat.plot(figsize=(12, 6))  # Custom figure size
            ```
        """
        plt.figure(figsize=figsize)  # Set the figure size
        plt.bar(range(len(self.values)), list(self.values.values()), tick_label=list(self.values.keys()), **kwargs)
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        if non_inform:
            plt.axhline(y=self.missing, linestyle="--", **kwargs)
        plt.show()


class DOverall(DesirabilityBase):
    """
    Combines multiple desirability objects into an overall desirability function.

    This class computes the overall desirability by combining individual desirability objects
    using the geometric mean of their desirability scores.

    Attributes:
        d_objs (list): A list of desirability objects (e.g., DMax, DMin, DTarget, etc.).

    Methods:
        predict(newdata, all=False):
            Predicts the overall desirability based on the individual desirability objects.

    References:
        Many thanks to Max Kuhn for his implementation of the 'desirability' package in R.
        This class is based on the 'desirability' package in R, see:
        https://cran.r-project.org/package=desirability

    Examples:
        ```{python}
        from spotdesirability import DOverall, DMax, DMin
        import numpy as np
        # Create individual desirability objects
        dmax = DMax(low=0, high=10, scale=1)
        dmin = DMin(low=5, high=15, scale=1)
        # Combine them into an overall desirability object
        doverall = DOverall(dmax, dmin)
        # Predict overall desirability for a set of inputs
        inputs = np.array([[5, 10], [0, 15], [10, 5]])
        overall_desirability = doverall.predict(inputs)
        print(overall_desirability)
        # Predict individual and overall desirabilities
        individual, overall = doverall.predict(inputs, all=True)
        print("Individual:", individual)
        print("Overall:", overall)
        ```
    """

    def __init__(self, *d_objs):
        """
        Initializes the DOverall object.

        Args:
            *d_objs (obj): Instances of desirability classes (e.g., DMax, DTarget, etc.).

        Raises:
            ValueError: If any object is not an instance of a valid desirability class.
        """
        super().__init__()
        valid_classes = (DMax, DMin, DTarget, DArb, DBox, DCategorical)

        if not all(isinstance(obj, valid_classes) for obj in d_objs):
            raise ValueError("All objects must be instances of valid desirability classes.")

        self.d_objs = d_objs  # Store the desirability objects

    def predict(self, newdata, all=False):
        """
        Predicts the overall desirability based on the individual desirability objects.

        Args:
            newdata (list or numpy.ndarray): A list or array of predicted outcomes, one for each desirability object.
            all (bool, optional): Whether to return individual desirabilities along with the overall desirability.
                Defaults to False.

        Returns:
            (float or tuple):
                The overall desirability score, or a tuple of individual and overall desirabilities if `all=True`.

        Raises:
            ValueError: If the number of columns in `newdata` does not match the number of desirability objects.

        Examples:
            ```{python}
            from spotdesirability import DOverall, DMax, DMin
            import numpy as np
            # Create individual desirability objects
            dmax = DMax(low=0, high=10, scale=1)
            dmin = DMin(low=5, high=15, scale=1)
            doverall = DOverall(dmax, dmin)
            inputs = np.array([[5, 10], [0, 15], [10, 5]])
            overall_desirability = doverall.predict(inputs)
            print(overall_desirability)
            ```
        """
        # Ensure newdata is a NumPy array
        newdata = np.array(newdata)

        # Validate the shape of newdata
        if newdata.ndim == 1 and len(newdata) != len(self.d_objs):
            raise ValueError("The number of columns in newdata must match the number of desirability objects.")

        if newdata.ndim == 1:
            newdata = newdata.reshape(1, -1)  # Reshape 1D array to 2D array with one row

        if newdata.shape[1] != len(self.d_objs):
            raise ValueError("The number of columns in newdata must match the number of desirability objects.")

        # Compute individual desirabilities
        individual_desirabilities = [obj.predict(newdata[:, i]) for i, obj in enumerate(self.d_objs)]

        # Compute the geometric mean of the individual desirabilities
        overall_desirability = np.prod(individual_desirabilities, axis=0) ** (1 / len(individual_desirabilities))

        if all:
            return individual_desirabilities, overall_desirability
        return overall_desirability
