afnio.autodiff.basic_ops

class Add(Function):
    """
    Implements an addition operation for [`Variable`][afnio.Variable] instances within
    the `afnio` framework, supporting automatic differentiation.

    This class inherits from [`Function`][afnio.autodiff.function.Function] and
    requires both the [`forward`][afnio.autodiff.function.Function.forward] and
    [`backward`][afnio.autodiff.function.Function.backward] methods to be defined.

    The `Add` function supports both scalar and list [`data`][afnio.Variable.data]
    fields:

    - **Scalars**: Adds numerical values (`int`, `float`) or concatenates strings.
    - **Lists**: Performs element-wise addition of corresponding elements from the lists.
      Lists must be of the same length.

    It automatically handles type-based operations:

    - For numerical data (`int`, `float`), it performs arithmetic addition.
    - For strings, it concatenates the values.
    - Mixed types (e.g., string and number) are converted appropriately before performing
      the addition.

    This operation also tracks [`Variable`][afnio.Variable] dependencies,
    enabling automatic gradient computation through backpropagation.

    Examples:
        Example with scalar inputs:
        >>> x = Variable(data="abc", role="first input", requires_grad=True)
        >>> y = Variable(data="def", role="second input", requires_grad=False)
        >>> result = Add.apply(x, y)
        >>> result.data
        'abcdef'
        >>> result.role
        'first input and second input'
        >>> result.requires_grad
        True
        >>> g = Variable(data="MY_FEEDBACK", role="add gradient")
        >>> result.backward(g)
        >>> x.grad[0].data
        'Here is the combined feedback we got for this specific first input and other variables: MY_FEEDBACK'
        >>> x.grad[0].role
        'feedback to first input'

        Example with batched inputs:
        >>> x = Variable(data=[1, 2, 3], role="first input", requires_grad=True)
        >>> y = Variable(data=[4, 5, 6], role="second input", requires_grad=False)
        >>> result = Add.apply(x, y)
        >>> result.data
        [5, 7, 9]
        >>> result.role
        'first input and second input'
        >>> result.requires_grad
        True
    """  # noqa: E501

    @staticmethod
    def forward(ctx, x: Variable, y: Variable) -> Variable:
        """
        Forward pass for element-wise addition.

        Warning:
            This method is invoked by
            [`apply()`][afnio.autodiff.function.Function.apply]
            and should not be called directly.

        Args:
            ctx: Context object used to save information for [`backward`][..backward]
                computation.
            x: The first input `Variable`.
            y: The second input `Variable`.

        Returns:
            A new `Variable` instance representing the result of the addition, \
            with appropriately aggregated [`data`][afnio.Variable.data], \
            [`role`][afnio.Variable.role], and \
            [`requires_grad`][afnio.Variable.requires_grad] attributes.

        Raises:
            TypeError: If either input is not an instance
                of [`Variable`][afnio.Variable].
            TypeError: If addition between the input types is not allowed.
            ValueError: If a scalar [`data`][afnio.Variable.data] is added
                to a list[`data`][afnio.Variable.data].
            ValueError: If list [`data`][afnio.Variable.data] fields
                have mismatched lengths.
        """
        raise NotImplementedError(
            "Add.forward is implemented on the server. "
            "Client-side execution is not supported."
        )

    @staticmethod
    def backward(
        ctx, grad_output: Variable
    ) -> Tuple[Optional[Variable], Optional[Variable]]:
        """
        Backward pass for elementwise addition.

        Warning:
            This method is invoked by the autodiff engine
            and should not be called directly.

        Args:
            ctx: Context object containing saved information from the
                [`forward`][..forward] pass.
            grad_output: The gradient of the output `Variable` with respect to
                the output of the `forward()` method.

        Returns:
            grad_x: The gradient of the output w.r.t. the first input `Variable`
                if it requires gradients, otherwise `None`.
            grad_y: The gradient of the output w.r.t. the second input `Variable`
                if it requires gradients, otherwise `None`.

        Raises:
            TypeError: If `grad_output` is not an instance
                of [`Variable`][afnio.Variable].
        """
        raise NotImplementedError(
            "Add.backward is implemented on the server. "
            "Client-side execution is not supported."
        )

class Sum(Function):
    """
    Implements a summation operation for a list of [`Variable`][afnio.Variable]
    instances within the `afnio` framework, supporting automatic differentiation.

    This class inherits from [`Function`][afnio.autodiff.function.Function] and
    requires both the [`forward`][afnio.autodiff.function.Function.forward] and
    [`backward`][afnio.autodiff.function.Function.backward] methods to be defined.

    The `Sum` function aggregates the [`data`][afnio.Variable.data],
    [`role`][afnio.Variable.role], and [`requires_grad`][afnio.Variable.requires_grad]
    attributes of all input [`Variable`][afnio.Variable] instances into a single
    [`Variable`][afnio.Variable]. It supports both scalar and list
    [`data`][afnio.Variable.data] fields:

    - **Scalars**: Computes the arithmetic sum for numerical data (`int`, `float`)
      or concatenates all string values, wrapping each in `<ITEM></ITEM>` tags.
    - **Lists**: Aggregates the corresponding elements of the lists. For numerical
      data, it sums the corresponding elements. For string data, it concatenates them,
      wrapping each element in `<ITEM></ITEM>` tags.

    During backpropagation, the function distributes the gradient to all input
    [`Variable`][afnio.Variable] instances that require gradients.

    Examples:
        Example with scalar inputs:
        >>> x = Variable(data="abc", role="first input", requires_grad=True)
        >>> y = Variable(data="def", role="second input", requires_grad=False)
        >>> result = Sum.apply([x, y])
        >>> result.data
        '<ITEM>abc</ITEM><ITEM>def</ITEM>'
        >>> result.role
        'first input and second input'
        >>> result.requires_grad
        True
        >>> g = Variable(data="MY_FEEDBACK", role="add gradient")
        >>> result.backward(g)
        >>> x.grad[0].data
        'Here is the combined feedback we got for this specific first input and other variables: MY_FEEDBACK'
        >>> x.grad[0].role
        'feedback to first input'

        Example with batched inputs:
        >>> x = Variable(data=[1, 2, 3.5], role="first input", requires_grad=True)
        >>> y = Variable(data=[4, 5, 6], role="second input", requires_grad=False)
        >>> result = Sum.apply([x, y])
        >>> result.data
        [5, 7, 9.5]
        >>> result.role
        'first input and second input'
        >>> result.requires_grad
        True
    """  # noqa: E501

    @staticmethod
    def forward(ctx, x: List[Variable]) -> Variable:
        """
        Forward pass for summation.

        Warning:
            This method is invoked by
            [`apply()`][afnio.autodiff.function.Function.apply]
            and should not be called directly.

        Args:
            ctx: Context object used to save information for [`backward`][..backward]
                computation.
            x: A list of `Variable` instances to be summed.

        Returns:
            A new `Variable` instance representing the result of the summation, \
            with appropriately aggregated [`data`][afnio.Variable.data], \
            [`role`][afnio.Variable.role], and \
            [`requires_grad`][afnio.Variable.requires_grad] attributes.

        Raises:
            TypeError: If any element in `x` is not an instance
                of [`Variable`][afnio.Variable] or a sequence
                of [`Variable`][afnio.Variable] instances, or if addition between
                the [`data`][afnio.Variable.data] types is not allowed.
        """
        raise NotImplementedError(
            "Sum.forward is implemented on the server. "
            "Client-side execution is not supported."
        )

    @staticmethod
    def backward(ctx, grad_output: Variable) -> Tuple[Optional[Variable], ...]:
        """
        Backward pass for summation.

        Warning:
            This method is invoked by the autodiff engine
            and should not be called directly.

        Args:
            ctx: Context object containing saved information from the
                [`forward`][..forward] pass.
            grad_output: The gradient of the output `Variable` with respect to
                the output of the `forward()` method.

        Returns:
            grad_inputs: A tuple of gradients corresponding to each input `Variable`
                in the same order as the inputs to `forward()`. Each gradient is a
                `Variable` if the corresponding input requires gradients,
                otherwise `None`.

        Raises:
            TypeError: If `grad_output` is not an instance
                of [`Variable`][afnio.Variable].
        """
        raise NotImplementedError(
            "Sum.backward is implemented on the server. "
            "Client-side execution is not supported."
        )

class Split(Function):
    """
    Implements a split operation for [`Variable`][afnio.Variable] instances within the
    `afnio` framework, supporting automatic differentiation.

    This class inherits from [`Function`][afnio.autodiff.function.Function] and
    requires both the [`forward`][afnio.autodiff.function.Function.forward] and
    [`backward`][afnio.autodiff.function.Function.backward] methods to be defined.

    The `Split` function divides the [`data`][afnio.Variable.data] of the input
    [`Variable`][afnio.Variable] into multiple parts using a specified delimiter `sep`.
    If `maxsplit` is specified, the split operation is limited to a maximum number of
    splits. It handles both scalar and list [`data`][afnio.Variable.data] fields:

    - **Scalars**: The scalar [`data`][afnio.Variable.data] (a single string) is split
        into substrings based on the specified `sep` and `maxsplit` parameters.
    - **Lists**: Each element of the list [`data`][afnio.Variable.data] (strings) is
        split individually. If splits of varying lengths occur, shorter splits are
        automatically padded with empty strings to ensure consistent dimensions.

    During backpropagation, feedback is collected and aggregated across all split parts.
    The combined feedback is propagated back to the original input
    [`Variable`][afnio.Variable], allowing for the proper computation of gradients.

    Examples:
        Example with scalar inputs:
        >>> x = Variable(data="afnio is great!", role="sentence", requires_grad=True)
        >>> result = Split.apply(x, sep=" ", maxsplit=1)
        >>> [var.data for var in result]
        ['afnio', 'is great!']
        >>> result[0].role
        'split part 0 of sentence'
        >>> g_1 = Variable(data="MY_FIRST_FEEDBACK", role="gradient")
        >>> g_2 = Variable(data="MY_SECOND_FEEDBACK", role="gradient")
        >>> result[0].backward(g_1, retain_graph=True)
        >>> result[1].backward(g_2)
        >>> x.grad[0].data
        'Here is the combined feedback we got for this specific sentence and other variables: <ITEM>MY_FIRST_FEEDBACK</ITEM><ITEM></ITEM>'
        >>> x.grad[0].role
        'feedback to sentence'
        >>> x.grad[1].data
        'Here is the combined feedback we got for this specific sentence and other variables: <ITEM></ITEM><ITEM>MY_SECOND_FEEDBACK</ITEM>'
        >>> x.grad[1].role
        'feedback to sentence'

        Example with batched inputs:
        >>> x = Variable(
        ...     data=["afnio is great!", "Deep learning"],
        ...     role="sentences",
        ...     requires_grad=True
        ... )
        >>> result = Split.apply(x, sep=" ", maxsplit=2)
        >>> [var.data for var in result]
        [['afnio', 'Deep'], ['is', 'learning'], ['great!', '']]
        >>> g = Variable(data="MY_FEEDBACK", role="gradient")
        >>> result[1].backward(g)
        >>> x.grad[0].data
        'Here is the combined feedback we got for this specific sentences and other variables: <ITEM></ITEM><ITEM>MY_FEEDBACK</ITEM><ITEM></ITEM>'
        >>> x.grad[0].role
        'feedback to sentences'
    """  # noqa: E501

    @staticmethod
    def forward(
        ctx,
        x: Variable,
        sep: Optional[Union[str, Variable]] = None,
        maxsplit: Optional[Union[int, Variable]] = -1,
    ) -> Tuple[Variable]:
        """
        Forward pass for splitting a `Variable`.

        Warning:
            This method is invoked by
            [`apply()`][afnio.autodiff.function.Function.apply]
            and should not be called directly.

        Args:
            ctx: Context object used to save information for [`backward`][..backward]
                computation.
            x: The input `Variable` to be split.
            sep: The delimiter to use for splitting the string. If `None`, splits on
                whitespace. Can be a string or a `Variable` containing a string.
            maxsplit: The maximum number of splits to perform. If `-1`, there is no
                limit on the number of splits. Can be an integer or a `Variable`
                containing an integer.

        Returns:
            A tuple of `Variable` instances resulting from the split operation, \
            each with appropriately assigned [`data`][afnio.Variable.data], \
            [`role`][afnio.Variable.role], and \
            [`requires_grad`][afnio.Variable.requires_grad] attributes.

        Raises:
            TypeError: If `x` is not an instance of [`Variable`][afnio.Variable] which
            [`data`][afnio.Variable.data] attribute is a string or a list of strings,
            or if `sep` is not a string or `Variable` containing a string,
            or if `maxsplit` is not an integer or `Variable` containing an integer.
        """
        raise NotImplementedError(
            "Split.forward is implemented on the server. "
            "Client-side execution is not supported."
        )

    # TODO: enable summarization of elements in `grad.data`` fields in `backward()`
    #       method using a LM (ideally a small and fast SLM)
    @staticmethod
    def backward(ctx, *grad_outputs: Variable) -> Tuple[Optional[Variable], None, None]:
        """
        Backward pass for splitting a `Variable`.

        Warning:
            This method is invoked by the autodiff engine
            and should not be called directly.

        Args:
            ctx: Context object containing saved information from the
                [`forward`][..forward] pass.
            grad_outputs: A variable number of gradients corresponding to each split
                output `Variable` with respect to the output of the `forward()` method.

        Returns:
            grad_input: The gradient of the output w.r.t. the input `Variable` if it
                requires gradients, otherwise `None`. The feedback from all split parts
                is combined and propagated back to the original input `Variable`.
            None: Placeholder for the `sep` argument, which does not require a gradient.
            None: Placeholder for the `maxsplit` argument, which does not require
                a gradient.

        Raises:
            TypeError: If any element in `x` is not an instance
                of [`Variable`][afnio.Variable] or a sequence
                of [`Variable`][afnio.Variable] instances.
        """
        raise NotImplementedError(
            "Split.backward is implemented on the server. "
            "Client-side execution is not supported."
        )