# mypy: allow-untyped-defs
"""CPU-specific implementations for flex attention"""

import copy
import os
import sys
from typing import Any

import sympy

import torch
from torch._inductor.virtualized import V
from torch.utils._ordered_set import OrderedSet
from torch.utils._sympy.numbers import int_oo
from torch.utils._sympy.value_ranges import ValueRanges

from ...codegen.cpp_flex_attention_template import CppFlexAttentionTemplate
from ...ir import Buffer, FixedLayout, TensorBox
from ...select_algorithm import autotune_select_algorithm
from .common import (
    build_subgraph_buffer,
    build_subgraph_module_buffer,
    contiguous_last_dim,
    create_placeholder,
    get_fwd_subgraph_outputs,
    infer_dense_strides,
    maybe_realize,
)


def check_cpu_supported():
    requires_avx2_on_cpu = (
        torch.cpu._is_avx2_supported() and os.getenv("ATEN_CPU_CAPABILITY") != "default"
    )
    supported = (
        requires_avx2_on_cpu
        and not torch.xpu.is_available()
        and sys.platform != "darwin"
    )
    return supported


def lower_cpu(
    query,
    key,
    value,
    subgraph,
    block_mask,
    scale,
    kernel_options,
    score_mod_other_buffers,
    mask_mod_other_buffers,
):
    """CPP based template for flex attention for x86 CPUs"""
    (
        _,  # q_length
        _,  # kv_length
        kv_num_blocks,
        kv_indices,
        full_kv_num_blocks,
        full_kv_indices,
        q_num_blocks,
        q_indices,
        full_q_num_blocks,
        full_q_indices,
        SPARSE_Q_BLOCK_SIZE,
        SPARSE_KV_BLOCK_SIZE,
        mask_graph,
    ) = block_mask

    if query.dtype != key.dtype or query.dtype != value.dtype:
        raise ValueError(
            f"Mixed query, key, and value dtype is not supported on this platform, "
            f"got query.dtype: {query.dtype}, key.dtype: {key.dtype}, "
            f"and value.dtype: {value.dtype}."
        )

    if kernel_options["OUTPUT_LOGSUMEXP"]:
        raise NotImplementedError(
            "torch.compile on CPU only supports inference and `return_lse` is not supported yet."
        )
    if not check_cpu_supported():
        raise NotImplementedError(
            "torch.compile on current platform is not supported for CPU."
        )

    fake_buffers: list[Buffer] = []  # noqa: F821

    # [Note] Handle the case where the split sizes are not statically known.
    # The value of cur_qSplitSize and cur_kvSplitSize are decided during runtime.
    # We use symbols to represent them during the compilation here.
    # They'll be replaced by the string "cur_qSplitSize" and "cur_kvSplitSize" in
    # the modification function of the CppFlexAttentionTemplate class.
    cur_qSplitSize = V.graph.sizevars.shape_env.create_unbacked_symint().node.expr
    cur_kvSplitSize = V.graph.sizevars.shape_env.create_unbacked_symint().node.expr
    shape_env = V.graph.sizevars.shape_env

    # We don't know the concrete value of cur_qSplitSize and cur_kvSplitSize during the compilation.
    # Mark symbols > 1 to ensure broadcasting is always applied.
    # This avoids treating them as equal when `eq(var, 1)` is evaluated in `broadcast_symbolic_shapes`.
    shape_env.var_to_range[cur_qSplitSize] = ValueRanges(2, int_oo)
    shape_env.var_to_range[cur_kvSplitSize] = ValueRanges(2, int_oo)

    score_dtype = torch.float
    placeholder_inps = [
        create_placeholder(name, dtype, query.get_device(), size)
        for name, dtype, size in [
            ("score", score_dtype, [cur_qSplitSize, cur_kvSplitSize]),
            ("b", torch.int64, []),
            ("h", torch.int64, []),
            ("q_idx", torch.int64, [cur_qSplitSize, 1]),
            ("kv_idx", torch.int64, [1, cur_kvSplitSize]),
        ]
    ]
    subgraph_buffer = build_subgraph_buffer(
        placeholder_inps + list(score_mod_other_buffers), subgraph
    )
    if subgraph_buffer is not None:
        if isinstance(subgraph_buffer, list):
            for _buf in subgraph_buffer:
                if _buf is not None:
                    _buf.freeze_layout()
        else:
            subgraph_buffer.freeze_layout()
    mask_graph_placeholder_inps = [
        create_placeholder(name, dtype, query.get_device(), size)
        for name, dtype, size in [
            ("score", score_dtype, [cur_qSplitSize, cur_kvSplitSize]),
            ("b", torch.int64, []),
            ("h", torch.int64, []),
            ("q_idx", torch.int64, [cur_qSplitSize, 1]),
            ("kv_idx", torch.int64, [1, cur_kvSplitSize]),
        ]
    ]

    # The original mask_graph works on a scalar and only includes
    # the logic of calculating the mask value.
    # We need to add the logic of applying the mark to the qk_data tensor
    # into the graph for the later codegen of this part.
    # Example:
    #   mask_graph:
    #   def mask_fn(b, h, q_idx, kv_idx):
    #       mask = q_idx >= kv_idx
    #       return mask
    #   The converted_mask_graph should be:
    #   def converted_mask_fn(qk_data, b, h, q_idx, kv_idx):
    #       mask = q_idx >= kv_idx
    #       qk_data = torch.where(mask, qk_data, torch.full_like(qk_data, -float("inf")))
    #       return qk_data
    def convert_mask_graph_module(mask_graph):
        gm = copy.deepcopy(mask_graph.graph_module)
        graph = gm.graph
        # Add qk_data as the first input
        with graph.inserting_before(next(iter(graph.nodes))):
            qk_data_node = graph.placeholder("qk_data")

        # Find the node that returns the mask
        output_node = None
        for node in graph.nodes:
            if node.op == "output":
                output_node = node
                break

        # Get the mask node
        assert output_node is not None
        mask_node = output_node.args[0]

        size_node = [cur_qSplitSize, cur_kvSplitSize]
        # Create a new node for torch.full
        with graph.inserting_after(mask_node):
            full_node = graph.call_function(
                torch.full,
                args=(size_node, -float("inf")),
                kwargs={"dtype": score_dtype},
            )

        # Create a new node for torch.where
        with graph.inserting_after(full_node):
            where_node = graph.call_function(
                torch.ops.aten.where, args=(mask_node, qk_data_node, full_node)
            )

        # Update the output node to return the result of torch.where
        output_node.args = (where_node,)

        graph.lint()
        converted = torch.fx.GraphModule(gm, graph)
        return converted

    converted_mask_graph_module = convert_mask_graph_module(mask_graph)

    mask_graph_buffer = build_subgraph_module_buffer(
        mask_graph_placeholder_inps + list(mask_mod_other_buffers),
        converted_mask_graph_module,
    )

    # Clear the pending fresh unbacked symbols that are created for cur_qSplitSize and cur_kvSplitSize in the current kernel.
    pending = V.graph.sizevars.shape_env.pending_fresh_unbacked_symbols
    V.graph.sizevars.shape_env.pending_fresh_unbacked_symbols = [
        x for x in pending if x not in (cur_qSplitSize, cur_kvSplitSize)
    ]

    buffer_list = (
        placeholder_inps
        + list(score_mod_other_buffers)
        + mask_graph_placeholder_inps
        + list(mask_mod_other_buffers)
    )
    for item in buffer_list:
        if isinstance(item, TensorBox):
            fake_buffers.append(item.data.data)  # type: ignore[attr-defined]

    # CPU kernel requires last dim to be contiguous
    query, key, value = map(contiguous_last_dim, [query, key, value])

    (
        query,
        key,
        value,
        kv_num_blocks,
        kv_indices,
        full_kv_num_blocks,
        full_kv_indices,
        q_num_blocks,
        q_indices,
        full_q_num_blocks,
        full_q_indices,
    ) = maybe_realize(
        [
            query,
            key,
            value,
            kv_num_blocks,
            kv_indices,
            full_kv_num_blocks,
            full_kv_indices,
            q_num_blocks,
            q_indices,
            full_q_num_blocks,
            full_q_indices,
        ]
    )

    if len(OrderedSet([query.get_name(), key.get_name(), value.get_name()])) != 3:
        raise NotImplementedError(
            "Unsupported for now if query, key, value are the same buffer."
        )
    if query.get_dtype() not in [torch.float, torch.bfloat16, torch.float16]:
        raise NotImplementedError(
            "`torch.float` , `torch.float16` and `torch.bfloat16` are supported in FlexAttention for CPU device. "
            f"Found input tensors are `{query.get_dtype()}`."
        )
    score_mod_other_buffers = maybe_realize(score_mod_other_buffers)
    mask_mod_other_buffers = maybe_realize(mask_mod_other_buffers)
    Bq, Hq, seq_len_q, qk_head_dim = query.get_size()
    Bkv, Hkv, seq_len_kv, v_head_dim = value.get_size()
    B = Bq

    # Construct output layout with strides matching the query.
    out_size = [B, Hq, seq_len_q, v_head_dim]
    out_strides = infer_dense_strides(out_size, query.get_stride())

    layout = FixedLayout(
        query.get_device(),
        query.get_dtype(),
        [B, Hq, seq_len_q, v_head_dim],
        stride=[sympy.sympify(s) for s in out_strides],
    )
    _choices: list[Any] = []
    input_nodes = [query, key, value, kv_num_blocks, kv_indices]
    if not full_kv_num_blocks:
        no_full_kv_block = True
    else:
        no_full_kv_block = False
        input_nodes += [full_kv_num_blocks]
        input_nodes += [full_kv_indices]
    has_other_buffer = False
    kernel_input_name_to_buffer = {}
    if score_mod_other_buffers or mask_mod_other_buffers:
        has_other_buffer = True

        for prefix, buffers in [
            ("score_others", score_mod_other_buffers),
            ("mask_others", mask_mod_other_buffers),
        ]:
            kernel_input_name_to_buffer.update(
                {f"{prefix}_{i}": buf for i, buf in enumerate(buffers)}
            )
        input_nodes += [
            value
            for value in kernel_input_name_to_buffer.values()
            if not isinstance(value, sympy.Symbol)
        ]

    skip_mask_score = kernel_options.get("SKIP_MASK_SCORE", False)
    # Mark SPARSE_KV_BLOCK_SIZE & SPARSE_Q_BLOCK_SIZE as static shapes and add guards.
    SPARSE_KV_BLOCK_SIZE = V.graph.sizevars.guard_int(SPARSE_KV_BLOCK_SIZE)
    SPARSE_Q_BLOCK_SIZE = V.graph.sizevars.guard_int(SPARSE_Q_BLOCK_SIZE)
    assert V.graph.sizevars.evaluate_expr(
        sympy.Le(seq_len_q, sympy.Mul(kv_indices.get_size()[-2], SPARSE_Q_BLOCK_SIZE))
    ), (
        "Q seqlen must be smaller than the block_mask size in the Q dimension, considering pass a larger block_mask."
    )
    assert V.graph.sizevars.evaluate_expr(
        sympy.Le(seq_len_kv, sympy.Mul(kv_indices.get_size()[-1], SPARSE_KV_BLOCK_SIZE))
    ), (
        "KV seqlen must be smaller than the block_mask size in the KV dimension, considering pass a larger block_mask."
    )
    CppFlexAttentionTemplate.add_choices(
        choices=_choices,
        input_nodes=input_nodes,
        layout=layout,
        scale=scale,
        score_mod=None if skip_mask_score else subgraph_buffer,
        mask_mod=None if skip_mask_score else mask_graph_buffer,
        kv_block_size=SPARSE_KV_BLOCK_SIZE,
        q_block_size=SPARSE_Q_BLOCK_SIZE,
        has_other_buffer=has_other_buffer,
        no_full_kv_block=no_full_kv_block,
        fake_buffers=fake_buffers,
        len_score_other=len(score_mod_other_buffers),
        len_mask_other=len(mask_mod_other_buffers),
        kernel_input_name_to_buffer=kernel_input_name_to_buffer,
        block_vars=(cur_qSplitSize, cur_kvSplitSize),
    )
    inputs_for_autotuning = [
        query,
        key,
        value,
    ]
    res = autotune_select_algorithm(
        "flex_attention",
        _choices,
        inputs_for_autotuning,
        layout,
    )

    # need subgraph inputs and outputs to analyze all symints used in flex attention
    res.data.data.subgraph_inps = list(score_mod_other_buffers) + list(
        mask_mod_other_buffers
    )
    res.data.data.subgraph_outs = get_fwd_subgraph_outputs(
        subgraph_buffer, mask_graph_buffer
    )

    return (res,)
