from typing import Dict, Optional, Tuple, cast

import torch
from torch import Tensor
from torch.nn import functional as F  # noqa: N812

from .base import VIModule
from .distributions import MeanFieldNormal
from .utils.common_types import VIkwargs, _dist_any_t


class VILinear(VIModule):
    """
    Applies an affine linear transformation to the incoming data: :math:`y = xA^T + b`.

    Equivalent of :class:`nn.Linear` with variational inference. See its
    `documentation <https://pytorch.org/docs/stable/generated/torch.nn.Linear.html>`__
    for usage.

    In addition to those arguments, this class accepts :class:`~.VIkwargs`.

    This module's random variables are

    - ("weight", "bias") if bias == True
    - ("weight", )       if bias == False

    Parameters
    ----------
    torch_args
        The same arguments and keyword arguments as the pytorch version
        :class:`~nn.Linear` (documentation
        `here <https://pytorch.org/docs/stable/generated/torch.nn.Linear.html>`__)
    VIkwargs
        Several standard keyword arguments. See :class:`~.VIkwargs` for details.
    """

    __constants__ = ["in_features", "out_features"]
    in_features: int
    out_features: int
    _fast_path_flag: Optional[bool] = None

    def __init__(
        self,
        in_features: int,
        out_features: int,
        variational_distribution: _dist_any_t = MeanFieldNormal(),
        prior: _dist_any_t = MeanFieldNormal(),
        bias: bool = True,
        rescale_prior: bool = False,
        kaiming_initialization: bool = True,
        prior_initialization: bool = False,
        return_log_probs: bool = True,
        device: Optional[torch.device] = None,
        dtype: Optional[torch.dtype] = None,
    ) -> None:
        vikwargs: VIkwargs = dict(
            variational_distribution=variational_distribution,
            prior=prior,
            rescale_prior=rescale_prior,
            kaiming_initialization=kaiming_initialization,
            prior_initialization=prior_initialization,
            return_log_probs=return_log_probs,
            device=device,
            dtype=dtype,
        )
        self.in_features = in_features
        self.out_features = out_features

        variable_shapes: Dict[str, Optional[Tuple[int, ...]]] = dict(
            weight=(out_features, in_features)
        )
        if bias:
            variable_shapes["bias"] = (out_features,)
        else:
            variable_shapes["bias"] = None

        super().__init__(variable_shapes=variable_shapes, **vikwargs)

    @property
    def _fast_path(self) -> bool:
        # If the variational distribution is stable we might be able to use the stable fast path
        if self._fast_path_flag is not None:
            return self._fast_path_flag

        if all(
            (dist is None) or isinstance(dist, MeanFieldNormal)
            for dist in self.variational_distribution.values()
        ):
            self._fast_path_flag = True
            return True
        else:
            self._fast_path_flag = False
            return False

    def forward(self, input_: Tensor) -> Tensor:
        r"""
        Forward computation.

        Parameters
        ----------
        input_: Tensor
            Input tensor of shape (\*, in_features).

        Returns
        -------
        output: Tensor
            Output tensor of shape (\*, out_features). Auto-sampling will add a sample
            dimension at the start for the overall output.
        """
        # Check for and perform fast path if possible:
        if (not self._return_log_probs) and self._fast_path:
            output = self._fast_forward(input_)
            return output

        output = F.linear(input_, self.weight, self.bias)

        return output

    def _fast_forward(self, input_: Tensor) -> Tensor:
        """Perform the stable fast path for Gaussian variational distribution."""
        weight_mean = self._weight_mean
        weight_variance = cast(Tensor, 2 * self._weight_log_std).exp()
        if self.variational_distribution["bias"] is not None:
            bias_mean = self._bias_mean
            bias_variance = cast(Tensor, 2 * self._bias_log_std).exp()
        else:
            bias_mean = None
            bias_variance = None
        output_mean = F.linear(input_, weight_mean, bias_mean)
        output_std = F.linear(input_.pow(2), weight_variance, bias_variance).sqrt()
        output = MeanFieldNormal._normal_sample(output_mean, output_std)
        return output
