"""Compare graphs generated by vLLM-Spyre vs AFTU

Run `python -m pytest tests/aftu/test_compare_graphs.py`.
"""

import os
import sys
import tempfile

import pytest
import torch
from graph_compare_utils import (
    collect_graph_files,
    compare_graphs,
    get_model_path,
    run_inference_py_and_get_graphs,
)
from pytest_mock.plugin import MockerFixture
from spyre_util import DecodeWarmupShapes, ModelInfo, patch_environment
from vllm import LLM

from vllm_spyre.model_executor.model_loader.spyre import SpyreCausalLM


# TODO: AFTU uses float32 for mask, while we use float16, which
# can generates different graphs. For now, until AFTU does not change
# this, we can just mock this method to match with AFTU behavior.
# NOTE: we need to set VLLM_ENABLE_V1_MULTIPROCESSING=0 otherwise this
# mock will not propagate to the child process of the model runner
def mock_get_mask_dtype(mocker: MockerFixture):
    mocker.patch.object(SpyreCausalLM, "get_mask_dtype", return_value=torch.float32)


@pytest.mark.spyre
@pytest.mark.cb
def test_compare_graphs_cb(
    model: ModelInfo,
    max_num_seqs: int,
    max_model_len: int,
    monkeypatch: pytest.MonkeyPatch,
    mocker: MockerFixture,
):
    """Test if the graphs generated by vllm-spyre matches with the ones by
    AFTU in continuous batching mode"""

    model_path = get_model_path(model)

    attn_type = "paged_fp8" if model.is_quantized else "paged"

    if model.is_quantized:
        mock_get_mask_dtype(mocker)

    # fmt: off
    inference_py_args = [
        sys.executable, "-m", "aiu_fms_testing_utils.scripts.inference",
        "--architecture", "hf_pretrained",
        "--model_path", model_path,
        "--tokenizer", model_path,
        "--unfuse_weights",
        "--device_type", "aiu",
        "--compile",
        "--cast_bf16_to_fp16",
        "--compile_dynamic",
        "--min_pad_length", "64",
        "--max_new_tokens", "5",
        "--batch_size", str(max_num_seqs),
        "--compile_dynamic_sendnn",
        "--attention_type", attn_type,
    ]
    # fmt: on

    if not model.is_quantized:
        inference_py_args += ["--default_dtype", "fp16"]

    extra_env = {
        "VLLM_DT_MAX_CONTEXT_LEN": str(max_model_len),
        "VLLM_DT_MAX_BATCH_SIZE": str(max_num_seqs),
        "VLLM_DT_MAX_BATCH_TKV_LIMIT": str(1024 * 128),
    }
    aftu_graphs = run_inference_py_and_get_graphs(inference_py_args, extra_env)

    assert len(aftu_graphs) > 0

    # VLLM

    monkeypatch.setenv("DEE_DUMP_GRAPHS", "vllm_cb")
    # Disable cache to produce the graphs
    monkeypatch.setenv("TORCH_SENDNN_CACHE_ENABLE", "0")

    # need for the mocker
    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", 0)
    patch_environment(use_cb=True, warmup_shapes=None, backend="sendnn", monkeypatch=monkeypatch)

    # NB: We need AFTU and vllm-spyre to have the same number of paged
    # attention blocks for the graph comparison.
    #
    # vllm-spyre adds 1 block to the full batch count for the padding block and
    # rounds up to a multiple of the batch size (see get_total_spyre_blocks in
    # spyre_model_runner.py).
    #
    # Adjusting the context length for vLLM by the block size counters the
    # padding block adjustment for alignment with AFTU
    adjusted_max_model_len = max_model_len - 64
    original_cwd = os.getcwd()
    try:
        # Change to temp dir to set the test environment clean
        with tempfile.TemporaryDirectory() as tmpdir:
            os.chdir(tmpdir)

            # We only need to load the model
            LLM(
                model=model.name,
                revision=model.revision,
                max_model_len=adjusted_max_model_len,
                tensor_parallel_size=1,
                max_num_seqs=max_num_seqs,
            )

            vllm_graphs = collect_graph_files(tmpdir)
    finally:
        # Restore in case of exception
        os.chdir(original_cwd)

    assert compare_graphs(vllm_graphs, aftu_graphs)


@pytest.mark.spyre
@pytest.mark.parametrize("warmup_shapes", [[(64, 4, 4)]])  # (prompt_length/new_tokens/batch_size)
def test_compare_graphs_static_batching(
    model: ModelInfo,
    warmup_shapes: DecodeWarmupShapes,
    monkeypatch: pytest.MonkeyPatch,
    mocker: MockerFixture,
) -> None:
    """Test if the graphs generated by vllm-spyre matches with the ones by
    AFTU in static batching mode"""

    attn_type = "math_fp8" if model.is_quantized else "sdpa"

    if model.is_quantized:
        mock_get_mask_dtype(mocker)

    model_path = get_model_path(model)

    # fmt: off
    inference_py_args = [
        sys.executable, "-m", "aiu_fms_testing_utils.scripts.inference",
        "--architecture", "hf_pretrained",
        "--model_path", model_path,
        "--tokenizer", model_path,
        "--unfuse_weights",
        "--device_type", "aiu",
        "--compile",
        "--cast_bf16_to_fp16",
        "--compile_dynamic",
        "--fixed_prompt_length", str(warmup_shapes[0][0]),
        "--max_new_tokens", str(warmup_shapes[0][1]),
        "--batch_size", str(warmup_shapes[0][2]),
        "--attention_type", attn_type,
    ]
    # fmt: on

    if not model.is_quantized:
        inference_py_args += ["--default_dtype", "fp16"]

    aftu_graphs = run_inference_py_and_get_graphs(inference_py_args)
    assert len(aftu_graphs) > 0
    # VLLM

    monkeypatch.setenv("DEE_DUMP_GRAPHS", "vllm_sb")
    # Disable cache to produce the graphs
    monkeypatch.setenv("TORCH_SENDNN_CACHE_ENABLE", "0")
    # needed for the mocker
    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", 0)
    patch_environment(
        use_cb=False, warmup_shapes=warmup_shapes, backend="sendnn", monkeypatch=monkeypatch
    )

    original_cwd = os.getcwd()
    try:
        # Change to temp dir to set the test environment clean
        with tempfile.TemporaryDirectory() as tmpdir:
            os.chdir(tmpdir)

            LLM(
                model=model.name,
                revision=model.revision,
                max_model_len=2048,
                tensor_parallel_size=1,
                max_num_seqs=warmup_shapes[0][2],
            )

            vllm_graphs = collect_graph_files(tmpdir)
    finally:
        # Restore in case of exception
        os.chdir(original_cwd)

    assert compare_graphs(vllm_graphs, aftu_graphs)


@pytest.mark.spyre
@pytest.mark.chunked_prefill
def test_compare_graphs_chunked_prefill(
    model: ModelInfo,
    max_num_seqs: int,
    max_model_len: int,
    monkeypatch: pytest.MonkeyPatch,
):
    """Test if the graphs generated by vllm-spyre matches with the ones by
    AFTU with chunked prefill enabled"""

    model_path = get_model_path(model)

    chunk_size = 128

    # fmt: off
    inference_py_args = [
        sys.executable, "-m", "aiu_fms_testing_utils.scripts.inference",
        "--architecture", "hf_pretrained",
        "--model_path", model_path,
        "--tokenizer", model_path,
        "--unfuse_weights",
        "--device_type", "aiu",
        "--compile",
        "--cast_bf16_to_fp16",
        "--compile_dynamic",
        "--min_pad_length", "64",
        "--max_new_tokens", "5",
        "--batch_size", str(max_num_seqs),
        "--compile_dynamic_sendnn",
        "--prefill_chunk_size", str(chunk_size),
    ]
    # fmt: on

    if not model.is_quantized:
        inference_py_args += ["--default_dtype", "fp16"]
        inference_py_args += ["--attention_type", "paged"]
    else:
        inference_py_args += ["--attention_type", "paged_fp8"]

    extra_env = {
        "VLLM_DT_MAX_CONTEXT_LEN": str(max_model_len),
        "VLLM_DT_MAX_BATCH_SIZE": str(max_num_seqs),
        "VLLM_DT_MAX_BATCH_TKV_LIMIT": str(1024 * 128),
        "VLLM_DT_CHUNK_LEN": str(chunk_size),
    }
    aftu_graphs = run_inference_py_and_get_graphs(inference_py_args, extra_env)

    assert len(aftu_graphs) > 0

    # VLLM

    monkeypatch.setenv("DEE_DUMP_GRAPHS", "vllm_cb_cp")
    # Disable cache to produce the graphs
    monkeypatch.setenv("TORCH_SENDNN_CACHE_ENABLE", "0")

    monkeypatch.setenv("VLLM_DT_CHUNK_LEN", str(chunk_size))
    patch_environment(
        use_cb=True,
        warmup_shapes=None,
        backend="sendnn",
        monkeypatch=monkeypatch,
        use_chunked_prefill=True,
    )

    # NB: We need AFTU and vllm-spyre to have the same number of paged
    # attention blocks for the graph comparison.
    #
    # vllm-spyre adds 1 block to the full batch count for the padding block and
    # rounds up to a multiple of the batch size (see get_total_spyre_blocks in
    # spyre_model_runner.py).
    #
    # Adjusting the context length for vLLM by the block size counters the
    # padding block adjustment for alignment with AFTU
    adjusted_max_model_len = max_model_len - 64
    original_cwd = os.getcwd()
    try:
        # Change to temp dir to set the test environment clean
        with tempfile.TemporaryDirectory() as tmpdir:
            os.chdir(tmpdir)

            # We only need to load the model
            LLM(
                model=model.name,
                revision=model.revision,
                max_model_len=adjusted_max_model_len,
                tensor_parallel_size=1,
                max_num_batched_tokens=chunk_size,
                max_num_seqs=max_num_seqs,
            )

            vllm_graphs = collect_graph_files(tmpdir)
    finally:
        # Restore in case of exception
        os.chdir(original_cwd)

    assert compare_graphs(vllm_graphs, aftu_graphs)
