import pytest  # used for our unit tests
import torch
from src.diffusers.schedulers.scheduling_ddpm import rescale_zero_terminal_snr

# unit tests

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

def test_basic_typical_betas():
    # Test with a typical linear schedule of betas
    betas = torch.linspace(0.0001, 0.02, 10)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    # The last SNR should be zero: i.e., the last cumulative alpha should be zero
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)

def test_basic_constant_betas():
    # Test with constant betas
    betas = torch.full((8,), 0.01)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)

def test_basic_small_tensor():
    # Test with a very small tensor (length=2)
    betas = torch.tensor([0.1, 0.2])
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)

def test_basic_dtype_preservation():
    # Test that dtype is preserved (float32)
    betas = torch.linspace(0.001, 0.01, 5, dtype=torch.float32)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output

def test_basic_device_preservation():
    # Test that device is preserved (cpu)
    betas = torch.linspace(0.001, 0.01, 5)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output

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

def test_edge_all_zeros():
    # All betas are zero (no noise)
    betas = torch.zeros(5)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output

def test_edge_all_ones():
    # All betas are one (full noise)
    betas = torch.ones(5)
    # This is an edge case: 1.0 - betas = 0, so cumulative product is zero, sqrt is zero, so division by zero may occur.
    # The code should not crash, but we expect NaNs or infs
    with pytest.raises(Exception):
        rescale_zero_terminal_snr(betas)

def test_edge_empty_tensor():
    # Empty tensor
    betas = torch.tensor([])
    # Should raise an error due to indexing [0] and [-1]
    with pytest.raises(IndexError):
        rescale_zero_terminal_snr(betas)

def test_edge_single_element():
    # Single element tensor
    betas = torch.tensor([0.5])
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    # Terminal SNR should be zero: 1-beta[0] == 0, so beta[0] == 1
    alphas = 1.0 - rescaled

def test_edge_negative_betas():
    # Negative betas are invalid, but let's test
    betas = torch.tensor([0.1, -0.2, 0.3])
    # Should not crash, but may produce NaNs due to sqrt of negative number
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output

def test_edge_large_betas():
    # Betas close to 1 (but not exactly 1)
    betas = torch.full((4,), 0.999)
    # 1-betas is very small, so cumprod will underflow to zero
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    # The last SNR should be zero (or numerically very close)
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)

def test_edge_high_precision():
    # Use float64 for high precision
    betas = torch.linspace(0.0001, 0.02, 10, dtype=torch.float64)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)

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

def test_large_scale_1000_elements():
    # Test with 1000 elements, linear schedule
    betas = torch.linspace(0.0001, 0.02, 1000)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)

def test_large_scale_random_betas():
    # Test with 999 random betas between 0 and 0.5
    torch.manual_seed(42)
    betas = torch.rand(999) * 0.5
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)

def test_large_scale_performance():
    # Test with a large tensor, close to 100MB (float32, 25_000_000 elements)
    # Each float32 is 4 bytes, so 25_000_000 * 4 = 100_000_000 bytes = 100MB
    size = 25_000_000
    betas = torch.linspace(0.0001, 0.01, size)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    # Check last SNR is zero (within reasonable tolerance for large tensors)
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)

# ------------------- ADDITIONAL FUNCTIONALITY TESTS -------------------

def test_functional_idempotence():
    # Applying the function twice should not change the result (idempotence)
    betas = torch.linspace(0.0001, 0.02, 50)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled1 = codeflash_output
    codeflash_output = rescale_zero_terminal_snr(rescaled1); rescaled2 = codeflash_output

def test_functional_monotonicity():
    # Test that monotonic increasing input betas produce monotonic increasing output betas
    betas = torch.linspace(0.0001, 0.02, 50)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    # The rescaled betas should be non-decreasing
    diffs = rescaled[1:] - rescaled[:-1]

def test_functional_nonnegativity():
    # Test that negative input betas produce NaNs or infs
    betas = torch.linspace(-0.1, 0.1, 10)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest  # used for our unit tests
import torch
from src.diffusers.schedulers.scheduling_ddpm import rescale_zero_terminal_snr

# unit tests

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

def test_basic_monotonic_increasing_betas():
    # Test with a small, monotonically increasing beta schedule
    betas = torch.linspace(0.01, 0.1, 10)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    # The last alpha_bar_sqrt should be zero (zero terminal SNR)
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    alphas_bar_sqrt = alphas_cumprod.sqrt()

def test_basic_constant_betas():
    # Test with constant betas
    betas = torch.full((5,), 0.05)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    # The last alpha_bar_sqrt should be zero
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    alphas_bar_sqrt = alphas_cumprod.sqrt()

def test_basic_random_betas():
    # Test with random betas in (0, 0.2)
    torch.manual_seed(0)
    betas = torch.rand(8) * 0.2
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    alphas_bar_sqrt = alphas_cumprod.sqrt()

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

def test_edge_single_element():
    # Test with a single beta value
    betas = torch.tensor([0.05])
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    # The last alpha_bar_sqrt should be zero
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    alphas_bar_sqrt = alphas_cumprod.sqrt()

def test_edge_all_zeros():
    # Test with all zeros (no noise)
    betas = torch.zeros(4)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    # The last alpha_bar_sqrt should be zero
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    alphas_bar_sqrt = alphas_cumprod.sqrt()

def test_edge_all_ones():
    # Test with all ones (full noise)
    betas = torch.ones(3)
    # This is a pathological case: 1 - betas = 0, so cumprod is 0, sqrt is 0
    # Should not produce NaNs or infs
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    alphas_bar_sqrt = alphas_cumprod.sqrt()

def test_edge_close_to_one():
    # Test with betas very close to 1
    betas = torch.full((5,), 0.999)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    alphas_bar_sqrt = alphas_cumprod.sqrt()

def test_edge_close_to_zero():
    # Test with betas very close to 0
    betas = torch.full((5,), 1e-8)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    alphas_bar_sqrt = alphas_cumprod.sqrt()


def test_edge_requires_grad():
    # Test with a tensor that requires grad
    betas = torch.linspace(0.01, 0.1, 6, requires_grad=True)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    alphas_bar_sqrt = alphas_cumprod.sqrt()

def test_edge_dtype_float64():
    # Test with float64 dtype
    betas = torch.linspace(0.01, 0.1, 7, dtype=torch.float64)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    alphas_bar_sqrt = alphas_cumprod.sqrt()

def test_edge_empty_tensor():
    # Test with an empty tensor (should raise an error)
    betas = torch.tensor([])
    with pytest.raises(IndexError):
        rescale_zero_terminal_snr(betas)

def test_edge_nan_input():
    # Test with NaN in input
    betas = torch.tensor([0.01, float('nan'), 0.02])
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output

def test_edge_inf_input():
    # Test with inf in input
    betas = torch.tensor([0.01, float('inf'), 0.02])
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output

def test_edge_negative_betas():
    # Test with negative betas (invalid, but should not crash)
    betas = torch.tensor([0.01, -0.1, 0.02])
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output

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

def test_large_scale_1d():
    # Test with a large 1D tensor (1000 elements, <100MB)
    betas = torch.linspace(0.0001, 0.02, 1000)
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    alphas_bar_sqrt = alphas_cumprod.sqrt()

def test_large_scale_random():
    # Test with a large random tensor (1000 elements, random values in (0, 0.2))
    torch.manual_seed(42)
    betas = torch.rand(1000) * 0.2
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    alphas_bar_sqrt = alphas_cumprod.sqrt()

def test_large_scale_performance():
    # Test with a large tensor and ensure it runs in reasonable time
    betas = torch.linspace(0.0001, 0.02, 1000)
    import time
    start = time.time()
    codeflash_output = rescale_zero_terminal_snr(betas); rescaled = codeflash_output
    elapsed = time.time() - start
    alphas = 1.0 - rescaled
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    alphas_bar_sqrt = alphas_cumprod.sqrt()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.