"""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
from graph_compare_utils import (collect_graph_files, compare_graphs,
                                 get_aftu_script_dir, get_model_path,
                                 run_inference_py_and_get_graphs)
from llm_cache import DecodeWarmupShapes
from spyre_util import generate_spyre_vllm_output, get_chicken_soup_prompts
from vllm import SamplingParams


@pytest.mark.spyre
def test_compare_graphs_cb(model: str, max_num_seqs: int,
                           monkeypatch: pytest.MonkeyPatch, use_llm_cache):
    """Test that the spyre worker correctly outputs
    continuous batches of requests by comparing to HF"""

    # AFTU
    script_dir = get_aftu_script_dir()
    if script_dir is None:
        pytest.skip("aiu-fms-testing-utils is required "
                    "and is not installed to run this test")
    max_model_len = 256
    model_path = get_model_path(model)

    inference_py_args = [
        sys.executable, "scripts/inference.py", "--architecture",
        "hf_configured", "--model_path", model_path, "--variant", model_path,
        "--tokenizer", model_path, "--unfuse_weights", "--model_source", "hf",
        "--device_type", "aiu", "--compile", "--default_dtype", "fp16",
        "--compile_dynamic", "--min_pad_length", "64", "--max_new_tokens", "5",
        "--batch_size",
        str(max_num_seqs), "--compile_dynamic_sendnn", "--attention_type=paged"
    ]

    extra_env = {
        "VLLM_DT_MAX_CONTEXT_LEN": str(max_model_len),
        "VLLM_DT_MAX_BATCH_SIZE": str(max_num_seqs)
    }
    aftu_graphs = run_inference_py_and_get_graphs(inference_py_args,
                                                  script_dir, extra_env)

    assert len(aftu_graphs) > 0
    # VLLM
    prompts = get_chicken_soup_prompts(4)

    max_new_tokens = 20

    monkeypatch.setenv("DEE_DUMP_GRAPHS", "vllm_static")
    # Disable cache to produce the graphs
    monkeypatch.setenv("TORCH_SENDNN_CACHE_ENABLE", "0")
    vllm_sampling_params = SamplingParams(
        max_tokens=max_new_tokens,
        temperature=0,
        logprobs=0,  # return logprobs of generated tokens only
        ignore_eos=True)

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

            generate_spyre_vllm_output(model=model,
                                       prompts=prompts,
                                       max_model_len=max_model_len,
                                       sampling_params=vllm_sampling_params,
                                       tensor_parallel_size=1,
                                       backend='sendnn',
                                       max_num_seqs=max_num_seqs,
                                       use_cb=True,
                                       monkeypatch=monkeypatch)

            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: str, warmup_shapes: DecodeWarmupShapes,
        monkeypatch: pytest.MonkeyPatch) -> None:

    # AFTU
    script_dir = get_aftu_script_dir()
    if script_dir is None:
        pytest.skip("aiu-fms-testing-utils is required "
                    "and is not installed to run this test")
    model_path = get_model_path(model)

    inference_py_args = [
        sys.executable,
        "scripts/inference.py",
        "--architecture",
        "hf_configured",
        "--model_path",
        model_path,
        "--variant",
        model_path,
        "--tokenizer",
        model_path,
        "--unfuse_weights",
        "--model_source",
        "hf",
        "--device_type",
        "aiu",
        "--compile",
        "--default_dtype",
        "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]),
    ]

    aftu_graphs = run_inference_py_and_get_graphs(inference_py_args,
                                                  script_dir)
    assert len(aftu_graphs) > 0
    # VLLM
    prompts = get_chicken_soup_prompts(4)

    max_new_tokens = warmup_shapes[0][1]

    monkeypatch.setenv("DEE_DUMP_GRAPHS", "vllm_static")
    # Disable cache to produce the graphs
    monkeypatch.setenv("TORCH_SENDNN_CACHE_ENABLE", "0")
    vllm_sampling_params = SamplingParams(max_tokens=max_new_tokens,
                                          temperature=0,
                                          logprobs=0,
                                          ignore_eos=True)

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

            generate_spyre_vllm_output(model=model,
                                       prompts=prompts,
                                       warmup_shapes=warmup_shapes,
                                       max_model_len=2048,
                                       sampling_params=vllm_sampling_params,
                                       tensor_parallel_size=1,
                                       backend='sendnn',
                                       monkeypatch=monkeypatch)

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

    assert compare_graphs(vllm_graphs, aftu_graphs)
