from pytensor import tensor as pt
from pytensor.assumptions import DIAGONAL, ORTHOGONAL, check_assumption
from pytensor.graph import Apply, FunctionGraph
from pytensor.graph.rewriting.basic import (
    node_rewriter,
)
from pytensor.graph.rewriting.unify import OpPattern
from pytensor.tensor.basic import alloc_diag
from pytensor.tensor.blockwise import Blockwise
from pytensor.tensor.elemwise import DimShuffle
from pytensor.tensor.linalg.constructors import BlockDiagonal, block_diag
from pytensor.tensor.linalg.decomposition.cholesky import Cholesky
from pytensor.tensor.linalg.inverse import MatrixInverse, MatrixPinv
from pytensor.tensor.linalg.products import KroneckerProduct, kron
from pytensor.tensor.linalg.solvers.general import solve
from pytensor.tensor.math import Dot
from pytensor.tensor.rewriting.basic import (
    register_canonicalize,
    register_specialize,
    register_stabilize,
)
from pytensor.tensor.rewriting.blockwise import blockwise_of
from pytensor.tensor.rewriting.linalg.utils import (
    ASSUME_A_OF_TRANSPOSE,
    MATRIX_INVERSE_OPS,
    get_assume_a,
)


@register_canonicalize
@node_rewriter([OpPattern(DimShuffle, is_left_expanded_matrix_transpose=True)])
def transpose_of_inv(fgraph, node):
    # TODO: Transpose is much more frequent that MatrixInverse, flip the rewrite pattern matching.
    [A] = node.inputs
    match A.owner_op_and_inputs:
        case (Blockwise(MatrixInverse()) as inv_op, X):
            return [inv_op(node.op(X))]


@register_stabilize
@node_rewriter([Dot, blockwise_of(Dot)])
def inv_to_solve(fgraph, node):
    """Replace inv(X) @ b with solve(X, b) and b @ inv(X) with solve(X.T, b.T).T."""
    l, r = node.inputs

    match l.owner_op_and_inputs:
        case (Blockwise(MatrixInverse()), X):
            return [solve(X, r, assume_a=get_assume_a(fgraph, X))]

    match r.owner_op_and_inputs:
        case (Blockwise(MatrixInverse()), X):
            assume_a = get_assume_a(fgraph, X)
            # X.mT == X for sym/pos, so reuse X and skip an unnecessary transpose.
            if assume_a in ("sym", "pos"):
                return [solve(X, l.mT, assume_a=assume_a).mT]
            return [solve(X.mT, l.mT, assume_a=ASSUME_A_OF_TRANSPOSE[assume_a]).mT]

    return None


@register_canonicalize
@register_stabilize
@node_rewriter([blockwise_of(MATRIX_INVERSE_OPS)])
def inv_of_inv(fgraph, node):
    """
    This rewrite takes advantage of the fact that if there are two consecutive inverse operations (inv(inv(input))),
    we get back our original input without having to compute inverse once.

    Here, we check for direct inverse operations (inv/pinv)  and allows for any combination of these "inverse" nodes to
    be simply rewritten.

    Parameters
    ----------
    fgraph: FunctionGraph
        Function graph being optimized
    node: Apply
        Node of the function graph to be optimized

    Returns
    -------
    list of Variable, optional
        List of optimized variables, or None if no optimization was performed
    """
    # Check if inner op is blockwise and possible inv
    match node.inputs[0].owner_op_and_inputs:
        case (Blockwise(MatrixInverse() | MatrixPinv()), X):
            return [X]


@register_canonicalize
@register_stabilize
@node_rewriter([blockwise_of(MATRIX_INVERSE_OPS)])
def inv_of_orthogonal_to_transpose(fgraph, node):
    """inv(Q) -> Q.T when Q is orthogonal."""
    [X] = node.inputs
    if not check_assumption(fgraph, X, ORTHOGONAL):
        return None
    return [X.mT]


@register_canonicalize
@register_stabilize
@node_rewriter([blockwise_of(MATRIX_INVERSE_OPS)])
def inv_of_diag_to_diag_reciprocal(fgraph, node):
    """inv(D) -> AllocDiag(1 / diagonal(D)) for diagonal D."""
    inp = node.inputs[0]

    if not check_assumption(fgraph, inp, DIAGONAL):
        return None

    diag_vals = pt.diagonal(inp, axis1=-2, axis2=-1)
    return [alloc_diag(1 / diag_vals, axis1=-2, axis2=-1)]


@register_specialize
@node_rewriter([blockwise_of(MatrixInverse | Cholesky | MatrixPinv)])
def lift_linalg_of_expanded_matrices(fgraph: FunctionGraph, node: Apply):
    """
    Rewrite compositions of linear algebra operations by lifting expensive operations (Cholesky, Inverse) through Ops
    that join matrices (KroneckerProduct, BlockDiagonal).

    This rewrite takes advantage of commutation between certain linear algebra operations to do several smaller matrix
    operations on component matrices instead of one large one. For example, when taking the inverse of Kronecker
    product, we can take the inverse of each component matrix and then take the Kronecker product of the inverses. This
    reduces the cost of the inverse from O((n*m)^3) to O(n^3 + m^3) where n and m are the dimensions of the component
    matrices.

    Parameters
    ----------
    fgraph: FunctionGraph
        Function graph being optimized
    node: Apply
        Node of the function graph to be optimized

    Returns
    -------
    list of Variable, optional
        List of optimized variables, or None if no optimization was performed
    """

    # TODO: Simplify this if we end up Blockwising KroneckerProduct
    outer_op = node.op
    [y] = node.inputs

    match y.owner_op_and_inputs:
        case (Blockwise(BlockDiagonal()), *inner_matrices):
            return [block_diag(*(outer_op(m) for m in inner_matrices))]
        case (KroneckerProduct(), *inner_matrices):  # type: ignore[unreachable, unused-ignore]
            return [kron(*(outer_op(m) for m in inner_matrices))]  # type: ignore[unreachable, unused-ignore]
