"""All optimizers, including the actual backpropagation code."""

import logging
from abc import ABC, abstractmethod
from collections import OrderedDict
from collections.abc import Callable, Iterable, ValuesView
from itertools import chain
from typing import Any

import numpy as np

from .backend import NDArray, TensorDevice
from .grad_mode import no_grad, no_grad_fn
from .tensor import Parameter, Tensor, ones_like, tensor, zeros_like
from .utils import traverse_attrs

logger = logging.getLogger(__name__)


def toposort(root: Tensor) -> list[Tensor]:
    """Performs topological sort on a graph.

    `root` is the starting point of the graph.

    Args:
        root (Tensor): The starting point of the graph.
            Expected to have an attribute `src` denoting
            a list of its neighbors, also being of type `Tensor`.

    Raises:
        ValueError: If the graph, with the starting point
            given by `root`, is not a DAG.

    Returns:
        list[Tensor]: The ordered nodes.
    """
    ordered_nodes: list[Tensor] = []
    currently_visiting: set[Tensor] = set()
    done: set[Tensor] = set()

    stack: list[tuple[Tensor, bool]] = [(root, False)]
    while stack:
        node, visited = stack.pop()
        if node in done:
            continue

        if visited:
            ordered_nodes.append(node)
            currently_visiting.discard(node)
            done.add(node)
            continue

        stack.append((node, True))
        currently_visiting.add(node)
        for neighbor in reversed(node.src):
            if neighbor in done:
                continue
            if neighbor in currently_visiting:
                raise ValueError("Cycle in computation graph detected, but only DAG allowed!")
            stack.append((neighbor, False))

    return ordered_nodes


class Optimizer(ABC):
    """Abstract base class for all optimizers.

    The optimizer state is stored in a single `OrderedDict` (`_state`) with
    two levels:

    - **Global entries** such as `lr`, `beta_1`, `t`, etc. sit at the top
      level.
    - **Per-parameter entries** live under the `"params"` key, keyed by the
      parameter's name in the model state dict (e.g. `"layers[0].W"`).

    Per-parameter state is created lazily: `_get_or_create_param_state` calls
    the subclass hook `_init_param_state` the first time a parameter is
    encountered during `step()`.  This means frozen parameters that are never
    optimized never allocate state, and parameters that become trainable later
    get their state on first use.
    """

    def __init__(self, params: OrderedDict[str, Parameter], *, lr: float = 1e-3) -> None:
        """Initialize the optimizer.

        Args:
            params (OrderedDict[str, Parameter]): Named parameters to optimize,
                typically obtained via `model.get_parameters()`.
            lr (float, optional): The learning rate. Defaults to 1e-3.
        """
        if len(params) == 0:
            raise ValueError("Must pass at least one parameter to optimize.")
        for name, param in params.items():
            if not isinstance(param, Parameter):
                raise TypeError(
                    f'Parameter "{name}" passed to the optimizer must be of type Parameter, '
                    f"got {type(param).__name__}."
                )
            if not param.keep_grad:
                raise ValueError(
                    f'Attribute keep_grad must always be True for parameter "{name}" '
                    "to avoid clearing gradients during the backward pass. "
                    "This is important for cases like gradient accumulation."
                )
            if len(param.src) > 0:
                raise ValueError(
                    f'Parameter "{name}" should be a leaf and therefore '
                    "should not have any parents/src "
                    "from which it was created."
                )

        self.params = params
        self._state: OrderedDict[str, Any] = OrderedDict(
            {
                "lr": tensor(lr),
                "params": OrderedDict(),
            }
        )

    def _init_param_state(self, param: Parameter) -> OrderedDict[str, Tensor]:  # noqa: ARG002
        """Create the initial per-parameter state for *param*.

        Override in subclasses that need per-parameter state (momentum,
        variance estimates, etc.).  The base implementation returns an empty
        dict, which is appropriate for stateless optimizers like vanilla SGD.

        Args:
            param (Parameter): The parameter to create state for.

        Returns:
            OrderedDict[str, Tensor]: The initial state for the parameter.
        """
        return OrderedDict()

    def _get_or_create_param_state(self, name: str, param: Parameter) -> OrderedDict[str, Tensor]:
        """Return the per-parameter state for *name*, creating it if absent.

        Args:
            name (str): The parameter's name in the model state dict.
            param (Parameter): The parameter itself (passed through to
                `_init_param_state` on first access).

        Returns:
            OrderedDict[str, Tensor]: The (possibly newly created) state dict
                for this parameter.
        """
        param_states = self._state["params"]
        if name not in param_states:
            param_states[name] = self._init_param_state(param)
        return param_states[name]

    @property
    def lr(self) -> Tensor:
        """The learning rate."""
        return self._state["lr"]

    @lr.setter
    def lr(self, value: Tensor) -> None:
        """Set the learning rate.

        Args:
            value (Tensor): The new value for the learning rate.
        """
        self._state["lr"] = value

    def traverse_state(
        self,
        on_tensor: Callable[[str, Tensor], Tensor | None],
    ) -> None:
        """Recursively traverse all Tensors in the optimizer state and optionally transform them.

        Walks `_state` (which `traverse_attrs` reaches as an instance
        attribute).  The model parameters stored in `self.params` are
        explicitly excluded.

        Args:
            on_tensor (Callable[[str, Tensor], Tensor | None]): Callback called for each
                Tensor with (path, tensor). If it returns a Tensor, the original is replaced in-place.
                If it returns None, no replacement occurs.
        """

        def on_tensor_wrapper(path: str, tensor: Tensor) -> Tensor | None:
            if path.startswith("params"):
                return None
            return on_tensor(path, tensor)

        traverse_attrs(
            self,
            target_type=Tensor,
            on_target=on_tensor_wrapper,
        )

    @property
    def device(self) -> tuple[TensorDevice, ...]:
        """The devices on which the optimizer state is currently located.

        Can be multiple if the optimizer state is sharded across multiple
        devices.

        Returns:
            tuple[TensorDevice, ...]: The devices on which the optimizer
                state is currently located. Empty, if the optimizer
                has no state tensors.
        """
        return tuple({t.device for t in self.state})

    @property
    def update_params(self) -> Iterable[tuple[str, Parameter, NDArray]]:
        """An iterator over `(name, param, grad)` triples that should be updated.

        Filters out parameters that are frozen (`requires_grad is False`)
        or that have no gradient yet.  Yielding `grad` as a separate element
        gives the type checker proof that it is not `None`.

        Returns:
            Iterable[tuple[str, Parameter, NDArray]]: Iterator over
                (name, param, grad) triples.
        """
        for name, param in self.params.items():
            if param.requires_grad and param.grad is not None:
                yield name, param, param.grad

    # "no_grad_fn" technically not needed here, as optimizer state Tensors
    # are leaves (not part of a computation graph). We still annotate
    # just to be safe.
    @no_grad_fn
    def copy_to_device(self, device: TensorDevice) -> "Optimizer":
        """Copy the optimizer state to the specified `device`.

        Args:
            device (TensorDevice): The device to copy the state to.

        Returns:
            Optimizer: self, for method chaining.
        """

        def copy_tensor(_path: str, tensor: Tensor) -> Tensor:
            return tensor.copy_to_device(device)

        self.traverse_state(copy_tensor)
        return self

    # "no_grad_fn" technically not needed here, as optimizer state Tensors
    # are leaves (not part of a computation graph). We still annotate
    # just to be safe.
    @no_grad_fn
    def get_state(self, to_device: TensorDevice | None = None) -> OrderedDict[str, Tensor]:
        """The state of the optimizer as a flat dict.

        Recursively collects all Tensors from `_state` (global and
        per-parameter) into a single `OrderedDict` keyed by their
        traversal path.  The `params` attribute (the parameters themselves)
        is **not** part of the optimizer state.

        Args:
            to_device (TensorDevice | None): If specified, copy each
                Tensor in the state to this device in the returned dict.
                If `None`, the device of the Tensors is not changed.
                Defaults to None.

        Returns:
            OrderedDict[str, Tensor]: Dict containing the state.
        """
        result: OrderedDict[str, Tensor] = OrderedDict()

        def collect_tensors(path: str, tensor: Tensor) -> None:
            result[path] = tensor.copy_to_device(to_device) if to_device is not None else tensor

        self.traverse_state(collect_tensors)
        return result

    @property
    def state(self) -> ValuesView[Tensor]:
        """The Tensors forming the state of the optimizer.

        Returns:
            ValuesView[Tensor]: A view over the state Tensors.
        """
        return self.get_state().values()

    # "no_grad_fn" technically not needed here, as optimizer state Tensors
    # are leaves (not part of a computation graph). We still annotate
    # just to be safe.
    @no_grad_fn
    def load_state(
        self,
        state: OrderedDict[str, Tensor],
        *,
        match_device: bool = False,
        partial: bool = False,
    ) -> "Optimizer":
        """Load/initialize the state of the optimizer.

        Args:
            state (OrderedDict[str, Tensor]): The state of the optimizer.
            match_device (bool): If True, copy each loaded Tensor
                to the target's device before assignment.
                If False, raises on device mismatch. Defaults to False.
            partial (bool): If True, allow missing keys in `state`.
                If False, raises on missing keys. Defaults to False.

        Returns:
            Optimizer: self, for method chaining.
        """

        def load_tensor(key: str, tensor: Tensor) -> None:
            init_data = state.get(key, None)
            if init_data is None:
                if not partial:
                    raise KeyError(f'Optimizer state "{key}" not found in passed state!')
                return

            if not isinstance(init_data, Tensor):
                raise TypeError(
                    'Data in passed state must be of type "Tensor", '
                    f'found "{type(init_data).__name__}" ({init_data})'
                )

            if tensor.shape != init_data.shape:
                raise ValueError(
                    f"Shape of seed Tensor does not align with shape of "
                    f'target parameter "{key}". Found "{tensor.shape}", '
                    f'expected "{init_data.shape}".'
                )

            if match_device:
                init_data = init_data.copy_to_device(tensor.device)
            elif tensor.device != init_data.device:
                raise ValueError(
                    f"Device of seed Tensor does not align with device of "
                    f'target parameter "{key}". Found "{init_data.device}", '
                    f'expected "{tensor.device}".'
                )

            tensor.data = init_data.data
            return  # Don't replace

        self.traverse_state(load_tensor)
        return self

    def _clear_graph(self, topo_nodes: Iterable[Tensor]) -> None:
        """Clears computation graph structure and gradients after backward pass.

        Removes references to parent tensors, backward functions, and operation
        context to free memory. Gradients are preserved for nodes with
        `keep_grad=True` (i.e. parameters).

        Args:
            topo_nodes (Iterator[Tensor]): Nodes in topological order from the
                computation graph to clear.
        """
        for node in topo_nodes:
            node.src = ()
            node.backward_fn = None
            node.op_ctx = {}
            if not node.keep_grad:
                node.grad = None

    def _clear_activation_gradients(self, topo_nodes: Iterable[Tensor]) -> None:
        """Clears gradients of activation tensors before backward pass.

        Ensures activations start with no gradients, preventing stale gradients
        from previous backward passes if tensors are reused. Parameters
        (keep_grad=True) retain their gradients for accumulation.

        Args:
            topo_nodes (Iterator[Tensor]): Nodes in topological order from the
                computation graph whose gradients should be cleared.
        """
        for node in topo_nodes:
            if not node.keep_grad:
                node.grad = None

    def backward(self, loss: Tensor) -> None:
        """Perform backpropagation on the computation graph with respect to `loss`.

        Note: Does **not** support accumulating multiple losses
        into the gradient of `loss`. This only works for parameters.
        Use multiple losses and sum them instead.

        Args:
            loss (Tensor): The loss on with respect to which we perform
                gradient calculations.

        Raises:
            ValueError: If the loss is not a scalar.
        """
        if not isinstance(loss, Tensor) or loss.size > 1:
            raise ValueError("Expected 'loss' argument to be to be a scalar Tensor.")
        if loss.keep_grad:
            raise ValueError(
                "keep_grad=True not supported for the loss Tensor. "
                "Use multiple losses and accumulate/sum them into a new "
                "loss Tensor instead."
            )

        node_order = toposort(loss)

        # Clear the gradients of all activations:
        # (this is necessary to make "node.grad is None" below well-defined)
        self._clear_activation_gradients(topo_nodes=node_order)

        # no "np.ones_like" here, because loss.grad needs to be on the same device as "loss"
        loss.grad = ones_like(loss).data  # seed gradient for loss

        with no_grad():
            for node in reversed(node_order):
                # "node.grad is None" -> Means: In the current computation graph,
                #   this node has not received any gradients, meaning it does not
                #   contribute to the loss => We can skip it
                # "node.is_leaf()" -> Means: In the current computation graph,
                #   this node has no parents to pass gradients to, meaning we can skip it
                # "not any(src_requires_grad)" -> Means: In the current computation graph,
                #   not parent requires a gradient to be passed to it, meaning we can skip
                #   the current node
                src_requires_grad = [t.requires_grad for t in node.src]
                if node.grad is None or node.is_leaf() or not any(src_requires_grad):
                    continue

                if node.backward_fn is None:
                    raise ValueError(
                        f'"backward_fn" for node "{id(node)}" in computation graph is "None"!'
                    )

                logger.debug(f'Calling backward function: "{node.backward_fn.__name__}"')

                array_src = [s.data for s in node.src]

                src_grads = node.backward_fn(
                    *array_src,
                    compute_grad=src_requires_grad,
                    grad_out=node.grad,
                    **node.op_ctx,
                )
                if len(src_grads) != len(node.src):
                    raise RuntimeError(
                        f"backward_fn returned {len(src_grads)} gradients but node has {len(node.src)} sources"
                    )

                for src, src_grad in zip(node.src, src_grads, strict=True):
                    if src_grad is None:
                        continue

                    if src.shape != src_grad.shape:
                        raise RuntimeError(
                            f"Gradient shape {src_grad.shape} does not match source shape {src.shape}"
                        )

                    src.grad = src_grad if src.grad is None else src.grad + src_grad

        self._clear_graph(topo_nodes=node_order)

    def zero_grad(self, additional_tensors: list[Tensor] | None = None) -> None:
        """Clears the gradients of all parameters that are optimized.

        Applied to all parameters in `self.params`.

        Args:
            additional_tensors (list[Tensor], optional): Extra Tensors for
            which to reset gradients. This could be activations that have
            explicitly set their `keep_grad` attribute to `True`, meaning
            they are not cleared before the backward pass of the graph
            they are part of. Defaults to None.
        """
        for param in chain(self.params.values(), additional_tensors or []):
            param.grad = None

    @no_grad_fn
    @abstractmethod
    def step(self) -> None:
        """The step function to update the parameters.

        Must be implemented by the specific optimizer.
        """


class SGD(Optimizer):
    """Stochastic gradient descent optimizer."""

    def __init__(
        self,
        params: OrderedDict[str, Parameter],
        *,
        lr: float = 1e-3,
        friction: float = 1,
        weight_decay: float = 0,
    ) -> None:
        """The stochastic gradient descent optimizer.

        Note: By default, vanilla SGD is used. However, when setting
        the arguments accordingly, it can become SGD with momentum and also
        apply weight decay.

        **Standard SGD:** `friction=1, weight_decay=0`
        **SGD w/ momentum:** `friction<1, weight_decay=0`
        **SGDW:** `friction<1, weight_decay>0`

        Args:
            params (OrderedDict[str, Parameter]): Named parameters to optimize,
                typically obtained via `model.get_parameters()`.
            lr (float, optional): The learning rate. Defaults to 1e-3.
            friction (float, optional): How much friction to apply on the
                momentum. If friction is 1 (100%), then we do not use momentum,
                as in every step all previous momentum is lost.
                If momentum is desired, set `friction<1`. A typical value is
                `0.1`, so 10% of momentum is lost every step due to friction.
                Defaults to 1.
            weight_decay (float, optional): Decay rate of the weights/parameters,
                equals the weight of L2-regularization on the loss.
                If SGDW is used, a typical value is `0.01`.
                Defaults to `0`.
        """
        super().__init__(params=params, lr=lr)

        if not 0 <= friction <= 1:
            raise ValueError(f"friction must be in [0, 1], got {friction}")

        self._state["friction"] = tensor(friction)
        self._state["weight_decay"] = tensor(weight_decay)
        self._use_momentum = friction < 1

    def _init_param_state(self, param: Parameter) -> OrderedDict[str, Tensor]:
        """Create the initial per-parameter state for *param*.

        When using momentum, i.e. `friction < 1` this creates the momentum
        store for `param`.

        Args:
            param (Parameter): The parameter to create state for.

        Returns:
            OrderedDict[str, Tensor]: The initial state for the parameter.
                Empty if no momentum is used.
        """
        if self._use_momentum:
            return OrderedDict({"m": zeros_like(param, dtype=param.dtype, requires_grad=False)})
        return OrderedDict()

    @no_grad_fn
    def step(self) -> None:
        """Performs a single gradient descent step."""
        lr = self.lr
        friction = self._state["friction"]
        weight_decay = self._state["weight_decay"]

        for name, param, grad in self.update_params:
            s = self._get_or_create_param_state(name, param)

            if "m" in s:
                # we lose momentum through "friction", the momentum remaining
                # from the previous step is (1-friction)
                # we add param.grad as new force from the current position
                s["m"] = (1 - friction) * s["m"] + grad
                effective_grad = s["m"].data
            else:
                effective_grad = grad

            param[...] = (
                (1 - lr * weight_decay) * param  # weight decay part
                - lr * effective_grad  # gradient part
            )


class Adam(Optimizer):
    """Adam optimizer."""

    def __init__(
        self,
        params: OrderedDict[str, Parameter],
        *,
        lr: float = 1e-3,
        beta_1: float = 0.9,
        beta_2: float = 0.999,
        epsilon: float = 1e-8,
        weight_decay: float = 0,
    ) -> None:
        """The Adam optimizer.

        Note: By setting `weight_decay` > 0 this becomes `AdamW`.

        Args:
            params (OrderedDict[str, Parameter]): Named parameters to optimize,
                typically obtained via `model.get_parameters()`.
            lr (float, optional): The learning rate, also called `alpha`
                in the paper. Defaults to 1e-3.
            beta_1 (float, optional): Exponential decay rate for
                the momentum. Defaults to 0.9.
            beta_2 (float, optional): Exponential decay rate for
                the noise. Defaults to 0.999.
            epsilon (float, optional): Value added to `v` to improve
                numerical stability and avoid division by zero. Defaults to 1e-8.
            weight_decay (float, optional): Decay rate of the weights/parameters,
                equals the weight of L2-regularization on the loss.
                If greater than `0`, the optimizer becomes **AdamW**.
                If AdamW is used, a typical value is `0.01`.
                Defaults to `0`, meaning vanilla Adam is used.
        """
        super().__init__(params=params, lr=lr)
        self._state["beta_1"] = tensor(beta_1)
        self._state["beta_2"] = tensor(beta_2)
        self._state["epsilon"] = tensor(epsilon)
        self._state["t"] = tensor(0)
        self._state["weight_decay"] = tensor(weight_decay)

    def _init_param_state(self, param: Parameter) -> OrderedDict[str, Tensor]:
        """Create the initial per-parameter state for *param*.

        Creates the moving average stores for `param`.

        Args:
            param (Parameter): The parameter to create state for.

        Returns:
            OrderedDict[str, Tensor]: The initial state for the parameter.
        """
        return OrderedDict(
            {
                "m": zeros_like(param, dtype=param.dtype, requires_grad=False),
                "v": zeros_like(param, dtype=param.dtype, requires_grad=False),
            }
        )

    @no_grad_fn
    def step(self) -> None:
        """Performs a single Adam step.

        Uses the slightly more efficient variant, which
        can be found at the end of section 2 in the paper: https://arxiv.org/pdf/1412.6980
        """
        self._state["t"] = self._state["t"] + 1

        t = self._state["t"]
        beta_1 = self._state["beta_1"]
        beta_2 = self._state["beta_2"]
        epsilon = self._state["epsilon"]
        weight_decay = self._state["weight_decay"]
        lr = self.lr

        for name, param, grad in self.update_params:
            s = self._get_or_create_param_state(name, param)

            s["m"] = beta_1 * s["m"] + (1 - beta_1) * grad
            s["v"] = beta_2 * s["v"] + (1 - beta_2) * grad**2

            lr_t = lr * np.sqrt(1 - beta_2**t) / (1 - beta_1**t)
            epsilon_hat = epsilon * np.sqrt(1 - beta_2**t)

            # [...] -> in-place assignment
            param[...] = (
                (1 - lr * weight_decay) * param  # weight decay part
                - lr_t * s["m"] / (np.sqrt(s["v"]) + epsilon_hat)  # gradient part
            )


__all__ = [
    "SGD",
    "Adam",
    "Optimizer",
    "toposort",
]
