# Copyright 2025-present the HuggingFace Inc. 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.

"""
Test PEFT method x quantization method matrix, focusing on basic tests.
"""

from dataclasses import dataclass

import pytest
import torch
from accelerate.utils.memory import clear_device_cache
from transformers import AutoModelForCausalLM, BitsAndBytesConfig, TorchAoConfig

from peft import BOFTConfig, MissConfig, VeraConfig, get_peft_model
from peft.import_utils import (
    is_bnb_4bit_available,
    is_bnb_available,
    is_gptqmodel_available,
    is_torchao_available,
)
from peft.tuners.tuners_utils import BaseTunerLayer
from peft.utils import infer_device
from peft.utils.quantization_utils import (
    Bnb4bitBackend,
    Bnb8bitBackend,
    ForwardOnlyQuantizationBackend,
    TorchaoBackend,
)

from .testing_utils import hub_online_once, set_init_weights_false


SEED = 0
DEVICE = infer_device()
MIN_CORR = 0.9
MAX_MSE = 1.0


@dataclass
class Bnb8bitLoader:
    name = "bnb_8bit"
    backend_cls = Bnb8bitBackend
    supports_merge = True
    supports_non_quantized_comparison = True
    model_id = "peft-internal-testing/opt-125m"
    expected_layer_count = 24  # (q_proj, v_proj) x 12 layers

    def load_model(self):
        quant_config = BitsAndBytesConfig(load_in_8bit=True)
        with hub_online_once(self.model_id):
            return AutoModelForCausalLM.from_pretrained(
                self.model_id, quantization_config=quant_config, device_map={"": DEVICE}
            )


@dataclass
class Bnb4bitLoader:
    name = "bnb_4bit"
    backend_cls = Bnb4bitBackend
    supports_merge = True
    supports_non_quantized_comparison = True
    model_id = "peft-internal-testing/opt-125m"
    expected_layer_count = 24  # (q_proj, v_proj) x 12 layers

    def load_model(self):
        quant_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_use_double_quant=False,
            bnb_4bit_compute_dtype=torch.float32,
        )
        with hub_online_once(self.model_id):
            return AutoModelForCausalLM.from_pretrained(
                self.model_id, quantization_config=quant_config, device_map={"": DEVICE}
            )


@dataclass
class TorchAoInt8WeightOnlyLoader:
    name = "torchao_int8_weight_only"
    backend_cls = TorchaoBackend
    supports_merge = True
    supports_non_quantized_comparison = True
    model_id = "peft-internal-testing/opt-125m"
    expected_layer_count = 24  # (q_proj, v_proj) x 12 layers

    def load_model(self):
        from torchao.quantization import Int8WeightOnlyConfig

        quant_config = TorchAoConfig(quant_type=Int8WeightOnlyConfig())
        with hub_online_once(self.model_id):
            return AutoModelForCausalLM.from_pretrained(
                self.model_id, quantization_config=quant_config, device_map={"": DEVICE}
            )


@dataclass
class TorchAoInt8DynamicActivationInt8WeightLoader:
    name = "torchao_int8_dynamic_activation_int8"
    backend_cls = TorchaoBackend
    supports_merge = False
    supports_non_quantized_comparison = True
    model_id = "peft-internal-testing/opt-125m"
    expected_layer_count = 24  # (q_proj, v_proj) x 12 layers

    def load_model(self):
        from torchao.quantization import Int8DynamicActivationInt8WeightConfig

        quant_config = TorchAoConfig(quant_type=Int8DynamicActivationInt8WeightConfig())
        with hub_online_once(self.model_id):
            return AutoModelForCausalLM.from_pretrained(
                self.model_id, quantization_config=quant_config, device_map={"": DEVICE}
            )


@dataclass
class Gptq4bitLoader:
    name = "gptq_4bit"
    backend_cls = ForwardOnlyQuantizationBackend
    supports_merge = False
    # No on-the-fly quantization path; the comparison would need a separate fp model.
    supports_non_quantized_comparison = False
    model_id = "marcsun13/opt-350m-gptq-4bit"
    expected_layer_count = 24  # (q_proj, v_proj) x 12 layers

    def load_model(self):
        from transformers import GPTQConfig

        quant_config = GPTQConfig(bits=4)
        with hub_online_once(self.model_id):
            return AutoModelForCausalLM.from_pretrained(
                self.model_id,
                quantization_config=quant_config,
                dtype=torch.float16,
                device_map={"": DEVICE},
            )


QUANTIZATION_BACKENDS = []
if is_bnb_available():
    QUANTIZATION_BACKENDS.append(Bnb8bitLoader())
if is_bnb_4bit_available():
    QUANTIZATION_BACKENDS.append(Bnb4bitLoader())
if is_torchao_available():
    QUANTIZATION_BACKENDS.append(TorchAoInt8WeightOnlyLoader())
    QUANTIZATION_BACKENDS.append(TorchAoInt8DynamicActivationInt8WeightLoader())
if is_gptqmodel_available():
    QUANTIZATION_BACKENDS.append(Gptq4bitLoader())


def _quant_id(backend):
    return backend.name


TEST_CASES = [
    (
        BOFTConfig,
        {"boft_block_size": 4, "target_modules": ["q_proj", "v_proj"]},
    ),
    (
        MissConfig,
        {"r": 2},
    ),
    (
        MissConfig,
        {"r": 2, "init_weights": "bat"},
    ),
    (
        VeraConfig,
        {"r": 8, "target_modules": ["q_proj", "v_proj"]},
    ),
]


def _peft_id(val):
    """Generate test id config_cls / config_kwargs."""
    if isinstance(val, dict):
        id_ = str(val).replace(" ", "")
    else:  # the PEFT config class
        id_ = val.__name__.removesuffix("Config").lower()
    return id_


def check_outputs_similar(x, y, min_corr=MIN_CORR, max_mse=MAX_MSE):
    # As quantization introduces a lot of error, use generous tolerances
    assert x.shape == y.shape
    corr = torch.corrcoef(torch.stack((x.flatten(), y.flatten())))
    mse = ((x - y) ** 2).mean()

    corr_checks = corr[0, 1] >= min_corr
    mse_checks = mse <= max_mse
    if not corr_checks and not mse_checks:
        assert False, f"both correlation ({corr[0, 1]:.4f}>={min_corr}) and MSE ({mse:.4f}<={max_mse}) check failed"
    if not corr_checks:
        assert False, f"correlation ({corr[0, 1]:.4f}>={min_corr}) check failed"
    if not mse_checks:
        assert False, f"MSE ({mse:.4f}<={max_mse}) check failed"


class TestQuantization:
    """Test for PEFT method x quantization method

    Note: It is recommended to keep the number of tests low, as the number of combinations is already large as is. This
    means testing multiple things per test, even if this is generally not desired. The reason is that we want to keep
    the number of model initializations to a minimum, as those take time.

    """

    @pytest.fixture(autouse=True)
    def set_seed(self):
        torch.manual_seed(SEED)

    @pytest.fixture(autouse=True)
    def cleanup(self):
        yield
        clear_device_cache(garbage_collection=True)

    @pytest.fixture
    def dummy_input(self):
        return torch.arange(10).view(1, -1).to(DEVICE)

    @pytest.mark.parametrize("quant", QUANTIZATION_BACKENDS, ids=_quant_id)
    @pytest.mark.parametrize("config_cls,config_kwargs", TEST_CASES, ids=_peft_id)
    def test_quantization_backend_is_set_and_repr(self, config_cls, config_kwargs, quant):
        """PEFT layers should have quantization_backend set"""
        model = quant.load_model()
        config = config_cls(**config_kwargs)
        model = get_peft_model(model, config)

        quantized_layers = [
            m for m in model.modules() if isinstance(m, BaseTunerLayer) and m.quantization_backend is not None
        ]
        assert len(quantized_layers) == quant.expected_layer_count
        for layer in quantized_layers:
            rep = repr(layer)
            assert "quantization_backend=" in rep

    @pytest.mark.parametrize("quant", QUANTIZATION_BACKENDS, ids=_quant_id)
    @pytest.mark.parametrize("config_cls,config_kwargs", TEST_CASES, ids=_peft_id)
    def test_forward_changes_output(self, config_cls, config_kwargs, quant, dummy_input):
        """Check that the forward pass works, also check if the results are affected"""
        config_kwargs = set_init_weights_false(config_cls, config_kwargs)
        model = quant.load_model()

        with torch.inference_mode():
            out_base = model(dummy_input).logits

        config = config_cls(**config_kwargs)
        model = get_peft_model(model, config)

        with torch.inference_mode():
            out_peft = model(dummy_input).logits

        atol, rtol = 1e-3, 1e-3
        assert not torch.allclose(out_base, out_peft, atol=atol, rtol=rtol)

    @pytest.mark.parametrize("quant", QUANTIZATION_BACKENDS, ids=_quant_id)
    @pytest.mark.parametrize("config_cls,config_kwargs", TEST_CASES, ids=_peft_id)
    def test_quantized_output_similar_to_non_quantized(self, config_cls, config_kwargs, quant, dummy_input):
        """Quantized PEFT output should be similar to non-quantized PEFT output.

        Both models use the same adapter config with non-identity init. The outputs won't match exactly due to
        quantization noise, but should be in the same ballpark.
        """
        if not quant.supports_non_quantized_comparison:
            pytest.skip(f"{quant.name} is pre-quantized; no on-the-fly non-quantized counterpart for comparison")

        config_kwargs = set_init_weights_false(config_cls, config_kwargs)

        # Quantized model
        model = quant.load_model()
        config = config_cls(**config_kwargs)
        torch.manual_seed(SEED)
        model = get_peft_model(model, config).eval()

        with torch.inference_mode():
            out_quant = model(dummy_input).logits

        del model

        # Non-quantized model
        with hub_online_once(quant.model_id):
            model = AutoModelForCausalLM.from_pretrained(quant.model_id, device_map={"": DEVICE})
        config = config_cls(**config_kwargs.copy())
        torch.manual_seed(SEED)
        model = get_peft_model(model, config).eval()

        with torch.inference_mode():
            out_non_quant = model(dummy_input).logits

        check_outputs_similar(out_non_quant, out_quant)

    @pytest.mark.parametrize("quant", QUANTIZATION_BACKENDS, ids=_quant_id)
    @pytest.mark.parametrize("config_cls,config_kwargs", TEST_CASES, ids=_peft_id)
    def test_merge_unmerge_unload(self, config_cls, config_kwargs, quant, dummy_input):
        """Check merge and unmerge roundtrip"""
        if not quant.supports_merge:
            pytest.skip(f"{quant.name} does not support merging")
        if (DEVICE == "cpu") and isinstance(quant, Bnb4bitLoader):
            pytest.skip("Bnb 4 bit quant with CPU results in high variance, skipping")

        config_kwargs = set_init_weights_false(config_cls, config_kwargs)
        model = quant.load_model()
        config = config_cls(**config_kwargs)
        torch.manual_seed(SEED)
        model = get_peft_model(model, config).eval()

        with torch.inference_mode():
            out_before = model(dummy_input).logits

        model.merge_adapter()
        with torch.inference_mode():
            out_merged = model(dummy_input).logits

        check_outputs_similar(out_before, out_merged)

        model.unmerge_adapter()
        with torch.inference_mode():
            out_unmerged = model(dummy_input).logits

        check_outputs_similar(out_before, out_unmerged)

        model.merge_adapter(safe_merge=True)
        with torch.inference_mode():
            out_merged_safe = model(dummy_input).logits

        check_outputs_similar(out_before, out_merged_safe)

        model.unmerge_adapter()
        model = model.merge_and_unload()
        with torch.inference_mode():
            out_unloaded = model(dummy_input).logits

        check_outputs_similar(out_before, out_unloaded)