import warnings
from collections.abc import Sequence
from copy import copy
from typing import cast

import numpy as np

import pytensor
from pytensor.configdefaults import config
from pytensor.graph.basic import Apply, Variable, equal_computations
from pytensor.graph.op import Op
from pytensor.graph.replace import _vectorize_node
from pytensor.misc.safe_asarray import _asarray
from pytensor.scalar import ScalarVariable
from pytensor.tensor.basic import (
    as_tensor_variable,
    concatenate,
    constant,
    get_vector_length,
    infer_static_shape,
)
from pytensor.tensor.blockwise import OpWithCoreShape
from pytensor.tensor.random.type import RandomGeneratorType, RandomType
from pytensor.tensor.random.utils import (
    compute_batch_shape,
    explicit_expand_dims,
    normalize_size_param,
)
from pytensor.tensor.shape import shape_tuple
from pytensor.tensor.type import TensorType
from pytensor.tensor.type_other import NoneConst, NoneTypeT
from pytensor.tensor.utils import _parse_gufunc_signature, safe_signature
from pytensor.tensor.variable import TensorVariable


class RandomVariable(Op):
    """An `Op` that produces a sample from a random variable.

    This is essentially `RandomFunction`, except that it removes the
    `outtype` dependency and handles shape dimension information more
    directly.

    """

    _output_type_depends_on_input_value = True

    __props__ = ("name", "signature", "dtype", "inplace")
    default_output = 1

    def __init__(
        self,
        name=None,
        ndim_supp=None,
        ndims_params=None,
        dtype: str | None = None,
        inplace=None,
        signature: str | None = None,
    ):
        """Create a random variable `Op`.

        Parameters
        ----------
        name: str
            The `Op`'s display name.
        signature: str
            Numpy-like vectorized signature of the random variable.
        dtype: str (optional)
            The default dtype of the sampled output.  If the value ``"floatX"`` is
            given, then ``dtype`` is set to ``pytensor.config.floatX``.  If
            ``None`` (the default), the `dtype` keyword must be set when
            `RandomVariable.make_node` is called.
        inplace: boolean (optional)
            Determine whether the underlying rng state is mutated or copied.

        """
        super().__init__()

        self.name = name or getattr(self, "name")

        ndim_supp = (
            ndim_supp if ndim_supp is not None else getattr(self, "ndim_supp", None)
        )
        if ndim_supp is not None:
            warnings.warn(
                "ndim_supp is deprecated. Provide signature instead.", FutureWarning
            )
            self.ndim_supp = ndim_supp
        ndims_params = (
            ndims_params
            if ndims_params is not None
            else getattr(self, "ndims_params", None)
        )
        if ndims_params is not None:
            warnings.warn(
                "ndims_params is deprecated. Provide signature instead.", FutureWarning
            )
            if not isinstance(ndims_params, Sequence):
                raise TypeError("Parameter ndims_params must be sequence type.")
            self.ndims_params = tuple(ndims_params)

        self.signature = signature or getattr(self, "signature", None)
        if self.signature is not None:
            # Assume a single output. Several methods need to be updated to handle multiple outputs.
            self.inputs_sig, [self.output_sig] = _parse_gufunc_signature(self.signature)
            self.ndims_params = [len(input_sig) for input_sig in self.inputs_sig]
            self.ndim_supp = len(self.output_sig)
        else:
            if (
                getattr(self, "ndim_supp", None) is None
                or getattr(self, "ndims_params", None) is None
            ):
                raise ValueError("signature must be provided")
            else:
                self.signature = safe_signature(self.ndims_params, [self.ndim_supp])

        self.dtype = dtype or getattr(self, "dtype", None)

        self.inplace = (
            inplace if inplace is not None else getattr(self, "inplace", False)
        )

        if self.inplace:
            self.destroy_map = {0: [0]}

    def _supp_shape_from_params(self, dist_params, param_shapes=None):
        """Determine the support shape of a multivariate `RandomVariable`'s output given its parameters.

        This does *not* consider the extra dimensions added by the `size` parameter
        or independent (batched) parameters.

        When provided, `param_shapes` should be given preference over `[d.shape for d in dist_params]`,
        as it will avoid redundancies in PyTensor shape inference.

        Examples
        --------
        Common multivariate `RandomVariable`s derive their support shapes implicitly from the
        last dimension of some of their parameters. For example `multivariate_normal` support shape
        corresponds to the last dimension of the mean or covariance parameters, `support_shape=(mu.shape[-1])`.
        For this case the helper `pytensor.tensor.random.utils.supp_shape_from_ref_param_shape` can be used.

        Other variables have fixed support shape such as `support_shape=(2,)` or it is determined by the
        values (not shapes) of some parameters. For instance, a `gaussian_random_walk(steps, size=(2,))`,
        might have `support_shape=(steps,)`.
        """
        if self.signature is not None:
            # Signature could indicate fixed numerical shapes
            # As per https://numpy.org/neps/nep-0020-gufunc-signature-enhancement.html
            output_sig = self.output_sig
            core_out_shape = {
                dim: int(dim) if str.isnumeric(dim) else None for dim in self.output_sig
            }

            # Try to infer missing support dims from signature of params
            for param, param_sig, ndim_params in zip(
                dist_params, self.inputs_sig, self.ndims_params
            ):
                if ndim_params == 0:
                    continue
                for param_dim, dim in zip(param.shape[-ndim_params:], param_sig):
                    if dim in core_out_shape and core_out_shape[dim] is None:
                        core_out_shape[dim] = param_dim

            if all(dim is not None for dim in core_out_shape.values()):
                # We have all we need
                return [core_out_shape[dim] for dim in output_sig]

        raise NotImplementedError(
            "`_supp_shape_from_params` must be implemented for multivariate RVs "
            "when signature is not sufficient to infer the support shape"
        )

    def rng_fn(self, rng, *args, **kwargs) -> int | float | np.ndarray:
        """Sample a numeric random variate."""
        return getattr(rng, self.name)(*args, **kwargs)

    def __str__(self):
        # Only show signature from core props
        if signature := self.signature:
            # inp, out = signature.split("->")
            # extended_signature = f"[rng],[size],{inp}->[rng],{out}"
            # core_props = [extended_signature]
            core_props = [f'"{signature}"']
        else:
            # Far back compat
            core_props = [str(self.ndim_supp), str(self.ndims_params)]

        # Add any extra props that the subclass may have
        extra_props = [
            str(getattr(self, prop))
            for prop in self.__props__
            if prop not in RandomVariable.__props__
        ]

        props_str = ", ".join(core_props + extra_props)
        return f"{self.name}_rv{{{props_str}}}"

    def _infer_shape(
        self,
        size: TensorVariable | Variable,
        dist_params: Sequence[TensorVariable],
        param_shapes: Sequence[tuple[Variable, ...]] | None = None,
    ) -> tuple[ScalarVariable | TensorVariable, ...]:
        """Compute the output shape given the size and distribution parameters.

        Parameters
        ----------
        size
            The size parameter specified for this `RandomVariable`.
        dist_params
            The symbolic parameter for this `RandomVariable`'s distribution.
        param_shapes
            The shapes of the `dist_params` as given by `ShapeFeature`'s
            via `Op.infer_shape`'s `input_shapes` argument.  This parameter's
            values are essentially more accurate versions of ``[d.shape for d
            in dist_params]``.

        """

        from pytensor.tensor.extra_ops import broadcast_shape_iter

        if self.ndim_supp == 0:
            supp_shape = ()
        else:
            supp_shape = tuple(
                self._supp_shape_from_params(dist_params, param_shapes=param_shapes)
            )

        if not isinstance(size.type, NoneTypeT):
            size_len = get_vector_length(size)

            # Fail early when size is incompatible with parameters
            for i, (param, param_ndim_supp) in enumerate(
                zip(dist_params, self.ndims_params)
            ):
                param_batched_dims = getattr(param, "ndim", 0) - param_ndim_supp
                if param_batched_dims > size_len:
                    raise ValueError(
                        f"Size length is incompatible with batched dimensions of parameter {i} {param}:\n"
                        f"len(size) = {size_len}, len(batched dims {param}) = {param_batched_dims}. "
                        f"Size length must be 0 or >= {param_batched_dims}"
                    )

            return tuple(size) + supp_shape

        # Size was not provided, we must infer it from the shape of the parameters
        if param_shapes is None:
            param_shapes = [shape_tuple(p) for p in dist_params]

        def extract_batch_shape(p, ps, n):
            shape = tuple(ps)

            if n == 0:
                return shape

            batch_shape = tuple(
                s if not b else constant(1, "int64")
                for s, b in zip(shape[:-n], p.type.broadcastable[:-n])
            )
            return batch_shape

        # These are versions of our actual parameters with the anticipated
        # dimensions (i.e. support dimensions) removed so that only the
        # independent variate dimensions are left.
        params_batch_shape = tuple(
            extract_batch_shape(p, ps, n)
            for p, ps, n in zip(dist_params, param_shapes, self.ndims_params)
        )

        if len(params_batch_shape) == 1:
            [batch_shape] = params_batch_shape
        elif len(params_batch_shape) > 1:
            # If there are multiple parameters, the dimensions of their
            # independent variates should broadcast together.
            batch_shape = broadcast_shape_iter(
                params_batch_shape,
                arrays_are_shapes=True,
            )
        else:
            # Distribution has no parameters
            batch_shape = ()

        shape = batch_shape + supp_shape

        return shape

    def infer_shape(self, fgraph, node, input_shapes):
        _, size, *dist_params = node.inputs
        _, _, *param_shapes = input_shapes

        shape = self._infer_shape(size, dist_params, param_shapes=param_shapes)

        return [None, list(shape)]

    def __call__(self, *args, size=None, name=None, rng=None, dtype=None, **kwargs):
        if dtype is None:
            dtype = self.dtype
        if dtype == "floatX":
            dtype = config.floatX

        # We need to recreate the Op with the right dtype
        if dtype != self.dtype:
            # Check we are not switching from float to int
            if self.dtype is not None:
                if dtype.startswith("float") != self.dtype.startswith("float"):
                    raise ValueError(
                        f"Cannot change the dtype of a {self.name} RV from {self.dtype} to {dtype}"
                    )
            props = self._props_dict()
            props["dtype"] = dtype
            new_op = type(self)(**props)
            return new_op.__call__(
                *args, size=size, name=name, rng=rng, dtype=dtype, **kwargs
            )

        res = super().__call__(rng, size, *args, **kwargs)

        if name is not None:
            res.name = name

        return res

    def make_node(self, rng, size, *dist_params):
        """Create a random variable node.

        Parameters
        ----------
        rng: RandomGeneratorType
            Existing PyTensor `Generator` object to be used.  Creates a new one, if `None`.
        size: int or Sequence
            NumPy-like size parameter.
        dtype: str
            The dtype of the sampled output.  If the value ``"floatX"`` is
            given, then `dtype` is set to ``pytensor.config.floatX``.  This value is
            only used when ``self.dtype`` isn't set.
        dist_params: list
            Distribution parameters.

        Results
        -------
        out: Apply
            A node with inputs ``(rng, size, dtype) + dist_args`` and outputs
            ``(rng_var, out_var)``.

        """
        size = normalize_size_param(size)

        dist_params = tuple(
            as_tensor_variable(p) if not isinstance(p, Variable) else p
            for p in dist_params
        )

        if rng is None:
            rng = pytensor.shared(np.random.default_rng())
        elif not isinstance(rng.type, RandomType):
            raise TypeError(
                "The type of rng should be an instance of RandomGeneratorType "
            )

        inferred_shape = self._infer_shape(size, dist_params)
        _, static_shape = infer_static_shape(inferred_shape)

        dist_params = explicit_expand_dims(
            dist_params,
            self.ndims_params,
            size_length=None if NoneConst.equals(size) else get_vector_length(size),
        )

        inputs = (rng, size, *dist_params)
        out_type = TensorType(dtype=self.dtype, shape=static_shape)
        outputs = (rng.type(), out_type())

        return Apply(self, inputs, outputs)

    def batch_ndim(self, node: Apply) -> int:
        return cast(int, node.default_output().type.ndim - self.ndim_supp)

    def rng_param(self, node) -> Variable:
        """Return the node input corresponding to the rng"""
        return node.inputs[0]

    def size_param(self, node) -> Variable:
        """Return the node input corresponding to the size"""
        return node.inputs[1]

    def dist_params(self, node) -> Sequence[Variable]:
        """Return the node inpust corresponding to dist params"""
        return node.inputs[2:]

    def perform(self, node, inputs, outputs):
        rng_var_out, smpl_out = outputs

        rng, size, *args = inputs

        # Draw from `rng` if `self.inplace` is `True`, and from a copy of `rng` otherwise.
        if not self.inplace:
            rng = copy(rng)

        rng_var_out[0] = rng

        if size is not None:
            size = tuple(size)
        smpl_val = self.rng_fn(rng, *([*args, size]))

        if not isinstance(smpl_val, np.ndarray) or str(smpl_val.dtype) != self.dtype:
            smpl_val = _asarray(smpl_val, dtype=self.dtype)

        smpl_out[0] = smpl_val

    def grad(self, inputs, outputs):
        return [
            pytensor.gradient.grad_undefined(
                self, k, inp, "No gradient defined for random variables"
            )
            for k, inp in enumerate(inputs)
        ]

    def R_op(self, inputs, eval_points):
        return [None for i in eval_points]


class AbstractRNGConstructor(Op):
    def make_node(self, seed=None):
        if seed is None:
            seed = NoneConst
        else:
            seed = as_tensor_variable(seed)
        inputs = [seed]
        outputs = [self.random_type()]
        return Apply(self, inputs, outputs)

    def perform(self, node, inputs, output_storage):
        (seed,) = inputs
        if seed is not None and seed.size == 1:
            seed = int(seed)
        output_storage[0][0] = getattr(np.random, self.random_constructor)(seed=seed)


class DefaultGeneratorMakerOp(AbstractRNGConstructor):
    random_type = RandomGeneratorType()
    random_constructor = "default_rng"


default_rng = DefaultGeneratorMakerOp()


@_vectorize_node.register(RandomVariable)
def vectorize_random_variable(
    op: RandomVariable, node: Apply, rng, size, *dist_params
) -> Apply:
    # If size was provided originally and a new size hasn't been provided,
    # We extend it to accommodate the new input batch dimensions.
    # Otherwise, we assume the new size already has the right values

    original_dist_params = op.dist_params(node)
    old_size = op.size_param(node)
    len_old_size = (
        None if isinstance(old_size.type, NoneTypeT) else get_vector_length(old_size)
    )

    if len_old_size and equal_computations([old_size], [size]):
        # If the original RV had a size variable and a new one has not been provided,
        # we need to define a new size as the concatenation of the original size dimensions
        # and the novel ones implied by new broadcasted batched parameters dimensions.
        new_ndim = dist_params[0].type.ndim - original_dist_params[0].type.ndim
        if new_ndim >= 0:
            new_size = compute_batch_shape(dist_params, ndims_params=op.ndims_params)
            new_size_dims = new_size[:new_ndim]
            size = concatenate([new_size_dims, size])

    return op.make_node(rng, size, *dist_params)


class RandomVariableWithCoreShape(OpWithCoreShape):
    """Generalizes a random variable `Op` to include a core shape parameter."""

    def __str__(self):
        [rv_node] = self.fgraph.apply_nodes
        return f"[{rv_node.op!s}]"
