from pytensor.compile.mode import optdb
from pytensor.graph import Constant, node_rewriter
from pytensor.graph.replace import vectorize_node
from pytensor.graph.rewriting.basic import copy_stack_trace, out2in
from pytensor.tensor.basic import Alloc, ARange, alloc, shape_padleft
from pytensor.tensor.blockwise import Blockwise
from pytensor.tensor.math import Dot
from pytensor.tensor.rewriting.basic import (
    register_canonicalize,
    register_specialize,
    register_stabilize,
)
from pytensor.tensor.shape import Reshape
from pytensor.tensor.subtensor import AdvancedIncSubtensor, AdvancedSubtensor, Subtensor


@node_rewriter([Blockwise])
def local_useless_blockwise(fgraph, node):
    """
    If there is a dispatch implementation that does not require Blockwise, use that instead.
    This means a user created a Blockwise manually when there was no need.

    Note: This rewrite is not registered by default anywhere
    """
    op = node.op
    inputs = node.inputs
    dummy_core_node = op._create_dummy_core_node(node.inputs)
    vect_node = vectorize_node(dummy_core_node, *inputs)
    if not isinstance(vect_node.op, Blockwise):
        return copy_stack_trace(node.outputs, vect_node.outputs)


@node_rewriter([Blockwise])
def local_useless_unbatched_blockwise(fgraph, node):
    """Remove Blockwise that don't have any batched dims."""
    op = node.op
    inputs = node.inputs

    batch_ndims = node.op.batch_ndim(node)
    if all(all(inp.type.broadcastable[:batch_ndims]) for inp in inputs):
        if batch_ndims:
            # Remove dummy batch dims
            axis = tuple(range(batch_ndims))
            inputs = [inp.squeeze(axis) for inp in inputs]
        new_outs = op.core_op.make_node(*inputs).outputs
        if batch_ndims:
            # Reintroduce dummy batch dims
            new_outs = [shape_padleft(out, batch_ndims) for out in new_outs]
        return copy_stack_trace(node.outputs, new_outs)


# We register this rewrite late, so that other rewrites need only target Blockwise Ops
optdb.register(
    "local_useless_unbatched_blockwise",
    out2in(local_useless_unbatched_blockwise, ignore_newtrees=True),
    "fast_run",
    "fast_compile",
    "blockwise",
    position=49,
)


# Avoid redundant cases early on for Ops whose default form is not Blockwised
@register_canonicalize
@register_stabilize
@register_specialize
@node_rewriter(tracks=[Blockwise])
def local_eager_useless_unbatched_blockwise(fgraph, node):
    if isinstance(
        node.op.core_op,
        Dot
        | Alloc
        | ARange
        | Subtensor
        | AdvancedSubtensor
        | AdvancedIncSubtensor
        | Reshape,
    ):
        # Many Dot-related rewrites (eg, all of BlasOpt) happen before specialize
        # These other Ops can't always be trivially vectorized at runtime,
        # since their inputs may imply non-rectangular shapes.
        return local_useless_unbatched_blockwise.fn(fgraph, node)


def _squeeze_left(x, stop_at_dim: int | None = None):
    """Squeeze any leading dims of `x` until a real dim or `stop_at_dim` (if not None) is reached."""
    x_dims = x.type.broadcastable
    squeeze_ndim = len(x_dims) if all(x_dims) else x_dims.index(False)
    if stop_at_dim is not None:
        squeeze_ndim = min(squeeze_ndim, stop_at_dim)
    if squeeze_ndim == 0:
        return x
    return x.squeeze(axis=tuple(range(squeeze_ndim)))


@register_specialize("shape_unsafe")
@node_rewriter([Blockwise])
def local_blockwise_alloc(fgraph, node):
    """Push Allocs from the inputs to the output of Blockwise Ops.

    BOp = Blockwise(Op, signature="(x),(x)->(x)")
    BOp(vector, alloc(vector, 10, 5)) -> alloc(BOp)(vector, vector), 10, 5)
    BOp(vector, alloc(scalar, 10, 5)) -> alloc(BOp)(vector, alloc(scalar, 5), 10, 5)
    BOp(matrix, alloc(vector, 10, 5)) -> BOp(matrix, vector)
    """

    op: Blockwise = node.op  # type: ignore

    batch_ndim = op.batch_ndim(node)
    if not batch_ndim:
        return None

    if not any(var.owner and isinstance(var.owner.op, Alloc) for var in node.inputs):
        return None

    new_inputs = []
    batch_shapes = []
    can_push_any_alloc = False
    for inp, inp_sig in zip(node.inputs, op.inputs_sig):
        if not all(inp.type.broadcastable[:batch_ndim]):
            if inp.owner and isinstance(inp.owner.op, Alloc):
                # Push batch dims from Alloc
                value, *shape = inp.owner.inputs

                # Check what to do with the value of the Alloc
                squeezed_value = _squeeze_left(value, batch_ndim)
                missing_ndim = len(shape) - value.type.ndim
                if (
                    (((1,) * missing_ndim + value.type.broadcastable)[batch_ndim:])
                    != inp.type.broadcastable[batch_ndim:]
                ):
                    # We still need an Alloc for the core dims
                    core_shape = shape[batch_ndim:]
                    # And the batch dims of the squeezed value
                    squeezed_value_batch_ndim = squeezed_value.type.ndim - len(
                        core_shape
                    )
                    batch_shape = [
                        1 if broadcastable else dim
                        for broadcastable, dim in zip(
                            squeezed_value.type.broadcastable[
                                :squeezed_value_batch_ndim
                            ],
                            tuple(squeezed_value.shape)[:squeezed_value_batch_ndim],
                        )
                    ]
                    squeezed_value = alloc(squeezed_value, *batch_shape, *core_shape)
                    if squeezed_value.type.broadcastable == inp.type.broadcastable:
                        # We can't change anything about this Alloc input
                        new_inputs.append(inp)
                        continue

                # We can push batch dims of this Alloc input
                batch_shapes.append(
                    tuple(
                        1 if broadcastable else dim
                        for broadcastable, dim in zip(
                            inp.type.broadcastable, shape[:batch_ndim]
                        )
                    )
                )
                new_inputs.append(squeezed_value)
                can_push_any_alloc = True
                continue

        # Nothing to do with this input other than removing dummy batch dims
        new_inputs.append(_squeeze_left(inp, batch_ndim))

    if not can_push_any_alloc:
        return None

    new_outs = node.op.make_node(*new_inputs).outputs

    new_out_type = new_outs[0].type
    old_out_type = node.outputs[0].type
    if new_out_type.broadcastable != old_out_type.broadcastable:
        # An Alloc is still needed to broadcast the new output to the original shape
        # We pick the most parsimonious batch dim from the pushed Alloc
        missing_ndim = old_out_type.ndim - new_out_type.ndim
        batch_shape = ([1] * missing_ndim + list(new_outs[0].shape))[:batch_ndim]
        for i, batch_dims in enumerate(zip(*batch_shapes)):  # Transpose shape tuples
            if old_out_type.broadcastable[i]:
                continue
            for batch_dim in batch_dims:
                if batch_dim == 1:
                    continue
                batch_shape[i] = batch_dim
                if isinstance(batch_dim, Constant):
                    # Give preference to Constants
                    break

        copy_stack_trace(node.outputs, new_outs)
        new_outs = [
            alloc(
                new_out,
                *batch_shape,
                *new_out.shape[batch_ndim - missing_ndim :],
            )
            for new_out in new_outs
        ]
    copy_stack_trace(node.outputs, new_outs)
    return new_outs


@register_specialize
@node_rewriter([Blockwise])
def local_blockwise_reshape(fgraph, node):
    """Rewrite away square Blockwise reshapes.

    Reshape is tricky to vectorize eagerly, because a graph like
    `x.reshape([x.shape[0] * x.shape[1], -1])` has many operations
    that must be vectorized before we arrize at the reshape operation.

    For the square Reshape case, we must wait for all the intemediate
    operations to be lifted as Allocs
    """
    if not isinstance(node.op.core_op, Reshape):
        return None

    x, output_shape = node.inputs
    batch_ndim = node.op.batch_ndim(node)
    if all(output_shape.type.broadcastable[:batch_ndim]):
        batched_shape = x.shape[:batch_ndim]
        core_reshape = _squeeze_left(output_shape, batch_ndim)
        new_out = x.reshape([*tuple(batched_shape), *tuple(core_reshape)])
        copy_stack_trace(node.outputs[0], new_out)
        return [new_out]
