"""
This file implement specialization optimization that break the
canonization form of the graph.

Currently there is problem with the order of optimization and the
definition of definition of canonized graph.

Right now there is a canonization optimization phase that try to make
all equivalent graph identical. This is not always the case, but it do
many of the basic stuff canonical. We need to extend the definition of
canonization to make this true more often.

The problem this file indent to fix in the future is that in the
"Equilibrium" specialization optimization phase, there is optimization
that request that the graph is canonical, some other request that this
is not true, and some other that break the canonicalization for some
optimization. As we can't control the order of those optimization, there
is case that some optimization requesting a canonical graph won't be
applied as optimization that break the canonicalization form of the
graph executed before.

To fix this, we need to split the specialization phase into a phase
where optimization can't break the canonicalization form and one where
this is allowed. This is also needed for the stabilized optimization
phase, but as it happen before the specialization phase, this cause less
problem.

Also, we should make the fgraph refuse optimization that break the
canonization of the graph in the optimizations phases where the graph is
supposed to be canonical.

"""

from pytensor import scalar as ps
from pytensor.graph.rewriting.basic import copy_stack_trace, node_rewriter
from pytensor.tensor.basic import Alloc, alloc, constant
from pytensor.tensor.elemwise import CAReduce, DimShuffle
from pytensor.tensor.math import Min, neg
from pytensor.tensor.rewriting.basic import register_uncanonicalize
from pytensor.tensor.shape import Reshape, reshape
from pytensor.tensor.subtensor import Subtensor, indices_from_subtensor


@register_uncanonicalize
@node_rewriter([neg])
def local_max_to_min(fgraph, node):
    """
    Change -(max(-x)) to min.

    This is tested in tensor/tests/test_basic.py:test_min_max.

    Notes
    -----
    We don't need an opt that will do the reverse as by default
    the interface put only Max into the graph.

    """
    if node.op == neg and node.inputs[0].owner:
        max = node.inputs[0]
        if (
            max.owner
            and isinstance(max.owner.op, CAReduce)
            and max.owner.op.scalar_op == ps.maximum
        ):
            neg_node = max.owner.inputs[0]
            if neg_node.owner and neg_node.owner.op == neg:
                new = Min(max.owner.op.axis)(neg_node.owner.inputs[0])
                return [copy_stack_trace(node.outputs[0], new)]

    return False


@register_uncanonicalize
@node_rewriter([Alloc])
def local_alloc_dimshuffle(fgraph, node):
    """
    If a dimshuffle is inside an alloc and only adds dimension to the
    left, remove it.

    Alloc(DimShuffle(x), ...) - > Alloc(x, ...)
    """
    input_ = node.inputs[0]
    if input_.owner and isinstance(input_.owner.op, DimShuffle):
        # check if it only adds dimension to the left
        new_order = input_.owner.op.new_order
        expected_new_order = ("x",) * (
            input_.ndim - input_.owner.inputs[0].ndim
        ) + tuple(range(input_.owner.inputs[0].ndim))
        if new_order != expected_new_order:
            return False
        return [alloc(input_.owner.inputs[0], *node.inputs[1:])]
    return False


@register_uncanonicalize
@node_rewriter([Reshape])
def local_reshape_dimshuffle(fgraph, node):
    """
    If a dimshuffle is inside a reshape and does not change the order
    of dimensions, remove it.

    Reshape(Dimshuffle(x), shp) -> Reshape(x, shp)
    """
    input_ = node.inputs[0]
    if input_.owner and isinstance(input_.owner.op, DimShuffle):
        new_order = input_.owner.op.new_order
        offset = 0
        for dim in new_order:
            if dim == "x":
                continue
            elif dim != offset:
                return False
            else:
                offset += 1
        return [
            reshape(input_.owner.inputs[0], node.inputs[1], ndim=node.outputs[0].ndim)
        ]
    return False


@register_uncanonicalize
@node_rewriter([DimShuffle])
def local_dimshuffle_alloc(fgraph, node):
    """
    If an alloc is inside a dimshuffle which only adds dimension to the left,
    scrap the dimshuffle and adds 1 into the alloc

    dimshuffle{x, 0, 1}(alloc([3 4], 3, 2) => alloc([3 4], 1, 3, 2)
    """
    if not node.inputs[0].owner:
        return False
    input_ = node.inputs[0]
    if isinstance(input_.owner.op, Alloc):
        # check if it only adds dimension to the left
        new_order = node.op.new_order
        expected_new_order = ("x",) * (len(new_order) - input_.ndim) + tuple(
            range(input_.ndim)
        )
        if new_order != expected_new_order:
            return False

        # count numbers of 'x'
        nb_new_dims = len(new_order) - input_.ndim
        new_shape_input = (1,) * nb_new_dims + tuple(input_.owner.inputs[1:])

        return [alloc(input_.owner.inputs[0], *new_shape_input)]
    return False


@register_uncanonicalize
@node_rewriter([DimShuffle])
def local_dimshuffle_subtensor(fgraph, node):
    """If a subtensor is inside a dimshuffle which only drop
    broadcastable dimensions, scrap the dimshuffle and index the
    subtensor in a way that avoids the degenerate dimension

    x[i:j, :, k:l].dimshuffle(0, 2) =>
        x[i:j, 0, k:l] if x.broadcastable == (False, True, False)

    x[i:j, k:l, :].dimshuffle(0, 2) => x[i:j, k, :]
    x[i:j, k:, :].dimshuffle(0, 2) => x[i:j, k, :]
    x[i:j, :l, :].dimshuffle(0, 2) => x[i:j, 0, :]

    """
    if not node.inputs[0].owner:
        return False
    # the dimshuffle can only drop dimensions (cannot reshape nor add 'x')
    if "x" in node.op.new_order:
        return False
    new_order = node.op.new_order
    # new order could be empty
    # Verif that we don't change dimensions order.
    if len(new_order) > 1:
        past_dim = new_order[0]
        for dim in new_order[1:]:
            if not dim > past_dim:
                return False
            else:
                past_dim = dim

    input_ = node.inputs[0]
    if isinstance(input_.owner.op, Subtensor):
        # the arguments missing from the dimshuffles must be dims
        # that are broadcastable
        broadcastable = input_.broadcastable

        missing_dims = list(range(input_.ndim))
        for dim in new_order:
            missing_dims.remove(dim)

        if not all(broadcastable[i] for i in missing_dims):
            return False

        # create a new index tuple for a new Subtensor
        # Reconstruct the full indices from the subtensor node, then replace
        # dimensions that are being dropped by dimshuffle with scalar index 0
        x = input_.owner.inputs[0]
        indices = list(
            indices_from_subtensor(input_.owner.inputs[1:], input_.owner.op.idx_list)
        )
        zero = constant(0)

        # Track which output dimension each index corresponds to
        # Scalar indices remove dimensions, slices keep them
        output_dim = 0
        for i, idx in enumerate(indices):
            if isinstance(idx, slice):
                # This slice produces an output dimension
                if output_dim in missing_dims:
                    # This output dimension is being dropped, so replace slice with scalar
                    if idx == slice(None):
                        indices[i] = zero
                    else:
                        # Use the start of the slice (or 0 if None)
                        indices[i] = idx.start if idx.start is not None else zero
                output_dim += 1
            # Scalar indices don't contribute to output dimensions

        # Handle trailing dimensions that weren't explicitly indexed
        for input_dim in range(len(indices), x.ndim):
            if output_dim in missing_dims:
                # This unindexed dimension is being dropped, index with 0
                indices.append(zero)
            else:
                # This unindexed dimension is kept, index with slice(None)
                indices.append(slice(None))
            output_dim += 1

        return [x[tuple(indices)]]
    return False
