#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
# This file is a part of the vllm-ascend project.
# Adapted from vllm-project/vllm/tests/spec_decode/e2e/test_integration.py
# Copyright 2023 The vLLM team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Tests which cover integration of the speculative decoding framework with
other features, e.g. cuda graphs.
"""

import pytest

from .conftest import run_equality_correctness_test

MAIN_MODEL = "JackFram/llama-68m"

# TODO: Open it when vllm-ascend support graph mode and
# support enforce_eager status is False to run model in graph mode
# @pytest.mark.parametrize(
#     "common_llm_kwargs",
#     [{

#         # Verify equality when cuda graphs allowed.
#         "enforce_eager": False,
#         "model_name": "JackFram/llama-68m",
#     }])
# @pytest.mark.parametrize(
#     "per_test_common_llm_kwargs",
#     [
#         {
#             # Identical models.
#             "speculative_model": "JackFram/llama-68m",
#             "num_speculative_tokens": 5,
#         },
#     ])
# @pytest.mark.parametrize("baseline_llm_kwargs", [{}])
# @pytest.mark.parametrize("test_llm_kwargs", [{}])
# @pytest.mark.parametrize("batch_size", [8])
# @pytest.mark.parametrize("output_len", [32])
# @pytest.mark.parametrize("seed", [1])
# def test_spec_decode_cuda_graph(vllm_runner, common_llm_kwargs,
#                                 per_test_common_llm_kwargs,
#                                 baseline_llm_kwargs, test_llm_kwargs,
#                                 batch_size: int, output_len: int, seed: int):
#     """Verify spec decode equality when cuda graphs are enabled.
#     """
#     run_equality_correctness_test(vllm_runner,
#                                   common_llm_kwargs,
#                                   per_test_common_llm_kwargs,
#                                   baseline_llm_kwargs,
#                                   test_llm_kwargs,
#                                   batch_size,
#                                   max_output_len=output_len,
#                                   seed=seed,
#                                   temperature=0.0)

# TODO: Open it when vllm-ascend support GPTQ-based model
# @pytest.mark.parametrize(
#     "common_llm_kwargs",
#     [{
#         "model_name": "JackFram/llama-160m",

#         # Skip cuda graph recording for fast test.
#         "enforce_eager": True,
#     }])
# @pytest.mark.parametrize("per_test_common_llm_kwargs", [
#     {
#         "speculative_model": "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit",
#         "num_speculative_tokens": 5,
#     },
# ])
# @pytest.mark.parametrize(
#     "test_llm_kwargs",
#     [
#         # Explicitly specify draft model quantization
#         {
#             "speculative_model_quantization": "gptq",
#         },
#         # Explicitly specify GPTQ-based draft model to use marlin quantization
#         {
#             "speculative_model_quantization": "marlin",
#         },
#         # Not explicitly specify draft model quantization
#         {
#             "speculative_model_quantization": None,
#         },
#     ])
# @pytest.mark.parametrize("baseline_llm_kwargs", [{}])
# @pytest.mark.parametrize("batch_size", [2])
# @pytest.mark.parametrize("seed", [1])
# def test_speculative_model_quantization_config(vllm_runner, common_llm_kwargs,
#                                                per_test_common_llm_kwargs,
#                                                baseline_llm_kwargs,
#                                                test_llm_kwargs,
#                                                batch_size: int, seed: int):
#     """Verify spec decode works well with draft model quantization configs.
#     """
#     run_equality_correctness_test(vllm_runner,
#                                   common_llm_kwargs,
#                                   per_test_common_llm_kwargs,
#                                   baseline_llm_kwargs,
#                                   test_llm_kwargs,
#                                   batch_size,
#                                   max_output_len=32,
#                                   seed=seed,
#                                   temperature=0.0)


@pytest.mark.parametrize(
    "common_llm_kwargs",
    [{
        "model_name": MAIN_MODEL,

        # Skip cuda graph recording for fast test.
        "enforce_eager": True,
        "speculative_model": "JackFram/llama-68m",
        "num_speculative_tokens": 3,
    }])
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
@pytest.mark.parametrize("test_llm_kwargs",
                         [{
                             "speculative_disable_mqa_scorer": True,
                         }])
@pytest.mark.parametrize("batch_size", [1, 5])
@pytest.mark.parametrize(
    "output_len",
    [
        # Use smaller output len for fast test.
        32,
    ])
@pytest.mark.parametrize("seed", [1])
def test_mqa_scorer(vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
                    baseline_llm_kwargs, test_llm_kwargs, batch_size: int,
                    output_len: int, seed: int):
    """Verify that ngram speculative decoding generates the same output 
    with batch expansion scorer and mqa scorer.
    """
    run_equality_correctness_test(vllm_runner,
                                  common_llm_kwargs,
                                  per_test_common_llm_kwargs,
                                  baseline_llm_kwargs,
                                  test_llm_kwargs,
                                  batch_size,
                                  max_output_len=output_len,
                                  seed=seed,
                                  temperature=0.0)
