from collections.abc import Callable
import os
import functools

import numpy as np
import pytest
import torch
from torch.testing import assert_close

if hasattr(torch.nn.functional, "scaled_mm"):
    from torch.nn.functional import ScalingType, SwizzleType

import thunder
import thunder.core.devices as devices
import thunder.core.dtypes as dtypes
from thunder.core.pytree import tree_flatten, tree_map
from thunder.tests.framework import assert_closer, ops, run_snippet, requiresCUDA
from thunder.tests.opinfos import OpInfo, SampleInput, opinfos
import thunder.tests.bf16

#
# Generic test templates for all operators
#


# NOTE err_msg_match=None will match any error message
def snippet_errors(op, sample, ex_type, err_msg_match=None):
    with pytest.raises(ex_type, match=err_msg_match):
        op(*sample.args, **sample.kwargs)


@ops(tuple(op for op in opinfos if op.error_input_generator is not None))
def test_errors(op, device, dtype, executor, comp):
    for sample, ex_type, err_msg in op.error_inputs(device):
        result = run_snippet(snippet_errors, op, device, None, executor.make_callable(op.op), sample, ex_type, err_msg)
        if result is not None:
            return result


def assert_consistency_of_compiletime_and_runtime(thunder_op, thunder_result):
    from thunder.core.baseutils import sequencify

    __tracebackhide__ = True

    extrace: thunder.TraceCtx = thunder.last_traces(thunder_op)[-1]
    for runtime, compiletime in zip(
        tree_flatten(sequencify(thunder_result))[0],
        tree_flatten(sequencify(extrace.output))[0],
    ):
        if isinstance(compiletime, thunder.TensorProxy):
            torch_device = devices.to_torch_device(compiletime.device)
            torch_dtype = dtypes.to_torch_dtype(compiletime.dtype)
            if not (
                (r_shape := tuple(runtime.shape)) == (c_shape := tuple(compiletime.shape))
                and runtime.device == torch_device
                and runtime.dtype == torch_dtype
            ):
                msg = (
                    f"Runtime output has shape of `{r_shape}`, device of `{runtime.device}`, and dtype of `{runtime.dtype}` "
                    f"but compiletime output has shape of `{c_shape}`, device of `{torch_device}`, and dtype of `{torch_dtype}`"
                )
                raise RuntimeError(msg)


# Snippets run a single test using a single sample
# TODO: should snippets be able to access the original opinfo? -- No?
# TODO: revisit atol/rtol, maybe be more selective about which ops need a more permissive check
def snippet_torch_consistency(op: OpInfo, torch_op, sample: SampleInput, comp: Callable):
    args, kwargs = sample.args, sample.kwargs

    thunder_result = op(*args, **kwargs)
    torch_result = torch_op(*args, **kwargs)

    # TODO Review how thunder.jit returns Exception information
    if isinstance(thunder_result, Exception):
        raise thunder_result

    assert_consistency_of_compiletime_and_runtime(op, thunder_result)

    # Try checking strictly, if that does not work, check against reference.
    try:
        comp(thunder_result, torch_result)
    except AssertionError:

        def upcast(x):
            if isinstance(x, torch.Tensor) and x.is_floating_point():
                return x.to(torch.double)

            # Some torch APIs (e.g. cumsum) take a `dtype` argument and use it
            # as the **compute** dtype. Replace that with torch.double so the
            # reference result is more likely to be computed in double.
            if isinstance(x, torch.dtype) and x.is_floating_point:
                return torch.double

            return x

        reference_args = tree_map(upcast, args)
        reference_kwargs = tree_map(upcast, kwargs)
        reference_result = torch_op(*reference_args, **reference_kwargs)

        assert_closer(
            reference=(reference_result,), candidate=(thunder_result,), competitor=(torch_result,), comparator=comp
        )


# TODO consider structuring tests like this to be autogenerated
#   using a snippet and an "extractor" that constructs the args and kwargs for the snippet
# TODO The name of this test is misleading as it may test operators from a variety of languages,
#   maybe we should cut it up so developers can test just torch operators or just core lang operators
# TODO Extend this test with some reproducible randomness (maybe using hypothesis)
# TODO Remove the atol and rtol defaults and rely on the given comparator to set them
@ops(tuple(op for op in opinfos if op.torch_reference is not None))
def test_core_vs_torch_consistency(op, device: str, dtype: dtypes.dtype, executor, comp):
    if dtypes.is_complex_dtype(dtype) and not op.instantiate_complex_tests:
        pytest.skip("Skipping complex operator tests in CI for speed")
    if (
        torch.device(device).type == "cuda"
        and dtype is dtypes.bfloat16
        and not thunder.tests.bf16.device_supports_bf16(device)
    ):
        pytest.skip("Your CUDA device does not support bfloat16")

    # Check if a specific sample index is requested via environment variable
    target_sample_idx = int(os.environ.get("THUNDER_OPTEST_SAMPLE_INDEX", -1))

    for i, sample in enumerate(op.sample_inputs(device, dtype)):
        # Skip samples that don't match the requested index
        if target_sample_idx >= 0 and i != target_sample_idx:
            continue

        sample_comp = sample.comp if sample.comp is not None else comp

        tfn = thunder.jit(
            op.op,
            executors=executor.executors_list(),
            cache="no caching",
            disable_torch_autograd=True,
        )

        path = os.path.relpath(__file__)
        repro_cmd = (
            f"command to reproduce the error: THUNDER_OPTEST_SAMPLE_INDEX={i} "
            f"pytest {path}::test_core_vs_torch_consistency_{op.name}_{executor.name}_{device}_{dtype}"
        )

        try:
            result = run_snippet(
                snippet_torch_consistency,
                op,
                device,
                dtype,
                tfn,
                op.torch_reference,
                sample,
                lambda a, b, **kwargs: sample_comp(a, b, equal_nan=True, **kwargs),
            )
        except Exception as e:
            if repro_cmd not in str(e):
                raise type(e)(f"{str(e)}\n{repro_cmd}") from e
            else:
                raise
        else:
            # This runs only if no exception was raised
            # See [NOTE] dynamo reset
            if any("torchcompile" in ex.name for ex in executor.executors_list()):
                torch._dynamo.reset()

            if result is not None:
                # Check if result is a tuple of exceptions and metadata
                if isinstance(result, tuple) and len(result) > 0 and isinstance(result[0], Exception):
                    # Add repro command to exception message
                    ex = result[0]
                    if repro_cmd not in str(ex):
                        new_ex = type(ex)(f"{str(ex)}\n{repro_cmd}")
                        # Preserve the original traceback
                        new_ex.__traceback__ = ex.__traceback__
                        result = (new_ex,) + result[1:]
                return result


def snippet_numpy_consistency(op: OpInfo, np_op, sample: SampleInput, comp: Callable):
    np_sample = sample.numpy()

    thunder_result = op(*sample.args, **sample.kwargs)
    np_result = np_op(*np_sample.args, **np_sample.kwargs)

    # Converts NumPy results to PyTorch.
    # NOTE This assumes PyTorch will return tensors where NumPy is aggressive about returning `np.number` objects.
    def convert_to_torch(x):
        if not isinstance(x, (np.ndarray, np.number, np.bool_)):
            return x

        if isinstance(x, (np.number, np.bool_)):
            return torch.tensor(x, device=thunder_result.device)
        elif x.shape == ():
            return torch.tensor(x.item(), device=thunder_result.device)
        else:
            return torch.asarray(x, device=thunder_result.device)

    np_result = tree_map(convert_to_torch, np_result)

    comp(thunder_result, np_result)


@ops(tuple(op for op in opinfos if op.numpy_reference is not None))
def test_core_vs_numpy_consistency(op: OpInfo, device: str, dtype: dtypes.dtype, executor, comp):
    if dtypes.is_complex_dtype(dtype):
        pytest.skip("Skipping complex operator tests in CI for speed")
    if dtype == dtypes.complex32:
        pytest.skip("NumPy does not support complex32")
    if dtype == dtypes.bfloat16:
        pytest.skip("NumPy does not support bfloat16")

    for sample in op.sample_inputs(device, dtype):
        comp = sample.comp if sample.comp is not None else comp

        result = run_snippet(
            snippet_numpy_consistency,
            op,
            device,
            dtype,
            executor.make_callable(op.op),
            op.numpy_reference,
            sample,
            # NOTE dtype is intentionally not checked because NumPy sometimes has slight dtype variances
            lambda a, b: comp(a, b, equal_nan=True, check_dtype=False),
        )
        if result is not None:
            return result


def test_interpolate_nearest_vs_nearest_exact():
    t0 = torch.tensor([[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]], dtype=torch.float16)

    def foo(mode):
        t1 = torch.nn.functional.interpolate(
            t0,
            scale_factor=1.5,
            mode=mode,
        )
        return t1

    def bar(mode):
        t1 = torch.nn.functional.interpolate(
            t0,
            size=4,
            mode=mode,
        )
        return t1

    tfoo = thunder.jit(foo)
    assert not torch.equal(tfoo("nearest"), tfoo("nearest-exact"))
    tbar = thunder.jit(bar)
    assert not torch.equal(tbar("nearest"), tbar("nearest-exact"))


def test_notimplemented_interpolate_align():
    def foo():
        t0 = torch.randn((22, 288, 15, 20), dtype=torch.float16)
        t1 = torch.nn.functional.interpolate(t0, scale_factor=2.0, mode="nearest", align_corners=True)
        return t1

    with pytest.raises(NotImplementedError, match="not yet support"):
        tfoo = thunder.jit(foo)
        tfoo()


def test_notimplemented_interpolate_recompute_scale():
    def foo():
        t0 = torch.randn((22, 288, 15, 20), dtype=torch.float16)
        t1 = torch.nn.functional.interpolate(t0, scale_factor=2.0, mode="nearest", recompute_scale_factor=True)
        return t1

    with pytest.raises(NotImplementedError, match="not yet support"):
        tfoo = thunder.jit(foo)
        tfoo()


def test_notimplemented_interpolate_antialias():
    def foo():
        t0 = torch.randn((22, 288, 15, 20), dtype=torch.float16)
        t1 = torch.nn.functional.interpolate(
            t0,
            scale_factor=2.0,
            mode="nearest",
            antialias=True,
        )
        return t1

    with pytest.raises(NotImplementedError, match="not yet support"):
        tfoo = thunder.jit(foo)
        tfoo()


def test_notimplemented_interpolate_modes():
    def foo(mode):
        t0 = torch.randn((22, 288, 15, 20), dtype=torch.float16)
        t1 = torch.nn.functional.interpolate(
            t0,
            scale_factor=2.0,
            mode=mode,
        )
        return t1

    tfoo = thunder.jit(foo)
    for mode in ["linear", "bicubic", "trilinear", "area"]:
        match = (
            f"only modes 'nearest', 'nearest-exact' and 'bilinear' are supported at the moment, but got mode='{mode}'"
        )
        with pytest.raises(NotImplementedError, match=match):
            tfoo(mode)


@pytest.mark.parametrize("requires_grad", (True, False))
def test_setitem(requires_grad):
    def _test_forward_and_backward(fn, a, value):
        a_ref = a.detach().clone()
        a_ref.requires_grad_(a.requires_grad)

        if isinstance(value, torch.Tensor):
            value_ref = value.detach().clone()
            value_ref.requires_grad_(value.requires_grad)
        else:
            value_ref = value

        out_ref = fn(a_ref, value_ref)
        jf = thunder.jit(fn)
        out = jf(a, value)
        assert_close(a, a_ref)
        assert_close(out, out_ref)

        if requires_grad:
            g = torch.randn_like(out)
            inputs = (a, value) if isinstance(value, torch.Tensor) else (a,)
            actual_grad = torch.autograd.grad(out, inputs, g)

            inputs_ref = (a_ref, value_ref) if isinstance(value, torch.Tensor) else (a_ref,)
            expected_grad = torch.autograd.grad(out_ref, inputs_ref, g)
            assert_close(actual_grad, expected_grad)

    def clone_if_requires_grad(a):
        if requires_grad:
            # Withou the clone
            # PyTorch eager errors with
            # `RuntimeError: a view of a leaf Variable that requires grad is being used in an in-place operation.`
            # and thunder has silent correctness issue - https://github.com/Lightning-AI/lightning-thunder/issues/1284
            return a.clone()
        return a

    def fn(a, value):
        a = clone_if_requires_grad(a)
        a[:3] = value
        return a * 2

    # set value: scalar
    _test_forward_and_backward(fn, torch.randn(5, requires_grad=requires_grad), 2.0)

    # set value: tensor which needs to be broadcasted
    _test_forward_and_backward(
        fn, torch.randn(5, requires_grad=requires_grad), torch.tensor(2.0, requires_grad=requires_grad)
    )

    def bcast_fn(a, value):
        a = clone_if_requires_grad(a)
        a[..., :3] = value
        return a * 2

    _test_forward_and_backward(
        bcast_fn, torch.randn(5, 3, 5, requires_grad=requires_grad), torch.randn(1, 3, requires_grad=requires_grad)
    )

    # set value: tensor of same rank
    _test_forward_and_backward(
        fn, torch.randn(5, requires_grad=requires_grad), torch.tensor([1.0, 2.0, 3.0], requires_grad=requires_grad)
    )


@requiresCUDA
def test_double_setitem():
    query_states = torch.zeros((1, 40, 16, 96), dtype=torch.bfloat16, device="cuda:0", requires_grad=False)
    q_nope = torch.ones((1, 40, 16, 64), dtype=torch.bfloat16, device="cuda:0", requires_grad=True)
    q_pe_1 = torch.zeros((1, 40, 16, 32), dtype=torch.bfloat16, device="cuda:0", requires_grad=True)

    @thunder.jit
    def computation(query_states, q_nope, q_pe_1):
        # Perform the in-place setitem operation
        query_states[:, :, :, :64] = q_nope
        query_states[:, :, :, 64:] = q_pe_1
        return None

    computation(query_states, q_nope, q_pe_1)
    assert query_states.sum() > 1


# TODO: Add random operator support to OpInfo
# https://github.com/Lightning-AI/lightning-thunder/issues/1163
@requiresCUDA
def test_exponential():
    def fn(a):
        return a.exponential_(1)

    size = 10
    seed = 1234

    # on cpu, aten.exponential_ is not decomposed to ops used in Thunder exponential_.
    with torch.device("cuda"):
        a_ref = torch.ones(size)
        b_ref = torch.ones((size, size, size))
        torch.manual_seed(seed)
        a_ref = fn(a_ref)
        b_ref = fn(b_ref)

        a = torch.ones(size)
        b = torch.ones((size, size, size))
        torch.manual_seed(seed)

        # nvfuser fuses prims.uniform, which is used by our exponential resulting in differing numerics.
        jf = thunder.jit(fn, executors={})
        a = jf(a)
        b = jf(b)

        assert_close(a, a_ref)
        assert_close(b, b_ref)


def _cuda_version_tuple() -> tuple[int, int] | None:
    if torch.version.cuda is None:
        return None
    parts = torch.version.cuda.split(".")
    try:
        major = int(parts[0])
        minor = int(parts[1]) if len(parts) > 1 else 0
        return major, minor
    except ValueError:
        return None


def _require_scaled_mm(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        if not hasattr(torch.nn.functional, "scaled_mm"):
            pytest.skip("torch.nn.functional.scaled_mm is not found in this PyTorch")
        return fn(*args, **kwargs)

    return wrapper


def _ensure_fp8_tensorwise(device: torch.device) -> None:
    if torch.cuda.get_device_capability(device) < (8, 9):
        pytest.skip("scaled_mm tensor-wise support requires SM89 or newer")


def _require_fp8_tensorwise(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        device = torch.device("cuda")
        _ensure_fp8_tensorwise(device)
        return fn(*args, **kwargs)

    return wrapper


def _require_fp8_rowwise(device: torch.device) -> None:
    _ensure_fp8_tensorwise(device)
    if torch.cuda.get_device_capability(device) < (9, 0):
        pytest.skip("row-wise scaled_mm requires SM90 or newer")
    cuda_version = _cuda_version_tuple()
    if cuda_version is not None and cuda_version < (12, 9):
        pytest.skip("row-wise scaled_mm requires CUDA 12.9 or newer")


def _require_fp8_blockwise(device: torch.device) -> None:
    _require_fp8_rowwise(device)


# Adapted from https://github.com/pytorch/pytorch/blob/b4403bfc62ca97eec554cdf815baab1fe93057d9/test/test_scaled_matmul_cuda.py#L645-L659
@requiresCUDA
@_require_fp8_tensorwise
@_require_scaled_mm
def test_scaled_mm_tensorwise_matches_torch():
    device = torch.device("cuda")

    def reference_fn(mat_a, mat_b, scale_a, scale_b):
        return torch.nn.functional.scaled_mm(
            mat_a,
            mat_b,
            scale_a,
            ScalingType.TensorWise,
            scale_b,
            ScalingType.TensorWise,
            swizzle_a=SwizzleType.NO_SWIZZLE,
            swizzle_b=SwizzleType.NO_SWIZZLE,
            output_dtype=torch.bfloat16,
        )

    M, K, N = 16, 32, 16
    mat_a = torch.randn(M, K, device=device, dtype=torch.float32)
    mat_b = torch.randn(K, N, device=device, dtype=torch.float32)
    mat_a_lp = mat_a.to(torch.float8_e4m3fn)
    mat_b_lp = mat_b.to(torch.float8_e4m3fn)
    scale_a = torch.tensor(1.0, device=device, dtype=torch.float32)
    scale_b = torch.tensor(1.0, device=device, dtype=torch.float32)

    try:
        expected = reference_fn(mat_a_lp, mat_b_lp, scale_a, scale_b)
    except (NotImplementedError, RuntimeError) as exc:
        pytest.skip(str(exc))

    jf = thunder.jit(reference_fn)
    result = jf(mat_a_lp, mat_b_lp, scale_a, scale_b)
    assert_close(result, expected)


# Adapted from https://github.com/pytorch/pytorch/blob/b4403bfc62ca97eec554cdf815baab1fe93057d9/test/test_scaled_matmul_cuda.py#L862-L910
@requiresCUDA
@_require_fp8_tensorwise
@_require_scaled_mm
def test_scaled_mm_matches_scaled_data():
    device = torch.device("cuda")

    def quantize_to_fp8(tensor):
        dtype = torch.float8_e4m3fn
        max_val = torch.finfo(dtype).max
        amax = tensor.abs().max()
        encode = (max_val / torch.clamp(amax, min=1e-12)).to(torch.float32)
        quant = torch.clamp(tensor * encode, min=-max_val, max=max_val).to(dtype)
        decode = encode.reciprocal()
        return quant, decode, encode

    def scaled_mm_fp8(mat_a, mat_b, scale_a, scale_b, *, out_dtype):
        return torch.nn.functional.scaled_mm(
            mat_a,
            mat_b,
            scale_a,
            ScalingType.TensorWise,
            scale_b,
            ScalingType.TensorWise,
            swizzle_a=SwizzleType.NO_SWIZZLE,
            swizzle_b=SwizzleType.NO_SWIZZLE,
            output_dtype=out_dtype,
        )

    M, K, N = 32, 64, 32
    mat_a = torch.randn(M, K, device=device, dtype=torch.float32)
    mat_b_base = torch.randn(N, K, device=device, dtype=torch.float32)

    mat_a_lp, decode_a, encode_a = quantize_to_fp8(mat_a)
    mat_b_lp_pre, decode_b, encode_b = quantize_to_fp8(mat_b_base)
    # To use cublaslt, the second matrix needs to be column-major.
    mat_b_lp = mat_b_lp_pre.t()

    try:
        reference = scaled_mm_fp8(mat_a_lp, mat_b_lp, decode_a, decode_b, out_dtype=torch.float32)
    except (NotImplementedError, RuntimeError) as exc:
        pytest.skip(str(exc))

    jf = thunder.jit(lambda a, b, sa, sb: scaled_mm_fp8(a, b, sa, sb, out_dtype=torch.float32))
    thunder_out = jf(mat_a_lp, mat_b_lp, decode_a, decode_b)

    assert_close(thunder_out, reference)


@requiresCUDA
@_require_scaled_mm
def test_scaled_mm_rowwise_matches_torch():
    device = torch.device("cuda")
    _require_fp8_rowwise(device)

    def reference_fn(mat_a, mat_b, scale_a, scale_b):
        return torch.nn.functional.scaled_mm(
            mat_a,
            mat_b,
            scale_a,
            ScalingType.RowWise,
            scale_b,
            ScalingType.RowWise,
            swizzle_a=SwizzleType.NO_SWIZZLE,
            swizzle_b=SwizzleType.NO_SWIZZLE,
            output_dtype=torch.bfloat16,
        )

    M, K, N = 16, 32, 16
    mat_a = torch.randn(M, K, device=device, dtype=torch.float32)
    mat_b_base = torch.randn(N, K, device=device, dtype=torch.float32)
    mat_a_lp = mat_a.to(torch.float8_e4m3fn)
    # To use cublaslt, the second matrix needs to be column-major.
    mat_b_lp = mat_b_base.to(torch.float8_e4m3fn).t()
    scale_a = torch.ones((M, 1), device=device, dtype=torch.float32)
    scale_b = torch.ones((1, N), device=device, dtype=torch.float32)

    try:
        expected = reference_fn(mat_a_lp, mat_b_lp, scale_a, scale_b)
    except (NotImplementedError, RuntimeError) as exc:
        pytest.skip(str(exc))

    jf = thunder.jit(reference_fn)
    result = jf(mat_a_lp, mat_b_lp, scale_a, scale_b)
    assert_close(result, expected)


@requiresCUDA
@_require_scaled_mm
def test_scaled_mm_rowwise_matches_scaled_data():
    device = torch.device("cuda")
    _require_fp8_rowwise(device)

    dtype_fp8 = torch.float8_e4m3fn
    max_val = torch.finfo(dtype_fp8).max

    def rowwise_quantize(tensor, *, dim):
        amax = tensor.abs().amax(dim=dim, keepdim=True)
        encode = (max_val / torch.clamp(amax, min=1e-12)).to(torch.float32)
        quant = torch.clamp(tensor * encode, min=-max_val, max=max_val).to(dtype_fp8)
        decode = encode.reciprocal()
        return quant, decode, encode

    def scaled_mm_rowwise(mat_a, mat_b, scale_a, scale_b, *, out_dtype):
        return torch.nn.functional.scaled_mm(
            mat_a,
            mat_b,
            scale_a,
            ScalingType.RowWise,
            scale_b,
            ScalingType.RowWise,
            swizzle_a=SwizzleType.NO_SWIZZLE,
            swizzle_b=SwizzleType.NO_SWIZZLE,
            output_dtype=out_dtype,
        )

    M, K, N = 32, 64, 32
    mat_a = torch.randn(M, K, device=device, dtype=torch.bfloat16)
    mat_b = torch.randn(K, N, device=device, dtype=torch.bfloat16)

    mat_a_lp, decode_a, encode_a = rowwise_quantize(mat_a.to(torch.float32), dim=1)
    mat_b_lp, decode_b, encode_b = rowwise_quantize(mat_b.to(torch.float32), dim=0)

    try:
        reference = scaled_mm_rowwise(mat_a_lp, mat_b_lp, decode_a, decode_b, out_dtype=torch.bfloat16)
    except (NotImplementedError, RuntimeError) as exc:
        pytest.skip(str(exc))

    jf = thunder.jit(lambda a, b, sa, sb: scaled_mm_rowwise(a, b, sa, sb, out_dtype=torch.bfloat16))
    thunder_out = jf(mat_a_lp, mat_b_lp, decode_a, decode_b)

    reference_f32 = reference.to(torch.float32)
    thunder_out_f32 = thunder_out.to(torch.float32)

    assert_close(thunder_out_f32, reference_f32, atol=3e-2, rtol=3e-2)


def _blockwise_quantize(tensor: torch.Tensor, block_rows: int, block_cols: int) -> tuple[torch.Tensor, torch.Tensor]:
    dtype_fp8 = torch.float8_e4m3fn
    max_val = torch.finfo(dtype_fp8).max

    M, K = tensor.shape
    assert M % block_rows == 0 and K % block_cols == 0

    reshaped = tensor.reshape(M // block_rows, block_rows, K // block_cols, block_cols)
    amax = reshaped.abs().amax(dim=(1, 3), keepdim=True)
    encode = (max_val / torch.clamp(amax, min=1e-12)).to(torch.float32)
    quant = torch.clamp(reshaped * encode, min=-max_val, max=max_val).to(dtype_fp8)

    return quant.reshape(M, K), encode.reshape(M // block_rows, K // block_cols).to(tensor.device)


@requiresCUDA
@_require_scaled_mm
@pytest.mark.parametrize("output_dtype", [torch.bfloat16])
@pytest.mark.parametrize("lhs_block,rhs_block", [(1, 1), (128, 1), (1, 128)])
def test_scaled_mm_blockwise_matches_torch(output_dtype, lhs_block, rhs_block):
    device = torch.device("cuda")
    _require_fp8_blockwise(device)

    M, K, N = 256, 256, 256
    mat_a = torch.randn(M, K, device=device, dtype=output_dtype).pow(3)
    mat_b_rows = torch.randn(N, K, device=device, dtype=output_dtype).pow(3)

    mat_a_lp, encode_a = _blockwise_quantize(mat_a.to(torch.float32), lhs_block, 128)
    mat_b_lp_rows, encode_b = _blockwise_quantize(mat_b_rows.to(torch.float32), rhs_block, 128)
    mat_b_lp = mat_b_lp_rows.t().contiguous()

    scale_a = encode_a.reciprocal().contiguous()
    scale_b = encode_b.reciprocal().t().contiguous()

    recipe_map = {
        1: ScalingType.BlockWise1x128,
        128: ScalingType.BlockWise128x128,
    }

    try:
        expected = torch.nn.functional.scaled_mm(
            mat_a_lp,
            mat_b_lp,
            scale_a,
            recipe_map[lhs_block],
            scale_b,
            recipe_map[rhs_block],
            swizzle_a=SwizzleType.NO_SWIZZLE,
            swizzle_b=SwizzleType.NO_SWIZZLE,
            output_dtype=output_dtype,
        )
    except (RuntimeError, NotImplementedError, ValueError) as exc:
        pytest.skip(str(exc))

    fn = thunder.jit(
        lambda a, b, sa, sb: torch.nn.functional.scaled_mm(
            a,
            b,
            sa,
            recipe_map[lhs_block],
            sb,
            recipe_map[rhs_block],
            swizzle_a=SwizzleType.NO_SWIZZLE,
            swizzle_b=SwizzleType.NO_SWIZZLE,
            output_dtype=output_dtype,
        )
    )
    thunder_out = fn(mat_a_lp, mat_b_lp, scale_a, scale_b)
    assert_close(thunder_out, expected)


# https://github.com/Lightning-AI/lightning-thunder/issues/1857
def test_max_with_int():
    def f(x, ids):
        x + x
        return ids[0].max()

    x = torch.rand([2, 2], requires_grad=True)
    ids = torch.randint(0, 10, size=(1, 512))

    thunder.jit(f)(x, ids)


def test_ltorch_cumsum_result_dtype_for_int_input():
    def f(a):
        return torch.cumsum(a, dim=0)

    x = torch.randint(0, 128, (4,), dtype=torch.int32)
    jitted = thunder.jit(f)
    out = jitted(x)
    # runtime check
    assert out.dtype is torch.int64

    trc = thunder.last_traces(jitted)[0]
    bsym_of_cumsum = trc.bound_symbols[1]
    assert bsym_of_cumsum.output.dtype is dtypes.int64


def test_ltorch_maximum_result_dtype_for_scalar_tensors():
    def foo(a, b):
        return torch.maximum(a, b)

    x = torch.rand(2, 2, dtype=torch.bfloat16)
    y = torch.tensor(0.12345)

    jfoo = thunder.jit(foo)
    out = jfoo(x, y)

    assert out.dtype == torch.bfloat16

    trc = thunder.last_traces(jfoo)[-1]

    bsym = trc.bound_symbols[2]
    assert bsym.output.dtype is dtypes.bfloat16


def test_multi_dot_optimization():
    def fn(tensors):
        return torch.linalg.multi_dot(tensors)

    # three matrices

    a = torch.randn(10, 100)
    b = torch.randn(100, 10)
    c = torch.randn(10, 100)

    jfn = thunder.jit(fn)
    jfn([a, b, c])
    trc = thunder.last_traces(jfn)[-1]

    # make sure that there is no (100 x 100) intermediates
    for bsym in trc.bound_symbols:
        if bsym.sym.id == "matmul":
            for flat_out in bsym.flat_outs:
                assert flat_out.shape != (100, 100)

    jfn([a.T, b.T, c.T])
    trc2 = thunder.last_traces(jfn)[-1]

    # make sure that there is no (100 x 100) intermediates
    for bsym in trc2.bound_symbols:
        if bsym.sym.id == "matmul":
            for flat_out in bsym.flat_outs:
                assert flat_out.shape != (100, 100)

    # five matrices

    d = torch.randn(100, 10)
    e = torch.randn(10, 100)

    jfn = thunder.jit(fn)
    jfn([a, b, c, d, e])
    trc = thunder.last_traces(jfn)[-1]

    # make sure that there is no (100 x 100) intermediates
    for bsym in trc.bound_symbols:
        if bsym.sym.id == "matmul":
            for flat_out in bsym.flat_outs:
                assert flat_out.shape != (100, 100)

    jfn([a.T, b.T, c.T, d.T, e.T])
    trc2 = thunder.last_traces(jfn)[-1]

    # make sure that there is no (100 x 100) intermediates
    for bsym in trc2.bound_symbols:
        if bsym.sym.id == "matmul":
            for flat_out in bsym.flat_outs:
                assert flat_out.shape != (100, 100)


def test_softmax_stacklevel():
    def fn(a):
        return torch.nn.functional.softmax(a, -1, _stacklevel=5)

    jfn = thunder.jit(fn)
    a = torch.randn(5, 5, requires_grad=True)  # trigger grad transform
    assert_close(fn(a), jfn(a))


def test_div_exact():
    # Test that division with integer inputs and requires_grad tensor output works correctly
    def fn(a, b, c):
        indices = torch.div(a, b, rounding_mode="trunc")
        # this would throw an error if indices are not ints
        return c[indices]

    jfn = thunder.jit(fn)
    a = torch.randint(1, 5, (5,))
    b = torch.ones(5, dtype=torch.int32)
    c = torch.randn(5, 5, requires_grad=True)
    result_eager = fn(a, b, c)
    result_jit = jfn(a, b, c)
    assert_close(result_eager, result_jit)


def test_view_as_complex_with_original_tensor_aliasing():
    """Test that using both original tensor and its complex view with inplace ops works correctly.

    This test specifically exercises the numel check in update_aliases.py:replace_args_with_alias_map.

    When a function has inplace operations AND receives aliased tensors as arguments, Thunder's
    insert_alias_updates() is called, which in turn calls replace_args_with_alias_map().

    The numel check (if arg.numel != arg_to_replace.numel: continue) prevents Thunder from
    incorrectly trying to replace tensors that share storage but have different element counts,
    such as a real tensor (48 elements) and its complex view (24 elements).

    Without this check, Thunder would attempt to replace one with a reshaped version of the other,
    causing shape mismatch errors.
    """

    def func_with_inplace(real_tensor, complex_tensor):
        real_tensor.add_(1.0)
        real_result = real_tensor.sum()
        complex_result = complex_tensor.abs().sum()
        return real_result + complex_result

    real_tensor = torch.randn((4, 6, 2), device="cpu", dtype=torch.float32)

    jfunc = thunder.jit(func_with_inplace)

    # Clone inputs for comparison (since inplace modifies them)
    real_tensor_test = real_tensor.clone()
    complex_tensor_test = torch.view_as_complex(real_tensor_test)
    real_tensor_expected = real_tensor.clone()
    complex_tensor_expected = torch.view_as_complex(real_tensor_expected)

    result = jfunc(real_tensor_test, complex_tensor_test)
    expected = func_with_inplace(real_tensor_expected, complex_tensor_expected)

    assert_close(result, expected)


def test_view_as_real_with_original_tensor_aliasing():
    """Test that using both original complex tensor and its real view with inplace ops works correctly.

    This tests the opposite case from test_view_as_complex_with_original_tensor_aliasing.
    """

    def func_with_inplace(complex_tensor, real_tensor):
        complex_tensor.add_(1.0 + 0.5j)
        complex_result = complex_tensor.real.sum()
        real_result = real_tensor.sum()
        return complex_result + real_result

    complex_tensor = torch.randn((4, 6), device="cpu", dtype=torch.complex64)

    jfunc = thunder.jit(func_with_inplace)

    # Clone inputs for comparison (since inplace modifies them)
    complex_tensor_test = complex_tensor.clone()
    real_tensor_test = torch.view_as_real(complex_tensor_test)
    complex_tensor_expected = complex_tensor.clone()
    real_tensor_expected = torch.view_as_real(complex_tensor_expected)

    result = jfunc(complex_tensor_test, real_tensor_test)
    expected = func_with_inplace(complex_tensor_expected, real_tensor_expected)

    assert_close(result, expected)


def test_rope_complex_vs_real_implementations():
    """Compare RoPE implementations: real arithmetic vs complex views.

    Tests that apply_rotary_emb (real arithmetic) and apply_rotary_emb_complex (complex views)
    produce identical results. This validates that view_as_complex/view_as_real work correctly.
    """
    from thunder.tests.llama2_model import apply_rotary_emb, apply_rotary_emb_with_complex_views

    # Model dimensions
    batch_size = 2
    seq_len = 8
    n_heads = 4
    head_dim = 16

    # Create query and key tensors
    xq = torch.randn(batch_size, seq_len, n_heads, head_dim, dtype=torch.float32)
    xk = torch.randn(batch_size, seq_len, n_heads, head_dim, dtype=torch.float32)

    # Create frequency tensors matching precompute_freqs_cis format
    # freqs_cos and freqs_sin should have shape (seq_len, head_dim//2)
    freqs = 1.0 / (10000.0 ** (torch.arange(0, head_dim, 2).float() / head_dim))
    positions = torch.arange(seq_len, dtype=torch.float32)
    angles = torch.outer(positions, freqs)  # (seq_len, head_dim//2)
    freqs_cos = torch.cos(angles)
    freqs_sin = torch.sin(angles)

    # Run both implementations
    xq_real, xk_real = apply_rotary_emb(xq.clone(), xk.clone(), freqs_cos, freqs_sin)
    xq_complex, xk_complex = apply_rotary_emb_with_complex_views(xq.clone(), xk.clone(), freqs_cos, freqs_sin)

    # Verify they produce identical results
    assert_close(xq_real, xq_complex)
    assert_close(xk_real, xk_complex)

    # Now test with Thunder JIT
    jit_apply_real = thunder.jit(apply_rotary_emb)
    jit_apply_complex = thunder.jit(apply_rotary_emb_with_complex_views)

    xq_thunder_real, xk_thunder_real = jit_apply_real(xq.clone(), xk.clone(), freqs_cos, freqs_sin)
    xq_thunder_complex, xk_thunder_complex = jit_apply_complex(xq.clone(), xk.clone(), freqs_cos, freqs_sin)

    # Verify Thunder results match PyTorch
    assert_close(xq_thunder_real, xq_real)
    assert_close(xk_thunder_real, xk_real)
    assert_close(xq_thunder_complex, xq_complex)
    assert_close(xk_thunder_complex, xk_complex)

    # Verify Thunder results match each other
    assert_close(xq_thunder_real, xq_thunder_complex)
    assert_close(xk_thunder_real, xk_thunder_complex)

    # Verify the complex version uses view_as_complex and view_as_real
    traces = thunder.last_traces(jit_apply_complex)
    extrace = traces[-1]
    symbol_names = {bsym.sym.name for bsym in extrace.bound_symbols}
    assert "view_as_complex" in symbol_names, "Complex RoPE should use view_as_complex"
    assert "view_as_real" in symbol_names, "Complex RoPE should use view_as_real"
