from typing import Any, List, Tuple

# imports
import pytest  # used for our unit tests
import torch
import torch.nn as nn
from src.diffusers.models.autoencoders.autoencoder_kl_wan import \
    AutoencoderKLWan

# function to test
# Copyright 2025 The Wan Team and The HuggingFace Team. All rights reserved.
#
# 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.


# Minimal stubs for required classes
class WanCausalConv3d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size):
        super().__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)

class WanEncoder3d(nn.Module):
    def __init__(self, base_dim, z_dim_2, dim_mult, num_res_blocks, attn_scales, temperal_downsample, dropout):
        super().__init__()
        # For testability, create a number of WanCausalConv3d modules
        self.layers = nn.ModuleList([
            WanCausalConv3d(base_dim, z_dim_2, 1)
            for _ in range(len(dim_mult))
        ])
    def modules(self):
        # Return self and all submodules for isinstance checks
        yield self
        for layer in self.layers:
            yield layer

class WanDecoder3d(nn.Module):
    def __init__(self, base_dim, z_dim, dim_mult, num_res_blocks, attn_scales, temperal_upsample, dropout):
        super().__init__()
        # For testability, create a number of WanCausalConv3d modules
        self.layers = nn.ModuleList([
            WanCausalConv3d(base_dim, z_dim, 1)
            for _ in range(len(dim_mult))
        ])
    def modules(self):
        # Return self and all submodules for isinstance checks
        yield self
        for layer in self.layers:
            yield layer

# Minimal stubs for mixins and utils
class ConfigMixin:
    pass
def register_to_config(fn):
    return fn
class FromOriginalModelMixin:
    pass
class PushToHubMixin:
    pass
CONFIG_NAME = "config.json"
def deprecate(*args, **kwargs):
    pass

class ModelMixin(nn.Module, PushToHubMixin):
    config_name = CONFIG_NAME
    _automatically_saved_args = ["_diffusers_version", "_class_name", "_name_or_path"]
    _supports_gradient_checkpointing = False
    _keys_to_ignore_on_load_unexpected = None
    _no_split_modules = None
    _keep_in_fp32_modules = None
    _skip_layerwise_casting_patterns = None
    _supports_group_offloading = True

    def __init__(self):
        super().__init__()
        self._gradient_checkpointing_func = None

    def __getattr__(self, name: str) -> Any:
        is_in_config = "_internal_dict" in self.__dict__ and hasattr(self.__dict__["_internal_dict"], name)
        is_attribute = name in self.__dict__
        if is_in_config and not is_attribute:
            deprecation_message = f"Accessing config attribute `{name}` directly via '{type(self).__name__}' object attribute is deprecated. Please access '{name}' over '{type(self).__name__}'s config object instead, e.g. 'unet.config.{name}'."
            deprecate("direct config name access", "1.0.0", deprecation_message, standard_warn=False, stacklevel=3)
            return self._internal_dict[name]
        return super().__getattr__(name)
from src.diffusers.models.autoencoders.autoencoder_kl_wan import \
    AutoencoderKLWan

# =========================
# Unit tests for clear_cache
# =========================

# --------
# Basic Test Cases
# --------

def test_clear_cache_basic_initialization():
    """Basic: After calling clear_cache, attributes are set with correct types and lengths."""
    model = AutoencoderKLWan()
    # Before clear_cache, attributes should not exist
    for attr in ['_conv_num', '_conv_idx', '_feat_map', '_enc_conv_num', '_enc_conv_idx', '_enc_feat_map']:
        pass
    model.clear_cache()









def test_clear_cache_does_not_affect_other_attributes():
    """Misc: clear_cache should not modify unrelated attributes."""
    model = AutoencoderKLWan()
    model.some_attr = 123
    model.clear_cache()



from typing import Any, List, Tuple

# imports
import pytest  # used for our unit tests
import torch
from src.diffusers.models.autoencoders.autoencoder_kl_wan import \
    AutoencoderKLWan

# function to test
# Copyright 2025 The Wan Team and The HuggingFace Team. All rights reserved.
#
# 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.


# Minimal stubs for missing dependencies
class ConfigMixin:
    pass

def register_to_config(fn):
    return fn

class FromOriginalModelMixin:
    pass

class PushToHubMixin:
    pass

CONFIG_NAME = "config.json"

def deprecate(*args, **kwargs):
    pass

# Minimal stub for WanCausalConv3d
class WanCausalConv3d(torch.nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size):
        super().__init__()
        self.conv = torch.nn.Conv3d(in_channels, out_channels, kernel_size)

# Minimal stub for WanEncoder3d
class WanEncoder3d(torch.nn.Module):
    def __init__(self, base_dim, z_dim_2, dim_mult, num_res_blocks, attn_scales, temperal_downsample, dropout):
        super().__init__()
        # For testing, just add a few WanCausalConv3d layers
        self.layers = torch.nn.ModuleList([
            WanCausalConv3d(base_dim, base_dim, 1) for _ in range(len(dim_mult))
        ])
    def modules(self):
        # Return self and submodules for isinstance checks
        yield self
        for l in self.layers:
            yield l

# Minimal stub for WanDecoder3d
class WanDecoder3d(torch.nn.Module):
    def __init__(self, base_dim, z_dim, dim_mult, num_res_blocks, attn_scales, temperal_upsample, dropout):
        super().__init__()
        self.layers = torch.nn.ModuleList([
            WanCausalConv3d(base_dim, base_dim, 1) for _ in range(len(dim_mult))
        ])
    def modules(self):
        yield self
        for l in self.layers:
            yield l

class ModelMixin(torch.nn.Module, PushToHubMixin):
    config_name = CONFIG_NAME
    _automatically_saved_args = ["_diffusers_version", "_class_name", "_name_or_path"]
    _supports_gradient_checkpointing = False
    _keys_to_ignore_on_load_unexpected = None
    _no_split_modules = None
    _keep_in_fp32_modules = None
    _skip_layerwise_casting_patterns = None
    _supports_group_offloading = True

    def __init__(self):
        super().__init__()
        self._gradient_checkpointing_func = None

    def __getattr__(self, name: str) -> Any:
        is_in_config = "_internal_dict" in self.__dict__ and hasattr(self.__dict__["_internal_dict"], name)
        is_attribute = name in self.__dict__
        if is_in_config and not is_attribute:
            deprecation_message = f"Accessing config attribute `{name}` directly via '{type(self).__name__}' object attribute is deprecated. Please access '{name}' over '{type(self).__name__}'s config object instead, e.g. 'unet.config.{name}'."
            deprecate("direct config name access", "1.0.0", deprecation_message, standard_warn=False, stacklevel=3)
            return self._internal_dict[name]
        return super().__getattr__(name)
from src.diffusers.models.autoencoders.autoencoder_kl_wan import \
    AutoencoderKLWan

# unit tests

# ----------- BASIC TEST CASES -----------

def test_clear_cache_basic_attributes():
    """
    Test that clear_cache creates the expected attributes with correct types and values on a default model.
    """
    model = AutoencoderKLWan()
    # Before calling clear_cache, these attributes should not exist
    for attr in ['_conv_num', '_conv_idx', '_feat_map', '_enc_conv_num', '_enc_conv_idx', '_enc_feat_map']:
        pass
    model.clear_cache()

def test_clear_cache_multiple_calls_idempotent():
    """
    Test that calling clear_cache multiple times does not cause errors and always resets the attributes.
    """
    model = AutoencoderKLWan()
    model.clear_cache()
    # Mutate the attributes
    model._conv_idx.append(1)
    model._feat_map[0] = "something"
    model._enc_feat_map[1] = "else"
    model.clear_cache()

def test_clear_cache_with_no_layers():
    """
    Test a model whose encoder and decoder have no WanCausalConv3d layers.
    """
    class EmptyEncoder(torch.nn.Module):
        def modules(self):  # no submodules
            yield self
    class EmptyDecoder(torch.nn.Module):
        def modules(self):
            yield self
    model = AutoencoderKLWan()
    model.encoder = EmptyEncoder()
    model.decoder = EmptyDecoder()
    model.clear_cache()

# ----------- EDGE TEST CASES -----------

def test_clear_cache_encoder_decoder_different_counts():
    """
    Test when encoder and decoder have different numbers of WanCausalConv3d layers.
    """
    class CustomEncoder(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.c1 = WanCausalConv3d(1, 1, 1)
            self.c2 = WanCausalConv3d(1, 1, 1)
        def modules(self):
            yield self
            yield self.c1
            yield self.c2
    class CustomDecoder(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.c1 = WanCausalConv3d(1, 1, 1)
        def modules(self):
            yield self
            yield self.c1
    model = AutoencoderKLWan()
    model.encoder = CustomEncoder()
    model.decoder = CustomDecoder()
    model.clear_cache()

def test_clear_cache_non_wanconv_layers():
    """
    Test that clear_cache ignores layers that are not WanCausalConv3d.
    """
    class DummyLayer(torch.nn.Module):
        pass
    class EncoderWithMix(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.c1 = WanCausalConv3d(1, 1, 1)
            self.d1 = DummyLayer()
        def modules(self):
            yield self
            yield self.c1
            yield self.d1
    class DecoderWithMix(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.d1 = DummyLayer()
        def modules(self):
            yield self
            yield self.d1
    model = AutoencoderKLWan()
    model.encoder = EncoderWithMix()
    model.decoder = DecoderWithMix()
    model.clear_cache()

def test_clear_cache_overwrites_existing_attributes():
    """
    Test that clear_cache overwrites previously set attributes, even if they have wrong types.
    """
    model = AutoencoderKLWan()
    # Set wrong types
    model._conv_num = "wrong"
    model._conv_idx = "wrong"
    model._feat_map = "wrong"
    model._enc_conv_num = "wrong"
    model._enc_conv_idx = "wrong"
    model._enc_feat_map = "wrong"
    model.clear_cache()

def test_clear_cache_encoder_decoder_are_same_object():
    """
    Test that clear_cache works if encoder and decoder refer to the same object instance.
    """
    class Shared(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.c1 = WanCausalConv3d(1, 1, 1)
        def modules(self):
            yield self
            yield self.c1
    shared = Shared()
    model = AutoencoderKLWan()
    model.encoder = shared
    model.decoder = shared
    model.clear_cache()

def test_clear_cache_with_large_feat_map():
    """
    Test that clear_cache can handle a large number of WanCausalConv3d layers (edge of allowed scale).
    """
    class LargeEncoder(torch.nn.Module):
        def __init__(self, n):
            super().__init__()
            self.layers = torch.nn.ModuleList([WanCausalConv3d(1, 1, 1) for _ in range(n)])
        def modules(self):
            yield self
            for l in self.layers:
                yield l
    class LargeDecoder(torch.nn.Module):
        def __init__(self, n):
            super().__init__()
            self.layers = torch.nn.ModuleList([WanCausalConv3d(1, 1, 1) for _ in range(n)])
        def modules(self):
            yield self
            for l in self.layers:
                yield l
    n = 999  # just under 1000
    model = AutoencoderKLWan()
    model.encoder = LargeEncoder(n)
    model.decoder = LargeDecoder(n)
    model.clear_cache()

# ----------- LARGE SCALE TEST CASES -----------

def test_clear_cache_performance_large_scale():
    """
    Test clear_cache performance and correctness with maximum allowed (999) layers.
    """
    class ManyLayers(torch.nn.Module):
        def __init__(self, n):
            super().__init__()
            self.layers = torch.nn.ModuleList([WanCausalConv3d(1, 1, 1) for _ in range(n)])
        def modules(self):
            yield self
            for l in self.layers:
                yield l
    n = 999
    model = AutoencoderKLWan()
    model.encoder = ManyLayers(n)
    model.decoder = ManyLayers(n)
    import time
    start = time.time()
    model.clear_cache()
    elapsed = time.time() - start

def test_clear_cache_no_side_effects_on_other_attributes():
    """
    Test that clear_cache does not modify unrelated attributes.
    """
    model = AutoencoderKLWan()
    model.some_random_attr = 12345
    model.clear_cache()

def test_clear_cache_handles_inheritance_and_subclassing():
    """
    Test that clear_cache works in a subclass that overrides encoder/decoder with custom logic.
    """
    class CustomAE(AutoencoderKLWan):
        def __init__(self):
            super().__init__()
            class Dummy(torch.nn.Module):
                def modules(self):
                    yield self
            self.encoder = Dummy()
            self.decoder = Dummy()
    model = CustomAE()
    model.clear_cache()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.