from afnio._variable import Variable
from afnio.autodiff.basic_ops import Add as AddOp

from .module import Module


class Add(Module):
    """
    Performs element-wise addition of two input [`Variable`][afnio.Variable]s.

    This module utilizes the [`Add`][afnio.autodiff.basic_ops.Add] operation from
    `afnio.autodiff.basic_ops`. The inputs must be instances of the
    [`Variable`][afnio.Variable] class. The `forward` method applies the addition
    operation to the [`data`][afnio.Variable.data] field of the inputs and returns
    the resulting [`Variable`][afnio.Variable].

    Note:
        This module does not have any trainable parameters.

    Examples:
        >>> from afnio import cognitive as cog
        >>> class Addition(cog.Module):
        ...     def __init__(self):
        ...         super().__init__()
        ...         self.add = cog.Add()
        >>>     def forward(self, x, y):
        ...         return self.add(x, y)
        >>> input1 = afnio.Variable(data="abc", role="input1")
        >>> input2 = afnio.Variable(data="def", role="input2")
        >>> addition = Addition()
        >>> result = addition(input1, input2)
        >>> print(result.data)
        'abcdef'
        >>> print(result.role)
        'input1 and input2'

    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.

    See Also:
        [`afnio.autodiff.basic_ops.Add`][afnio.autodiff.basic_ops.Add]
        for the underlying operation.
    """

    def __init__(self):
        super().__init__()

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

        Warning:
            Users should not call this method directly. Instead, they should call the
            module instance itself, which will internally invoke this `forward` method.

        Args:
            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.
        """
        return AddOp.apply(x, y)
