import logging
import operator
from typing import Any

import torch
import torch.fx
from torch.fx.experimental.const_fold import split_const_subgraphs
from torch.fx.passes.infra.pass_base import PassResult

_LOGGER = logging.getLogger(__name__)

# Create an alias for module input type to avoid littering pyre-ignore for Any
# throughout the file.
Input = Any


def run_const_fold(traced_mod: torch.fx.GraphModule) -> torch.fx.GraphModule:
    # Now we do constant folding on traced module.
    def skip_folding(node: torch.fx.Node):
        if node.target == torch.ops.aten.sym_size:
            return True

    const_split_mod = split_const_subgraphs(
        traced_mod, skip_folding_node_fn=skip_folding
    )
    const_split_mod.run_folding()
    return const_split_mod


def replace_inplace_ops(
    module: torch.fx.GraphModule,
) -> torch.fx.GraphModule:
    """
    Remove this func after functionalization is workable
    """
    modified = False
    map_func = {
        torch.ops.aten.relu_.default: torch.ops.aten.relu.default,
        torch.ops.aten.hardtanh_.default: torch.ops.aten.hardtanh.default,
        torch.ops.aten.add_.Tensor: torch.ops.aten.add.Tensor,
    }
    for n in module.graph.nodes:
        if n.op == "call_function" and n.target in map_func.keys():
            modified = True
            node = n
            with module.graph.inserting_after(node):
                new_args = node.args
                new_node = module.graph.create_node(
                    "call_function",
                    map_func[node.target],
                    args=new_args,
                    kwargs=None,
                )
                node.replace_all_uses_with(new_node)
                module.graph.erase_node(node)
    module.graph.eliminate_dead_code()
    module.recompile()
    return PassResult(module, modified)


def replace_native_layernorm_with_layernorm(
    module: torch.fx.GraphModule,
) -> torch.fx.GraphModule:
    modified = False
    for n in module.graph.nodes:
        if (
            n.op == "call_function"
            and n.target == torch.ops.aten.native_layer_norm.default
        ):
            for v in n.users:
                if v.op == "call_function" and v.target == operator.getitem:
                    if v.args[1] != 0:
                        raise RuntimeError(
                            f"Got args[{v.args[1]}]!!\n"
                            "layernorm can only generate output (args[0]), "
                            "not mean (args[1]) or std (args[2])!"
                        )
                    new_op = torch.ops.aten.layer_norm.default
                    new_args = (*n.args, True)  # cudnn_enable=True
                    modified = True
                else:
                    continue

                with module.graph.inserting_after(v):
                    new_node = module.graph.create_node(
                        "call_function",
                        new_op,
                        args=new_args,
                        kwargs=v.kwargs,
                    )
                    v.replace_all_uses_with(new_node)

    module.graph.eliminate_dead_code()
    module.recompile()
    return PassResult(module, modified)


def replace_transpose_mm_op_with_linear(
    module: torch.fx.GraphModule,
) -> torch.fx.GraphModule:
    modified = False
    for n in module.graph.nodes:
        if n.op == "call_function" and n.target == torch.ops.aten.t.default:
            to_erase = []
            for v in n.users:
                if v.op == "call_function" and v.target == torch.ops.aten.addmm.default:
                    new_op = torch.ops.aten.linear
                    bias, inp, _ = list(v.args)
                    weight = list(n.args)[0]
                    new_args = (inp, weight, bias)
                    modified = True
                elif v.op == "call_function" and v.target == torch.ops.aten.mm.default:
                    new_op = torch.ops.aten.linear
                    inp, _ = list(v.args)
                    weight = list(n.args)[0]
                    new_args = (inp, weight, None)
                    modified = True
                # this pass should be after `compose_bmm`
                elif v.op == "call_function" and v.target == aten_compose_bmm_2d:
                    new_op = torch.ops.aten.linear
                    inp, _ = list(v.args)
                    weight = list(n.args)[0]
                    new_args = (inp, weight, None)
                    modified = True
                else:
                    continue

                with module.graph.inserting_after(v):
                    new_node = module.graph.create_node(
                        "call_function",
                        new_op,
                        args=new_args,
                        kwargs=v.kwargs,
                    )
                    v.replace_all_uses_with(new_node)
                    to_erase.append(v)
            for v in to_erase:
                module.graph.erase_node(v)
    module.graph.eliminate_dead_code()
    module.recompile()
    # handle the linear with multiple dim, remove the extra reshape
    for n in module.graph.nodes:
        if n.op == "call_function" and n.target == torch.ops.aten.linear:
            before = n.args[0]
            after = next(iter(n.users))
            if (len(n.users) == 1 and after.target == torch.ops.aten.view.default) and (
                before.target == torch.ops.aten.view.default and len(before.users) == 1
            ):
                real_input = before.args[0]
                new_args = list(n.args)
                new_args[0] = real_input
                n.args = tuple(new_args)
                after.replace_all_uses_with(n)
                module.graph.eliminate_dead_code()
                module.recompile()

    return PassResult(module, modified)


def replace_aten_op_with_indices(module: torch.fx.GraphModule) -> torch.fx.GraphModule:
    modified = False
    for n in module.graph.nodes:
        if n.op == "call_function" and n.target in (
            torch.ops.aten.max_pool2d_with_indices.default,
            torch.ops.aten.max_pool3d_with_indices.default,
            torch.ops.aten.native_batch_norm.default,
            torch.ops.aten._native_batch_norm_legit.default,
            torch.ops.aten._native_batch_norm_legit_no_training.default,
        ):
            modified = True
            if len(n.users) != 1:
                raise RuntimeError(
                    f"{n.target} has users={len(n.users)}. We can only handle it with 1 user"
                )
            if n.target == torch.ops.aten.max_pool2d_with_indices.default:
                new_op = torch.ops.aten.max_pool2d
                new_args = n.args
            elif n.target == torch.ops.aten.max_pool3d_with_indices.default:
                new_op = torch.ops.aten.max_pool3d
                new_args = n.args
            elif (
                n.target == torch.ops.aten.native_batch_norm.default
                or n.target == torch.ops.aten._native_batch_norm_legit.default
            ):
                new_op = torch.ops.aten.batch_norm
                new_args = list(n.args)
                new_args.append(False)
                new_args = tuple(new_args)
            elif (
                n.target == torch.ops.aten._native_batch_norm_legit_no_training.default
            ):
                new_op = torch.ops.aten.batch_norm
                new_args = list(n.args)
                new_args.append(False)
                # _native_batch_norm_legit_no_training doesn't take in a training arg (assumed to be false)
                # but batchnorm takes in a training arg at position 5.
                new_args.insert(5, False)
                new_args = tuple(new_args)

            getitem_node = next(iter(n.users))
            with module.graph.inserting_after(getitem_node):
                new_node = module.graph.create_node(
                    "call_function",
                    new_op,
                    args=new_args,
                    kwargs=n.kwargs,
                )
                getitem_node.replace_all_uses_with(new_node)
                module.graph.erase_node(getitem_node)
    module.graph.eliminate_dead_code()
    module.recompile()
    return PassResult(module, modified)


def replace_aten_reshape_alias_with_replace(
    module: torch.fx.GraphModule,
) -> torch.fx.GraphModule:
    # The stride parameter is not used. Replace with reshape without stride
    modified = False
    for n in module.graph.nodes:
        if n.op == "call_function" and n.target in (
            torch.ops.aten._reshape_alias.default,
        ):
            modified = True
            node = n
            with module.graph.inserting_after(node):
                new_args = (node.args[0], node.args[1])
                new_node = module.graph.create_node(
                    "call_function",
                    torch.ops.aten.reshape,
                    args=new_args,
                    kwargs=None,
                )
                node.replace_all_uses_with(new_node)
                module.graph.erase_node(node)
            break
    module.graph.eliminate_dead_code()
    module.recompile()
    return PassResult(module, modified)


def remove_ops(
    module: torch.fx.GraphModule,
) -> torch.fx.GraphModule:
    """
    1. Remove clone, _unsafe_view node. #TODO Remove this func after functionalization is workable
    2. Remove inefficient op getitem(index=slice) P561572458
    """
    modified = False
    for n in module.graph.nodes:
        if n.op == "call_function" and n.target in (torch.ops.aten.clone.default,):
            modified = True
            node = n
            input_n = node.all_input_nodes[0]
            node.replace_all_uses_with(input_n)
    module.graph.eliminate_dead_code()
    module.recompile()
    for n in module.graph.nodes:
        if n.op == "call_function" and n.target in (
            torch.ops.aten._unsafe_view.default,
        ):
            modified = True
            node = n
            with module.graph.inserting_after(node):
                new_node = module.graph.create_node(
                    "call_function",
                    torch.ops.aten.reshape,
                    args=node.args,
                    kwargs=node.kwargs,
                )
                node.replace_all_uses_with(new_node)
                module.graph.erase_node(node)
    module.graph.eliminate_dead_code()
    module.recompile()
    return PassResult(module, modified)


def aten_operator_getitem(*args):
    return operator.getitem(*args)


def replace_builtin_ops(
    module: torch.fx.GraphModule,
) -> torch.fx.GraphModule:
    """
    To differential the same op in fx2ait as they are registered in the same dictionary
    """

    modified = False
    for n in module.graph.nodes:
        if n.op == "call_function" and n.target in (operator.getitem,):
            modified = True
            n.target = aten_operator_getitem
    module.graph.eliminate_dead_code()
    module.recompile()

    module.graph.eliminate_dead_code()
    module.recompile()
    return PassResult(module, modified)


###############
"""
Trace compose. For some ops, we do not want to decompose further but want coarse granularity
For ex:
1. bmm
2. chunk
3. getitem(input, idx=(slice(),slice()...))
"""


def aten_compose_getitem_slice(input, list_args):
    for _, args in enumerate(list_args):
        input = torch.ops.aten.slice.Tensor(input, *args)
    return input


def compose_getitem_slice(
    module: torch.fx.GraphModule,
) -> torch.fx.GraphModule:
    """
    combine decomposed getitem(input, idx=(slice(),slice()...))
    """

    def match_pattern(module, node):
        if node.op == "call_function" and node.target == torch.ops.aten.slice.Tensor:
            holder = []
            holder.append(node)
            while (
                len(node.users.keys()) == 1
                and next(iter(node.users)).target == torch.ops.aten.slice.Tensor
                and node.args[1] + 1 == next(iter(node.users)).args[1]
            ):
                node = next(iter(node.users))
                holder.append(node)
            if len(holder) == 1:
                return (False,)
            else:
                return (True, holder)
        return (False,)

    modified = False
    for node in module.graph.nodes:
        res = match_pattern(module, node)
        if res[0]:
            modified = True
            holder = res[1]
            input_n = holder[0].args[0]
            last_n = holder[-1]
            list_args = []
            for h_n in holder:
                list_args.append(h_n.args[1:])

            with module.graph.inserting_after(last_n):
                new_args = (input_n, list_args)
                new_node = module.graph.create_node(
                    "call_function",
                    aten_compose_getitem_slice,
                    args=new_args,
                    kwargs=None,
                )
            last_n.replace_all_uses_with(new_node)
    module.graph.eliminate_dead_code()
    module.recompile()
    return PassResult(module, modified)


def aten_compose_bmm_2d(flat_args_1, flat_args_2):
    sym_size = torch.ops.aten.sym_size(flat_args_1, 0)
    sym_size_1 = torch.ops.aten.sym_size(flat_args_1, 1)
    sym_size_2 = torch.ops.aten.sym_size(flat_args_1, 2)
    expand = torch.ops.aten.expand.default(
        flat_args_1, [sym_size, sym_size_1, sym_size_2]
    )
    view = torch.ops.aten.view.default(expand, [sym_size, sym_size_1, sym_size_2])
    sym_size_3 = torch.ops.aten.sym_size(flat_args_2, 0)
    sym_size_4 = torch.ops.aten.sym_size(flat_args_2, 1)
    expand_1 = torch.ops.aten.expand.default(
        flat_args_2, [sym_size, sym_size_3, sym_size_4]
    )
    view_1 = torch.ops.aten.view.default(expand_1, [sym_size, sym_size_3, sym_size_4])
    bmm = torch.ops.aten.bmm.default(view, view_1)
    view_2 = torch.ops.aten.view.default(bmm, [sym_size, sym_size_1, sym_size_4])
    return view_2


def aten_compose_bmm_3d(flat_args_1, flat_args_2):
    sym_size = torch.ops.aten.sym_size(flat_args_1, 0)
    sym_size_1 = torch.ops.aten.sym_size(flat_args_1, 1)
    sym_size_2 = torch.ops.aten.sym_size(flat_args_1, 2)
    expand = torch.ops.aten.expand.default(
        flat_args_1, [sym_size, sym_size_1, sym_size_2]
    )
    view = torch.ops.aten.view.default(expand, [sym_size, sym_size_1, sym_size_2])
    sym_size_3 = torch.ops.aten.sym_size(flat_args_2, 1)
    sym_size_4 = torch.ops.aten.sym_size(flat_args_2, 2)
    expand_1 = torch.ops.aten.expand.default(
        flat_args_2, [sym_size, sym_size_3, sym_size_4]
    )
    view_1 = torch.ops.aten.view.default(expand_1, [sym_size, sym_size_3, sym_size_4])
    bmm = torch.ops.aten.bmm.default(view, view_1)
    view_2 = torch.ops.aten.view.default(bmm, [sym_size, sym_size_1, sym_size_4])
    return view_2


def compose_bmm(
    module: torch.fx.GraphModule,
) -> torch.fx.GraphModule:
    """
    combine decomposed bmm (matmul)
    """
    modified = False
    for n in module.graph.nodes:
        if n.op == "call_function" and n.target in (torch.ops.aten.bmm.default,):
            modified = True
            node = n
            input_n = node.all_input_nodes[0]
            other_n = node.all_input_nodes[1]
            output = next(iter(node.users))
            input_input_n = input_n.all_input_nodes[0]
            if (
                input_input_n.target != torch.ops.aten.expand.default
                and input_n.target != torch.ops.aten.view.default
            ):
                raise RuntimeError(
                    "Bmm is addressed in fixed pattern. A new pattern is met!"
                )
            real_input = input_input_n.all_input_nodes[0]
            input_other_n = other_n.all_input_nodes[0]
            if (
                input_other_n.target != torch.ops.aten.expand.default
                and other_n.target != torch.ops.aten.view.default
            ):
                raise RuntimeError(
                    "Bmm is addressed in fixed pattern. A new pattern is met!"
                )
            real_other = input_other_n.all_input_nodes[0]
            if len(real_other.meta["val"].size()) == 2:
                new_func = aten_compose_bmm_2d
            if len(real_other.meta["val"].size()) == 3:
                new_func = aten_compose_bmm_3d

            with module.graph.inserting_after(node):
                new_args = (real_input, real_other)
                new_node = module.graph.create_node(
                    "call_function",
                    new_func,
                    args=new_args,
                    kwargs=None,
                )
            output.replace_all_uses_with(new_node)

    module.graph.eliminate_dead_code()
    module.recompile()
    return PassResult(module, modified)


def aten_compose_chunk(flat_args_1, chunk, dim):
    sym_size = torch.ops.aten.sym_size(flat_args_1, dim)
    add = operator.add(sym_size, chunk)
    sub = operator.sub(add, 1)
    floordiv = operator.floordiv(sub, chunk)
    split = torch.ops.aten.split.Tensor(flat_args_1, floordiv, dim)
    return split


def compose_chunk(
    module: torch.fx.GraphModule,
) -> torch.fx.GraphModule:
    """
    combine decomposed chunk
    """

    def match_pattern(module, node):
        if node.op == "call_function" and node.target in (torch.ops.aten.split.Tensor,):
            div = node.args[1]
            input = node.args[0]
            if isinstance(div, int):
                return (False,)
            if div.target != operator.floordiv:
                return (False,)
            else:
                div_const = div.args[1]
                sub = div.args[0]
                if sub.target != operator.sub:
                    return (False,)
                else:
                    add = sub.args[0]
                    if add.target != operator.add:
                        return (False,)
                    else:
                        add_const = add.args[1]
                        if add_const != div_const:
                            return (False,)
                        symsize = add.args[0]
                        if symsize.target != torch.ops.aten.sym_size:
                            return (False,)
                        else:
                            symsize_input = symsize.args[0]
                            dim = symsize.args[1]
                            if symsize_input != input:
                                return (False,)

            return (True, div_const, dim)
        else:
            return (False,)

    modified = False
    for node in module.graph.nodes:
        res = match_pattern(module, node)
        if res[0]:
            modified = True
            with module.graph.inserting_after(node):
                new_args = (node.args[0], res[1], res[2])
                new_node = module.graph.create_node(
                    "call_function",
                    aten_compose_chunk,
                    args=new_args,
                    kwargs=None,
                )
            node.replace_all_uses_with(new_node)

    module.graph.eliminate_dead_code()
    module.recompile()
    return PassResult(module, modified)