from math import gcd
from typing import Optional

# imports
import pytest  # used for our unit tests
import torch
from src.diffusers.models.controlnets.controlnet_xs import \
    get_mid_block_adapter
from torch import nn

# function to test
# Copyright 2024 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.


# Dummy MidBlockControlNetXSAdapter for test purposes
class MidBlockControlNetXSAdapter(nn.Module):
    def __init__(self, base_to_ctrl, midblock, ctrl_to_base):
        super().__init__()
        self.base_to_ctrl = base_to_ctrl
        self.midblock = midblock
        self.ctrl_to_base = ctrl_to_base

# Dummy UNetMidBlock2DCrossAttn for test purposes
class UNetMidBlock2DCrossAttn(nn.Module):
    def __init__(
        self,
        transformer_layers_per_block,
        in_channels,
        out_channels,
        temb_channels,
        resnet_groups,
        cross_attention_dim,
        num_attention_heads,
        use_linear_projection,
        upcast_attention,
    ):
        super().__init__()
        self.transformer_layers_per_block = transformer_layers_per_block
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.temb_channels = temb_channels
        self.resnet_groups = resnet_groups
        self.cross_attention_dim = cross_attention_dim
        self.num_attention_heads = num_attention_heads
        self.use_linear_projection = use_linear_projection
        self.upcast_attention = upcast_attention
from src.diffusers.models.controlnets.controlnet_xs import \
    get_mid_block_adapter


def zero_module(module):
    for p in module.parameters():
        nn.init.zeros_(p)
    return module

# unit tests

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

def test_basic_shapes_and_types():
    # Test with typical small channel values
    base_channels = 8
    ctrl_channels = 16
    codeflash_output = get_mid_block_adapter(base_channels, ctrl_channels); adapter = codeflash_output

def test_basic_with_custom_args():
    # Test with all arguments provided
    codeflash_output = get_mid_block_adapter(
        base_channels=4,
        ctrl_channels=12,
        temb_channels=24,
        max_norm_num_groups=6,
        transformer_layers_per_block=2,
        num_attention_heads=3,
        cross_attention_dim=128,
        upcast_attention=True,
        use_linear_projection=False,
    ); adapter = codeflash_output
    # Check midblock parameters
    mid = adapter.midblock

def test_zero_module_initialization():
    # All weights in base_to_ctrl and ctrl_to_base should be zero
    codeflash_output = get_mid_block_adapter(8, 12); adapter = codeflash_output
    for param in adapter.base_to_ctrl.parameters():
        pass
    for param in adapter.ctrl_to_base.parameters():
        pass

def test_find_largest_factor_basic():
    # Test the factor calculation logic used for resnet_groups
    # gcd(12, 8+12) = gcd(12, 20) = 4
    # max_norm_num_groups = 32
    # Should use min(gcd, max_norm_num_groups) = 4
    codeflash_output = get_mid_block_adapter(8, 12); adapter = codeflash_output

# -------------------------------
# Edge Test Cases
# -------------------------------

def test_base_and_ctrl_channels_equal():
    # When base and ctrl channels are equal
    codeflash_output = get_mid_block_adapter(16, 16); adapter = codeflash_output



def test_max_norm_num_groups_larger_than_gcd():
    # max_norm_num_groups is larger than gcd, should return gcd
    codeflash_output = get_mid_block_adapter(6, 12, max_norm_num_groups=100); adapter = codeflash_output

def test_max_norm_num_groups_smaller_than_gcd():
    # max_norm_num_groups is smaller than gcd, should return max_norm_num_groups if divides gcd
    codeflash_output = get_mid_block_adapter(10, 30, max_norm_num_groups=5); adapter = codeflash_output

def test_max_norm_num_groups_one():
    # max_norm_num_groups=1 should always return 1
    codeflash_output = get_mid_block_adapter(8, 12, max_norm_num_groups=1); adapter = codeflash_output

def test_negative_channels():
    # Negative channels should raise error
    with pytest.raises(Exception):
        get_mid_block_adapter(-8, 12)
    with pytest.raises(Exception):
        get_mid_block_adapter(8, -12)

def test_non_integer_channels():
    # Non-integer channels should raise error
    with pytest.raises(Exception):
        get_mid_block_adapter(8.5, 12)
    with pytest.raises(Exception):
        get_mid_block_adapter(8, "12")

def test_ctrl_channels_not_divisible_by_max_norm_num_groups():
    # When max_norm_num_groups does not divide gcd, should find next lower factor
    # gcd(14, 18) = 2, max_norm_num_groups=3, so should return 2
    codeflash_output = get_mid_block_adapter(4, 14, max_norm_num_groups=3); adapter = codeflash_output

# -------------------------------
# Large Scale Test Cases
# -------------------------------

def test_large_channels():
    # Test with large but reasonable channel sizes
    base_channels = 128
    ctrl_channels = 256
    codeflash_output = get_mid_block_adapter(base_channels, ctrl_channels); adapter = codeflash_output

def test_large_max_norm_num_groups():
    # Large max_norm_num_groups with small gcd, should return gcd
    codeflash_output = get_mid_block_adapter(250, 500, max_norm_num_groups=999); adapter = codeflash_output


def test_large_batch_conv2d_forward():
    # Test that the zero convs can process a large batch of data
    codeflash_output = get_mid_block_adapter(32, 64); adapter = codeflash_output
    x = torch.randn(8, 32, 16, 16)  # batch size 8, 32 channels, 16x16
    y = adapter.base_to_ctrl(x)

def test_maximum_tensor_size_under_100mb():
    # Ensure no tensors exceed 100MB (roughly 25 million float32 elements)
    # 100MB / 4 bytes = 25,000,000 elements
    # Let's use a 128x128x32x16 tensor (8,388,608 elements, ~32MB)
    codeflash_output = get_mid_block_adapter(32, 64); adapter = codeflash_output
    x = torch.randn(16, 32, 64, 8)
    y = adapter.base_to_ctrl(x)



def test_zero_module_sets_all_weights_to_zero():
    # Test that zero_module sets all weights to zero
    conv = nn.Conv2d(3, 5, 1)
    conv = zero_module(conv)
    for param in conv.parameters():
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from math import gcd
from typing import Optional

# imports
import pytest  # used for our unit tests
import torch
from src.diffusers.models.controlnets.controlnet_xs import \
    get_mid_block_adapter
from torch import nn

# function to test
# Copyright 2024 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.


# Mocks for external dependencies (since we do not have the real implementations)
class UNetMidBlock2DCrossAttn(nn.Module):
    def __init__(
        self,
        transformer_layers_per_block,
        in_channels,
        out_channels,
        temb_channels,
        resnet_groups,
        cross_attention_dim,
        num_attention_heads,
        use_linear_projection,
        upcast_attention,
    ):
        super().__init__()
        self.transformer_layers_per_block = transformer_layers_per_block
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.temb_channels = temb_channels
        self.resnet_groups = resnet_groups
        self.cross_attention_dim = cross_attention_dim
        self.num_attention_heads = num_attention_heads
        self.use_linear_projection = use_linear_projection
        self.upcast_attention = upcast_attention

class MidBlockControlNetXSAdapter(nn.Module):
    def __init__(self, base_to_ctrl, midblock, ctrl_to_base):
        super().__init__()
        self.base_to_ctrl = base_to_ctrl
        self.midblock = midblock
        self.ctrl_to_base = ctrl_to_base
from src.diffusers.models.controlnets.controlnet_xs import \
    get_mid_block_adapter

# unit tests

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

def test_basic_shapes_and_types():
    # Check that the returned object has the correct submodules and parameter shapes
    base_channels = 8
    ctrl_channels = 16
    codeflash_output = get_mid_block_adapter(base_channels, ctrl_channels); adapter = codeflash_output

def test_basic_zero_module_weights():
    # Ensure that all weights in base_to_ctrl and ctrl_to_base are initialized to zero
    codeflash_output = get_mid_block_adapter(4, 4); adapter = codeflash_output
    for param in adapter.base_to_ctrl.parameters():
        pass
    for param in adapter.ctrl_to_base.parameters():
        pass

def test_default_parameters():
    # Test that default parameters are set correctly in midblock
    codeflash_output = get_mid_block_adapter(6, 10); adapter = codeflash_output
    midblock = adapter.midblock

def test_custom_parameters():
    # Test that custom parameters are passed through
    codeflash_output = get_mid_block_adapter(
        base_channels=5,
        ctrl_channels=7,
        temb_channels=123,
        max_norm_num_groups=3,
        transformer_layers_per_block=2,
        num_attention_heads=4,
        cross_attention_dim=256,
        upcast_attention=True,
        use_linear_projection=False,
    ); adapter = codeflash_output
    midblock = adapter.midblock

def test_resnet_groups_factor():
    # Check that resnet_groups is the largest factor of gcd(ctrl_channels, ctrl_channels+base_channels) <= max_norm_num_groups
    base_channels = 6
    ctrl_channels = 10
    max_norm_num_groups = 4
    codeflash_output = get_mid_block_adapter(base_channels, ctrl_channels, max_norm_num_groups=max_norm_num_groups); adapter = codeflash_output

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

def test_base_ctrl_channels_equal():
    # Edge: base_channels == ctrl_channels
    base_channels = 12
    ctrl_channels = 12
    codeflash_output = get_mid_block_adapter(base_channels, ctrl_channels); adapter = codeflash_output



def test_max_norm_num_groups_larger_than_gcd():
    # max_norm_num_groups > gcd(ctrl_channels, ctrl_channels+base_channels)
    base_channels = 7
    ctrl_channels = 14
    max_norm_num_groups = 100
    codeflash_output = get_mid_block_adapter(base_channels, ctrl_channels, max_norm_num_groups=max_norm_num_groups); adapter = codeflash_output

def test_max_norm_num_groups_smaller_than_gcd():
    # max_norm_num_groups < gcd(ctrl_channels, ctrl_channels+base_channels)
    base_channels = 8
    ctrl_channels = 12
    max_norm_num_groups = 2  # gcd(12,20)=4, but only factors <=2 are allowed, so should be 2
    codeflash_output = get_mid_block_adapter(base_channels, ctrl_channels, max_norm_num_groups=max_norm_num_groups); adapter = codeflash_output

def test_find_largest_factor_returns_1():
    # Edge: test when only 1 is a factor (prime numbers)
    base_channels = 2
    ctrl_channels = 3
    max_norm_num_groups = 1
    codeflash_output = get_mid_block_adapter(base_channels, ctrl_channels, max_norm_num_groups=max_norm_num_groups); adapter = codeflash_output

def test_find_largest_factor_exact_match():
    # max_norm_num_groups exactly matches gcd
    base_channels = 4
    ctrl_channels = 8
    max_norm_num_groups = 4
    codeflash_output = get_mid_block_adapter(base_channels, ctrl_channels, max_norm_num_groups=max_norm_num_groups); adapter = codeflash_output



def test_large_channels():
    # Test with large but not excessive channel numbers
    base_channels = 512
    ctrl_channels = 768
    codeflash_output = get_mid_block_adapter(base_channels, ctrl_channels, max_norm_num_groups=32); adapter = codeflash_output

def test_many_different_channel_combinations():
    # Test a range of channel sizes for robustness
    for base_channels in [1, 2, 4, 8, 16, 32, 64, 128]:
        for ctrl_channels in [1, 2, 4, 8, 16, 32, 64, 128]:
            codeflash_output = get_mid_block_adapter(base_channels, ctrl_channels); adapter = codeflash_output

def test_performance_large_but_safe():
    # Large scale, but under 100MB tensor size
    base_channels = 128
    ctrl_channels = 256
    codeflash_output = get_mid_block_adapter(base_channels, ctrl_channels); adapter = codeflash_output
    # Try a forward pass through base_to_ctrl and ctrl_to_base with a reasonable input
    x_base = torch.randn(1, base_channels, 8, 8)
    x_ctrl = torch.randn(1, ctrl_channels, 8, 8)
    # base_to_ctrl should output shape (1, base_channels, 8, 8)
    y = adapter.base_to_ctrl(x_base)
    # ctrl_to_base should output shape (1, base_channels, 8, 8)
    y2 = adapter.ctrl_to_base(x_ctrl)

def test_extreme_max_norm_num_groups():
    # max_norm_num_groups is very large, but should be capped by gcd
    base_channels = 128
    ctrl_channels = 256
    max_norm_num_groups = 1000
    codeflash_output = get_mid_block_adapter(base_channels, ctrl_channels, max_norm_num_groups=max_norm_num_groups); adapter = codeflash_output

def test_all_zero_weights_large():
    # For a large model, all weights should still be zero
    codeflash_output = get_mid_block_adapter(128, 256); adapter = codeflash_output
    for param in adapter.base_to_ctrl.parameters():
        pass
    for param in adapter.ctrl_to_base.parameters():
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.