import copy
import gc
import os
import tempfile

import pytest

from diffusers import NVIDIAModelOptConfig, SD3Transformer2DModel, StableDiffusion3Pipeline
from diffusers.utils import is_nvidia_modelopt_available, is_torch_available

from ...testing_utils import (
    backend_empty_cache,
    backend_reset_peak_memory_stats,
    enable_full_determinism,
    nightly,
    numpy_cosine_similarity_distance,
    require_accelerate,
    require_big_accelerator,
    require_modelopt_version_greater_or_equal,
    require_torch_cuda_compatibility,
    torch_device,
)


if is_nvidia_modelopt_available():
    import modelopt.torch.opt as mto
    import modelopt.torch.quantization as mtq

if is_torch_available():
    import torch

    from ..utils import LoRALayer, get_memory_consumption_stat

enable_full_determinism()


@nightly
@require_big_accelerator
@require_accelerate
@require_modelopt_version_greater_or_equal("0.33.1")
class ModelOptBaseTesterMixin:
    model_id = "hf-internal-testing/tiny-sd3-pipe"
    model_cls = SD3Transformer2DModel
    pipeline_cls = StableDiffusion3Pipeline
    torch_dtype = torch.bfloat16
    expected_memory_reduction = 0.0
    keep_in_fp32_module = ""
    modules_to_not_convert = ""
    _test_torch_compile = False

    @pytest.fixture(autouse=True)
    def _setup(self):
        backend_reset_peak_memory_stats(torch_device)
        backend_empty_cache(torch_device)
        gc.collect()
        yield
        backend_reset_peak_memory_stats(torch_device)
        backend_empty_cache(torch_device)
        gc.collect()

    def get_dummy_init_kwargs(self):
        return {"quant_type": "FP8"}

    def get_dummy_model_init_kwargs(self):
        return {
            "pretrained_model_name_or_path": self.model_id,
            "torch_dtype": self.torch_dtype,
            "quantization_config": NVIDIAModelOptConfig(**self.get_dummy_init_kwargs()),
            "subfolder": "transformer",
        }

    def test_modelopt_layers(self):
        model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
        for name, module in model.named_modules():
            if isinstance(module, torch.nn.Linear):
                assert mtq.utils.is_quantized(module)

    def test_modelopt_memory_usage(self):
        inputs = self.get_dummy_inputs()
        inputs = {
            k: v.to(device=torch_device, dtype=torch.bfloat16) for k, v in inputs.items() if not isinstance(v, bool)
        }

        unquantized_model = self.model_cls.from_pretrained(
            self.model_id, torch_dtype=self.torch_dtype, subfolder="transformer"
        )
        unquantized_model.to(torch_device)
        unquantized_model_memory = get_memory_consumption_stat(unquantized_model, inputs)

        quantized_model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
        quantized_model.to(torch_device)
        quantized_model_memory = get_memory_consumption_stat(quantized_model, inputs)

        assert unquantized_model_memory / quantized_model_memory >= self.expected_memory_reduction

    def test_keep_modules_in_fp32(self):
        _keep_in_fp32_modules = self.model_cls._keep_in_fp32_modules
        self.model_cls._keep_in_fp32_modules = self.keep_in_fp32_module

        model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
        model.to(torch_device)

        for name, module in model.named_modules():
            if isinstance(module, torch.nn.Linear):
                if name in model._keep_in_fp32_modules:
                    assert module.weight.dtype == torch.float32
        self.model_cls._keep_in_fp32_modules = _keep_in_fp32_modules

    def test_modules_to_not_convert(self):
        init_kwargs = self.get_dummy_model_init_kwargs()
        quantization_config_kwargs = self.get_dummy_init_kwargs()
        quantization_config_kwargs.update({"modules_to_not_convert": self.modules_to_not_convert})
        quantization_config = NVIDIAModelOptConfig(**quantization_config_kwargs)
        init_kwargs.update({"quantization_config": quantization_config})

        model = self.model_cls.from_pretrained(**init_kwargs)
        model.to(torch_device)

        for name, module in model.named_modules():
            if name in self.modules_to_not_convert:
                assert not mtq.utils.is_quantized(module)

    def test_dtype_assignment(self):
        model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())

        with pytest.raises(ValueError):
            model.to(torch.float16)

        with pytest.raises(ValueError):
            device_0 = f"{torch_device}:0"
            model.to(device=device_0, dtype=torch.float16)

        with pytest.raises(ValueError):
            model.float()

        with pytest.raises(ValueError):
            model.half()

        model.to(torch_device)

    def test_serialization(self):
        model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
        inputs = self.get_dummy_inputs()

        model.to(torch_device)
        with torch.no_grad():
            model_output = model(**inputs)

        with tempfile.TemporaryDirectory() as tmp_dir:
            model.save_pretrained(tmp_dir)
            saved_model = self.model_cls.from_pretrained(
                tmp_dir,
                torch_dtype=torch.bfloat16,
            )

        saved_model.to(torch_device)
        with torch.no_grad():
            saved_model_output = saved_model(**inputs)

        assert torch.allclose(model_output.sample, saved_model_output.sample, rtol=1e-5, atol=1e-5)

    def test_torch_compile(self):
        if not self._test_torch_compile:
            return

        model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs())
        compiled_model = torch.compile(model, mode="max-autotune", fullgraph=True, dynamic=False)

        model.to(torch_device)
        with torch.no_grad():
            model_output = model(**self.get_dummy_inputs()).sample

        compiled_model.to(torch_device)
        with torch.no_grad():
            compiled_model_output = compiled_model(**self.get_dummy_inputs()).sample

        model_output = model_output.detach().float().cpu().numpy()
        compiled_model_output = compiled_model_output.detach().float().cpu().numpy()

        max_diff = numpy_cosine_similarity_distance(model_output.flatten(), compiled_model_output.flatten())
        assert max_diff < 1e-3

    def test_device_map_error(self):
        with pytest.raises(ValueError):
            _ = self.model_cls.from_pretrained(
                **self.get_dummy_model_init_kwargs(),
                device_map={0: "8GB", "cpu": "16GB"},
            )

    def get_dummy_inputs(self):
        batch_size = 1
        seq_len = 16
        height = width = 32
        num_latent_channels = 4
        caption_channels = 8

        torch.manual_seed(0)
        hidden_states = torch.randn((batch_size, num_latent_channels, height, width)).to(
            torch_device, dtype=torch.bfloat16
        )
        encoder_hidden_states = torch.randn((batch_size, seq_len, caption_channels)).to(
            torch_device, dtype=torch.bfloat16
        )
        timestep = torch.tensor([1.0]).to(torch_device, dtype=torch.bfloat16).expand(batch_size)

        return {
            "hidden_states": hidden_states,
            "encoder_hidden_states": encoder_hidden_states,
            "timestep": timestep,
        }

    def test_model_cpu_offload(self):
        init_kwargs = self.get_dummy_init_kwargs()
        transformer = self.model_cls.from_pretrained(
            self.model_id,
            quantization_config=NVIDIAModelOptConfig(**init_kwargs),
            subfolder="transformer",
            torch_dtype=torch.bfloat16,
        )
        pipe = self.pipeline_cls.from_pretrained(self.model_id, transformer=transformer, torch_dtype=torch.bfloat16)
        pipe.enable_model_cpu_offload(device=torch_device)
        _ = pipe("a cat holding a sign that says hello", num_inference_steps=2)

    def test_training(self):
        quantization_config = NVIDIAModelOptConfig(**self.get_dummy_init_kwargs())
        quantized_model = self.model_cls.from_pretrained(
            self.model_id,
            subfolder="transformer",
            quantization_config=quantization_config,
            torch_dtype=torch.bfloat16,
        ).to(torch_device)

        for param in quantized_model.parameters():
            param.requires_grad = False
            if param.ndim == 1:
                param.data = param.data.to(torch.float32)

        for _, module in quantized_model.named_modules():
            if hasattr(module, "to_q"):
                module.to_q = LoRALayer(module.to_q, rank=4)
            if hasattr(module, "to_k"):
                module.to_k = LoRALayer(module.to_k, rank=4)
            if hasattr(module, "to_v"):
                module.to_v = LoRALayer(module.to_v, rank=4)

        with torch.amp.autocast(str(torch_device), dtype=torch.bfloat16):
            inputs = self.get_dummy_inputs()
            output = quantized_model(**inputs)[0]
            output.norm().backward()

        for module in quantized_model.modules():
            if isinstance(module, LoRALayer):
                assert module.adapter[1].weight.grad is not None


class TestSanaTransformerFP8Weights(ModelOptBaseTesterMixin):
    expected_memory_reduction = 0.6

    def get_dummy_init_kwargs(self):
        return {"quant_type": "FP8"}

    @require_modelopt_version_greater_or_equal("0.44.0")
    def test_prequantized_serialization_with_device_map(self):
        mto.enable_huggingface_checkpointing()
        model = self.model_cls.from_pretrained(
            self.model_id,
            subfolder="transformer",
            torch_dtype=self.torch_dtype,
            quantization_config=NVIDIAModelOptConfig(
                quant_type="FP8", modelopt_config=copy.deepcopy(mtq.FP8_DEFAULT_CFG)
            ),
        )
        model.to(torch_device)

        with tempfile.TemporaryDirectory() as tmp_dir:
            model.save_pretrained(tmp_dir)
            assert os.path.isfile(os.path.join(tmp_dir, "modelopt_state.pth"))
            saved_model = self.model_cls.from_pretrained(
                tmp_dir,
                torch_dtype=self.torch_dtype,
                device_map=torch_device,
            )

        named_parameters = list(saved_model.named_parameters())
        named_buffers = list(saved_model.named_buffers())
        assert any(name.endswith(("_amax", "_scale")) for name, _ in named_buffers), (
            "The restored model did not contain ModelOpt quantizer buffers."
        )

        for tensor_kind, named_tensors in (("parameter", named_parameters), ("buffer", named_buffers)):
            for name, tensor in named_tensors:
                assert not tensor.is_meta, f"{tensor_kind} {name} was not materialized from meta."


class TestSanaTransformerINT8Weights(ModelOptBaseTesterMixin):
    expected_memory_reduction = 0.6
    _test_torch_compile = True

    def get_dummy_init_kwargs(self):
        return {"quant_type": "INT8"}


@require_torch_cuda_compatibility(8.0)
class TestSanaTransformerINT4Weights(ModelOptBaseTesterMixin):
    expected_memory_reduction = 0.55

    def get_dummy_init_kwargs(self):
        return {
            "quant_type": "INT4",
            "block_quantize": 128,
            "channel_quantize": -1,
            "disable_conv_quantization": True,
        }


@require_torch_cuda_compatibility(8.0)
class TestSanaTransformerNF4Weights(ModelOptBaseTesterMixin):
    expected_memory_reduction = 0.65

    def get_dummy_init_kwargs(self):
        return {
            "quant_type": "NF4",
            "block_quantize": 128,
            "channel_quantize": -1,
            "scale_block_quantize": 8,
            "scale_channel_quantize": -1,
            "modules_to_not_convert": ["conv"],
        }


@require_torch_cuda_compatibility(8.0)
class TestSanaTransformerNVFP4Weights(ModelOptBaseTesterMixin):
    expected_memory_reduction = 0.65

    def get_dummy_init_kwargs(self):
        return {
            "quant_type": "NVFP4",
            "block_quantize": 128,
            "channel_quantize": -1,
            "scale_block_quantize": 8,
            "scale_channel_quantize": -1,
            "modules_to_not_convert": ["conv"],
        }