import libcst as cst

from core.languages.python.testgen.postprocessing.tensor_limit import modify_tensors


def test_modify_tensors_large_tensor() -> None:
    """Test that modify_tensors reduces the size of large tensors with scaling factor."""
    code = """
def test_function():
    x = torch.zeros(1000, 1000, 1000)  # Above threshold
    return x
"""
    module = cst.parse_module(code)
    result = modify_tensors(module, threshold=100000)

    # Convert back to code and check
    modified_code = result.code
    expected_code = """
def test_function():
    x = torch.zeros(46, 46, 46)  # Above threshold
    return x
"""
    assert modified_code.strip() == expected_code.strip()


def test_modify_tensors_small_tensor() -> None:
    """Test that modify_tensors doesn't change small tensors."""
    code = """
def test_function():
    x = torch.rand(2, 3)  # Below threshold
    return x
"""
    module = cst.parse_module(code)
    result = modify_tensors(module, threshold=100)

    # Convert back to code and check
    modified_code = result.code
    expected_code = """
def test_function():
    x = torch.rand(2, 3)  # Below threshold
    return x
"""
    assert modified_code.strip() == expected_code.strip()


def test_modify_tensors_mixed_tensors() -> None:
    """Test that modify_tensors handles code with both large and small tensors correctly."""
    code = """
def test_function():
    x = torch.zeros(1000, 1000, 1000)  # Above threshold
    y = torch.rand(2, 3)  # Below threshold
    return x, y
"""
    module = cst.parse_module(code)
    result = modify_tensors(module, threshold=100000)

    # Convert back to code and check
    modified_code = result.code
    expected_code = """
def test_function():
    x = torch.zeros(46, 46, 46)  # Above threshold
    y = torch.rand(2, 3)  # Below threshold
    return x, y
"""
    assert modified_code.strip() == expected_code.strip()


def test_modify_tensors_multiple_torch_functions() -> None:
    """Test that modify_tensors handles different torch tensor creation functions."""
    code = """
def test_function():
    x = torch.zeros(1000, 1000, 1000)  # Above threshold
    y = torch.ones(500, 500, 500)  # Above threshold
    z = torch.rand(800, 800, 800)  # Above threshold
    return x, y, z
"""
    module = cst.parse_module(code)
    result = modify_tensors(module, threshold=100000)

    # Convert back to code and check
    modified_code = result.code
    expected_code = """
def test_function():
    x = torch.zeros(46, 46, 46)  # Above threshold
    y = torch.ones(46, 46, 46)  # Above threshold
    z = torch.rand(46, 46, 46)  # Above threshold
    return x, y, z
"""
    assert modified_code.strip() == expected_code.strip()


def test_modify_tensors_with_expressions() -> None:
    """Test that modify_tensors can evaluate expressions in tensor dimensions."""
    code = """
def test_function():
    x = torch.zeros(10**3, 10**3, 10**3)  # Above threshold using power
    y = torch.ones(500+500, 500+500, 500+500)  # Above threshold using addition
    return x, y
"""
    module = cst.parse_module(code)
    result = modify_tensors(module, threshold=100000)

    # Convert back to code and check
    modified_code = result.code
    expected_code = """
def test_function():
    x = torch.zeros(46, 46, 46)  # Above threshold using power
    y = torch.ones(46, 46, 46)  # Above threshold using addition
    return x, y
"""
    assert modified_code.strip() == expected_code.strip()


def test_modify_tensors_non_torch_functions() -> None:
    """Test that modify_tensors doesn't affect non-torch functions."""
    code = """
def test_function():
    x = np.zeros((1000, 1000, 1000))  # Non-torch function
    y = create_tensor(1000, 1000, 1000)  # Non-torch function
    return x, y
"""
    module = cst.parse_module(code)
    result = modify_tensors(module, threshold=100000)

    # Convert back to code and check
    modified_code = result.code
    expected_code = """
def test_function():
    x = np.zeros((1000, 1000, 1000))  # Non-torch function
    y = create_tensor(1000, 1000, 1000)  # Non-torch function
    return x, y
"""
    assert modified_code.strip() == expected_code.strip()


def test_modify_tensors_with_default_threshold() -> None:
    """Test that modify_tensors uses the default threshold correctly."""
    code = """
def test_function():
    # Default threshold is usually 10,000,000
    small_tensor = torch.zeros(10, 10, 10)  # 1,000 elements
    large_tensor = torch.zeros(1000, 1000, 1000)  # 1,000,000,000 elements
    return small_tensor, large_tensor
"""
    module = cst.parse_module(code)
    # Use default threshold (assuming it's 10,000,000)
    result = modify_tensors(module)

    # Convert back to code and check
    modified_code = result.code
    expected_code = """
def test_function():
    # Default threshold is usually 10,000,000
    small_tensor = torch.zeros(10, 10, 10)  # 1,000 elements
    large_tensor = torch.zeros(215, 215, 215)  # 1,000,000,000 elements
    return small_tensor, large_tensor
"""
    assert modified_code.strip() == expected_code.strip()


def test_modify_tensors_asymmetric_dimensions() -> None:
    """Test that modify_tensors correctly scales tensors with asymmetric dimensions."""
    code = """
def test_function():
    x = torch.zeros(10, 1000, 1000)  # Above threshold, asymmetric
    return x
"""
    module = cst.parse_module(code)
    result = modify_tensors(module, threshold=100000)

    # Convert back to code and check
    modified_code = result.code
    expected_code = """
def test_function():
    x = torch.zeros(2, 215, 215)  # Above threshold, asymmetric
    return x
"""
    assert modified_code.strip() == expected_code.strip()


def test_modify_tensors_very_large_dimensions() -> None:
    """Test that modify_tensors handles very large tensor dimensions correctly."""
    code = """
def test_function():
    x = torch.zeros(10000, 10000)  # Very large 2D tensor
    return x
"""
    module = cst.parse_module(code)
    result = modify_tensors(module, threshold=100000)

    # Convert back to code and check
    modified_code = result.code
    expected_code = """
def test_function():
    x = torch.zeros(316, 316)  # Very large 2D tensor
    return x
"""
    assert modified_code.strip() == expected_code.strip()


def test_modify_tensors_one_dimension() -> None:
    """Test that modify_tensors handles one-dimensional tensors correctly."""
    code = """
def test_function():
    x = torch.zeros(1000000)  # Large 1D tensor
    return x
"""
    module = cst.parse_module(code)
    result = modify_tensors(module, threshold=100000)

    # Convert back to code and check
    modified_code = result.code
    expected_code = """
def test_function():
    x = torch.zeros(100000)  # Large 1D tensor
    return x
"""
    assert modified_code.strip() == expected_code.strip()


def test_modify_tensors_scaling_preserves_dimension_ratios() -> None:
    """Test that modify_tensors preserves the ratio between dimensions when scaling."""
    code = """
def test_function():
    x = torch.zeros(100, 200, 300)  # Above threshold with different dimensions
    return x
"""
    module = cst.parse_module(code)
    result = modify_tensors(module, threshold=100000)

    # Convert back to code and check
    modified_code = result.code
    expected_code = """
def test_function():
    x = torch.zeros(25, 51, 76)  # Above threshold with different dimensions
    return x
"""
    assert modified_code.strip() == expected_code.strip()


def test_modify_tensors_randn() -> None:
    """Test that modify_tensors handles torch.randn (commonly used for random normal tensors)."""
    code = """
def test_function():
    x = torch.randn(1000, 1000, 1000)  # Above threshold
    return x
"""
    module = cst.parse_module(code)
    result = modify_tensors(module, threshold=100000)

    modified_code = result.code
    expected_code = """
def test_function():
    x = torch.randn(46, 46, 46)  # Above threshold
    return x
"""
    assert modified_code.strip() == expected_code.strip()


def test_modify_tensors_empty() -> None:
    """Test that modify_tensors handles torch.empty."""
    code = """
def test_function():
    x = torch.empty(1000, 1000, 1000)  # Above threshold
    return x
"""
    module = cst.parse_module(code)
    result = modify_tensors(module, threshold=100000)

    modified_code = result.code
    expected_code = """
def test_function():
    x = torch.empty(46, 46, 46)  # Above threshold
    return x
"""
    assert modified_code.strip() == expected_code.strip()


def test_modify_tensors_preserves_dtype_keyword() -> None:
    """Test that modify_tensors preserves dtype keyword argument."""
    code = """
def test_function():
    x = torch.randn(1000, 1000, dtype=torch.float64)  # Above threshold with dtype
    return x
"""
    module = cst.parse_module(code)
    result = modify_tensors(module, threshold=100000)

    modified_code = result.code
    expected_code = """
def test_function():
    x = torch.randn(316, 316, dtype=torch.float64)  # Above threshold with dtype
    return x
"""
    assert modified_code.strip() == expected_code.strip()


def test_modify_tensors_preserves_device_keyword() -> None:
    """Test that modify_tensors preserves device keyword argument."""
    code = """
def test_function():
    x = torch.zeros(1000, 1000, 1000, device='cpu')  # Above threshold with device
    return x
"""
    module = cst.parse_module(code)
    result = modify_tensors(module, threshold=100000)

    modified_code = result.code
    expected_code = """
def test_function():
    x = torch.zeros(46, 46, 46, device='cpu')  # Above threshold with device
    return x
"""
    assert modified_code.strip() == expected_code.strip()


def test_modify_tensors_preserves_multiple_keywords() -> None:
    """Test that modify_tensors preserves multiple keyword arguments."""
    code = """
def test_function():
    g = torch.Generator().manual_seed(42)
    x = torch.randn(1000, 1000, dtype=torch.float64, device='cpu', generator=g)
    return x
"""
    module = cst.parse_module(code)
    result = modify_tensors(module, threshold=100000)

    modified_code = result.code
    expected_code = """
def test_function():
    g = torch.Generator().manual_seed(42)
    x = torch.randn(316, 316, dtype=torch.float64, device='cpu', generator=g)
    return x
"""
    assert modified_code.strip() == expected_code.strip()


def test_modify_tensors_small_with_keywords_unchanged() -> None:
    """Test that small tensors with keywords are not modified."""
    code = """
def test_function():
    x = torch.randn(10, 10, dtype=torch.float64)  # Below threshold
    return x
"""
    module = cst.parse_module(code)
    result = modify_tensors(module, threshold=100000)

    modified_code = result.code
    # Should be unchanged
    assert modified_code.strip() == code.strip()
