# coding=utf-8
# Copyright 2025 HuggingFace Inc.
#
# 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.

import inspect

import torch

from diffusers import FasterCacheConfig, PyramidAttentionBroadcastConfig
from diffusers.hooks.first_block_cache import FirstBlockCacheConfig
from diffusers.hooks.mag_cache import MagCacheConfig
from diffusers.hooks.pyramid_attention_broadcast import PyramidAttentionBroadcastHook
from diffusers.hooks.taylorseer_cache import TaylorSeerCacheConfig

from ...testing_utils import assert_tensors_close, is_cache, torch_device
from .common import BasePipelineOutputMixin


class CacheTesterMixin(BasePipelineOutputMixin):
    """
    Shared machinery for cache-hook tester mixins. Each cache backend subclasses this and supplies its own config,
    mirroring the model-level `cache.py` layout. Backends store their config *kwargs* as a dict class attribute and
    build a fresh config instance per test via `_get_cache_config()`; a shared config instance would leak per-test
    mutations (e.g. `current_timestep_callback`) across tests. The denoiser-level enable/disable inference comparison
    is shared via `_test_cache_inference`; backend-specific state/layer checks live on the subclasses.
    """

    def _get_cache_config(self):
        raise NotImplementedError("Subclass must implement `_get_cache_config`.")

    def _test_cache_inference(self, cache_config, num_inference_steps, expected_atol=0.1, set_timestep_callback=False):
        device = "cpu"  # ensure determinism for the device-dependent torch.Generator

        def create_pipe():
            torch.manual_seed(0)
            return self.get_pipeline(**self.get_dummy_components(num_layers=2)).to(device)

        def run_forward(pipe):
            torch.manual_seed(0)
            inputs = self.get_dummy_inputs()
            inputs["num_inference_steps"] = num_inference_steps
            return pipe(**inputs)[0]

        # Run inference without cache
        pipe = create_pipe()
        output = run_forward(pipe).flatten()
        original_image_slice = torch.cat((output[:8], output[-8:]))

        # Run inference with cache enabled
        pipe = create_pipe()
        if set_timestep_callback:
            cache_config.current_timestep_callback = lambda: pipe.current_timestep
        pipe.transformer.enable_cache(cache_config)
        output = run_forward(pipe).flatten()
        image_slice_enabled = torch.cat((output[:8], output[-8:]))

        # Run inference with cache disabled
        pipe.transformer.disable_cache()
        output = run_forward(pipe).flatten()
        image_slice_disabled = torch.cat((output[:8], output[-8:]))

        assert_tensors_close(
            image_slice_enabled,
            original_image_slice,
            atol=expected_atol,
            rtol=1e-5,
            msg="Cached outputs should not differ much in the specified timestep range.",
        )
        assert_tensors_close(
            image_slice_disabled,
            original_image_slice,
            atol=1e-4,
            rtol=1e-5,
            msg="Outputs from normal inference and after disabling cache should not differ.",
        )


@is_cache
class PyramidAttentionBroadcastTesterMixin(CacheTesterMixin):
    PAB_CONFIG = {
        "spatial_attention_block_skip_range": 2,
        "spatial_attention_timestep_skip_range": (100, 800),
        "spatial_attention_block_identifiers": ["transformer_blocks"],
    }

    def _get_cache_config(self):
        return PyramidAttentionBroadcastConfig(**self.PAB_CONFIG)

    def test_pyramid_attention_broadcast_layers(self):
        num_layers = 0
        num_single_layers = 0
        dummy_component_kwargs = {}
        dummy_component_parameters = inspect.signature(self.get_dummy_components).parameters
        if "num_layers" in dummy_component_parameters:
            num_layers = 2
            dummy_component_kwargs["num_layers"] = num_layers
        if "num_single_layers" in dummy_component_parameters:
            num_single_layers = 2
            dummy_component_kwargs["num_single_layers"] = num_single_layers

        pipe = self.get_pipeline(**self.get_dummy_components(**dummy_component_kwargs))

        pab_config = self._get_cache_config()
        pab_config.current_timestep_callback = lambda: pipe.current_timestep
        denoiser = pipe.transformer if hasattr(pipe, "transformer") else pipe.unet
        denoiser.enable_cache(pab_config)

        expected_hooks = 0
        if pab_config.spatial_attention_block_skip_range is not None:
            expected_hooks += num_layers + num_single_layers
        if pab_config.temporal_attention_block_skip_range is not None:
            expected_hooks += num_layers + num_single_layers
        if pab_config.cross_attention_block_skip_range is not None:
            expected_hooks += num_layers + num_single_layers

        denoiser = pipe.transformer if hasattr(pipe, "transformer") else pipe.unet
        count = 0
        for module in denoiser.modules():
            if hasattr(module, "_diffusers_hook"):
                hook = module._diffusers_hook.get_hook("pyramid_attention_broadcast")
                if hook is None:
                    continue
                count += 1
                assert isinstance(hook, PyramidAttentionBroadcastHook), (
                    "Hook should be of type PyramidAttentionBroadcastHook."
                )
                assert hook.state.cache is None, "Cache should be None at initialization."
        assert count == expected_hooks, "Number of hooks should match the expected number."

        # Perform dummy inference step to ensure state is updated
        def pab_state_check_callback(pipe, i, t, kwargs):
            for module in denoiser.modules():
                if hasattr(module, "_diffusers_hook"):
                    hook = module._diffusers_hook.get_hook("pyramid_attention_broadcast")
                    if hook is None:
                        continue
                    assert hook.state.cache is not None, "Cache should have updated during inference."
                    assert hook.state.iteration == i + 1, "Hook iteration state should have updated during inference."
            return {}

        inputs = self.get_dummy_inputs()
        inputs["num_inference_steps"] = 2
        inputs["callback_on_step_end"] = pab_state_check_callback
        pipe(**inputs)[0]

        # After inference, reset_stateful_hooks is called within the pipeline, which should have reset the states
        for module in denoiser.modules():
            if hasattr(module, "_diffusers_hook"):
                hook = module._diffusers_hook.get_hook("pyramid_attention_broadcast")
                if hook is None:
                    continue
                assert hook.state.cache is None, "Cache should be reset to None after inference."
                assert hook.state.iteration == 0, "Iteration should be reset to 0 after inference."

    def test_pyramid_attention_broadcast_inference(self, base_pipe_output, expected_atol: float = 0.2):
        # We need to use higher tolerance because we are using a random model. With a converged/trained model, the
        # tolerance can be lower. The no-PAB reference is the shared, class-cached `base_pipe_output`, so the PAB
        # runs below reuse the same components/inputs the fixture does (default `get_dummy_components()` and
        # `get_dummy_inputs()`) and reseed the global RNG before each forward, so the disabled run reproduces the
        # reference and the only remaining differences come from PAB itself.
        pipe = self.get_pipeline().to(torch_device)

        # Reference (no PAB) output, computed once per test class by the `base_pipe_output` fixture.
        original_image_slice = base_pipe_output.flatten()
        original_image_slice = torch.cat((original_image_slice[:8], original_image_slice[-8:]))

        # Run inference with PAB enabled
        pab_config = self._get_cache_config()
        pab_config.current_timestep_callback = lambda: pipe.current_timestep
        denoiser = pipe.transformer if hasattr(pipe, "transformer") else pipe.unet
        denoiser.enable_cache(pab_config)

        torch.manual_seed(0)
        output = pipe(**self.get_dummy_inputs())[0]
        image_slice_pab_enabled = output.flatten()
        image_slice_pab_enabled = torch.cat((image_slice_pab_enabled[:8], image_slice_pab_enabled[-8:]))

        # Run inference with PAB disabled — should reproduce the reference output
        denoiser.disable_cache()

        torch.manual_seed(0)
        output = pipe(**self.get_dummy_inputs())[0]
        image_slice_pab_disabled = output.flatten()
        image_slice_pab_disabled = torch.cat((image_slice_pab_disabled[:8], image_slice_pab_disabled[-8:]))

        assert_tensors_close(
            image_slice_pab_enabled,
            original_image_slice,
            atol=expected_atol,
            rtol=1e-5,
            msg="PAB outputs should not differ much in specified timestep range.",
        )
        assert_tensors_close(
            image_slice_pab_disabled,
            original_image_slice,
            atol=1e-4,
            rtol=1e-5,
            msg="Outputs from normal inference and after disabling cache should not differ.",
        )


@is_cache
class FasterCacheTesterMixin(CacheTesterMixin):
    FASTER_CACHE_CONFIG = {
        "spatial_attention_block_skip_range": 2,
        "spatial_attention_timestep_skip_range": (-1, 901),
        "unconditional_batch_skip_range": 2,
        "attention_weight_callback": lambda _: 0.5,
    }

    def _get_cache_config(self):
        return FasterCacheConfig(**self.FASTER_CACHE_CONFIG)

    def test_faster_cache_inference(self, expected_atol: float = 0.1):
        self._test_cache_inference(
            self._get_cache_config(), num_inference_steps=4, expected_atol=expected_atol, set_timestep_callback=True
        )

    def test_faster_cache_state(self):
        from diffusers.hooks.faster_cache import _FASTER_CACHE_BLOCK_HOOK, _FASTER_CACHE_DENOISER_HOOK

        num_layers = 0
        num_single_layers = 0
        dummy_component_kwargs = {}
        dummy_component_parameters = inspect.signature(self.get_dummy_components).parameters
        if "num_layers" in dummy_component_parameters:
            num_layers = 2
            dummy_component_kwargs["num_layers"] = num_layers
        if "num_single_layers" in dummy_component_parameters:
            num_single_layers = 2
            dummy_component_kwargs["num_single_layers"] = num_single_layers

        pipe = self.get_pipeline(**self.get_dummy_components(**dummy_component_kwargs))

        faster_cache_config = self._get_cache_config()
        faster_cache_config.current_timestep_callback = lambda: pipe.current_timestep
        pipe.transformer.enable_cache(faster_cache_config)

        # Hook registration/removal is covered at the model level (`_test_cache_hooks_registered`). Here we only
        # verify the pipeline-specific behaviour: that hook state evolves during the denoising loop and is reset
        # afterwards.
        denoiser = pipe.transformer if hasattr(pipe, "transformer") else pipe.unet

        # Perform inference to ensure that states are updated correctly
        def faster_cache_state_check_callback(pipe, i, t, kwargs):
            for name, module in denoiser.named_modules():
                if not hasattr(module, "_diffusers_hook"):
                    continue
                if name == "":
                    # Root denoiser module
                    state = module._diffusers_hook.get_hook(_FASTER_CACHE_DENOISER_HOOK).state
                    if not faster_cache_config.is_guidance_distilled:
                        assert state.low_frequency_delta is not None, "Low frequency delta should be set."
                        assert state.high_frequency_delta is not None, "High frequency delta should be set."
                else:
                    # Internal blocks
                    state = module._diffusers_hook.get_hook(_FASTER_CACHE_BLOCK_HOOK).state
                    assert state.cache is not None and len(state.cache) == 2, "Cache should be set."
                assert state.iteration == i + 1, "Hook iteration state should have updated during inference."
            return {}

        inputs = self.get_dummy_inputs()
        inputs["num_inference_steps"] = 4
        inputs["callback_on_step_end"] = faster_cache_state_check_callback
        _ = pipe(**inputs)[0]

        # After inference, reset_stateful_hooks is called within the pipeline, which should have reset the states
        for name, module in denoiser.named_modules():
            if not hasattr(module, "_diffusers_hook"):
                continue

            if name == "":
                # Root denoiser module
                state = module._diffusers_hook.get_hook(_FASTER_CACHE_DENOISER_HOOK).state
                assert state.iteration == 0, "Iteration should be reset to 0."
                assert state.low_frequency_delta is None, "Low frequency delta should be reset to None."
                assert state.high_frequency_delta is None, "High frequency delta should be reset to None."
            else:
                # Internal blocks
                state = module._diffusers_hook.get_hook(_FASTER_CACHE_BLOCK_HOOK).state
                assert state.iteration == 0, "Iteration should be reset to 0."
                assert state.batch_size is None, "Batch size should be reset to None."
                assert state.cache is None, "Cache should be reset to None."


# TODO(aryan, dhruv): the cache tester mixins should probably be rewritten so that more models can be tested out
# of the box once there is better cache support/implementation
@is_cache
class FirstBlockCacheTesterMixin(CacheTesterMixin):
    # threshold is intentionally set higher than usual values since we're testing with random unconverged models
    # that will not satisfy the expected properties of the denoiser for caching to be effective
    FIRST_BLOCK_CACHE_CONFIG = {"threshold": 0.8}

    def _get_cache_config(self):
        return FirstBlockCacheConfig(**self.FIRST_BLOCK_CACHE_CONFIG)

    def test_first_block_cache_inference(self, expected_atol: float = 0.1):
        self._test_cache_inference(self._get_cache_config(), num_inference_steps=4, expected_atol=expected_atol)


@is_cache
class TaylorSeerCacheTesterMixin(CacheTesterMixin):
    TAYLORSEER_CACHE_CONFIG = {
        "cache_interval": 5,
        "disable_cache_before_step": 10,
        "max_order": 1,
        "taylor_factors_dtype": torch.bfloat16,
        "use_lite_mode": True,
    }

    def _get_cache_config(self):
        return TaylorSeerCacheConfig(**self.TAYLORSEER_CACHE_CONFIG)

    def test_taylorseer_cache_inference(self, expected_atol: float = 0.1):
        self._test_cache_inference(self._get_cache_config(), num_inference_steps=50, expected_atol=expected_atol)


@is_cache
class MagCacheTesterMixin(CacheTesterMixin):
    MAG_CACHE_CONFIG = {
        "threshold": 0.06,
        "max_skip_steps": 3,
        "retention_ratio": 0.2,
        "num_inference_steps": 50,
        "mag_ratios": torch.ones(50),
    }

    def _get_cache_config(self):
        return MagCacheConfig(**self.MAG_CACHE_CONFIG)

    def test_mag_cache_inference(self, expected_atol: float = 0.1):
        self._test_cache_inference(self._get_cache_config(), num_inference_steps=50, expected_atol=expected_atol)