from typing import List, Optional, Union

from afnio._variable import Variable


class Parameter(Variable):
    """
    A kind of `Variable` that is to be considered a module parameter.

    Parameters are [`Variable`][afnio.Variable] subclasses, that have a very special
    property when used with [`Module`][afnio.cognitive.modules.Module]s - when assigned
    as an attribute to a `Module`, they are automatically added to the list of its
    parameters, and will appear e.g. in
    [`parameters()`][afnio.cognitive.modules.Module.parameters] iterator.
    Assigning a `Variable` doesn't have such effect. This is because one might want to
    cache some temporary state, like a piece of context memory for
    [`ChatCompletion`][afnio.cognitive.modules.ChatCompletion], in
    the agent. If there was no such class as [`Parameter`][.], these temporaries would
    get registered too.

    Attributes:
        data: The raw data, which can be a scalar string or number, or a sequence of
            such scalars.
        role: A specific description of the role of the parameter in the agent that
            provides context to the optimizer.
        requires_grad: Whether to track operations for automatic differentiation and
            compute gradients.
    """

    def __init__(
        self,
        data: Optional[Union[str, int, float, List[Union[str, int, float]]]] = None,
        role: str = None,
        requires_grad: bool = True,
    ):
        """
        Initializes a `Parameter` instance.

        Args:
            data: The raw data for the parameter, which can be a scalar string or
                number, or a sequence of such scalars.
            role: A specific description of the role of the parameter in the agent that
                provides context to the [`Optimizer`][afnio.optim.optimizer.Optimizer].
            requires_grad: A boolean indicating whether to compute gradients for this
                parameter during backpropagation. Defaults to `True`.
        """
        super().__init__(data=data, role=role, requires_grad=requires_grad)

    def __deepcopy__(self, memo):
        if id(self) in memo:
            return memo[id(self)]
        else:
            result = type(self)(self.data, self.role, self.requires_grad)
            memo[id(self)] = result
            return result

    def __repr__(self):
        return "Parameter containing:\n  " + super().__repr__()
