import math

import pytest
import torch
import torch.nn.functional as F

from mini_transformer import BasicEncoderDecoderTransformer
from mini_transformer.configs import ModelCfg
from mini_transformer.utils import (
    calculate_attention,
    combine_masks,
    create_causal_mask,
    extract_all_attention_maps,
    join_heads,
    sample_from_logits,
    sinusoidal_positional_encoding,
    split_heads,
)

###___split_heads___###


def test_happy_path_shape_dtype_device_cpu():
    x = torch.randn(2, 10, 16)  # B=2, S=10, D=16
    out = split_heads(x, num_heads=4)
    assert out.shape == (2, 4, 10, 4)
    assert out.dtype == x.dtype
    assert out.device == x.device


def test_smallest_valid():
    x = torch.randn(1, 1, 4)
    out = split_heads(x, num_heads=2)
    assert out.shape == (1, 2, 1, 2)


def test_zero_length_seq_allowed():
    x = torch.randn(2, 0, 8)
    out = split_heads(x, num_heads=2)
    assert out.shape == (2, 2, 0, 4)


def test_non_3d_input_raises_value_error():
    x2 = torch.randn(10, 16)
    with pytest.raises(ValueError) as ei2:
        split_heads(x2, 4)
    assert "(B, S, D)" in str(ei2.value)

    x4 = torch.randn(2, 3, 4, 5)
    with pytest.raises(ValueError) as ei4:
        split_heads(x4, 4)
    assert "(B, S, D)" in str(ei4.value)


def test_invalid_num_heads_type_raises_type_error():
    x = torch.randn(2, 10, 16)
    with pytest.raises(TypeError):
        split_heads(x, 4.0)  # float
    with pytest.raises(TypeError):
        split_heads(x, torch.tensor(4))  # Tensor


def test_num_heads_leq_zero_raises_value_error():
    x = torch.randn(2, 10, 16)
    with pytest.raises(ValueError):
        split_heads(x, 0)
    with pytest.raises(ValueError):
        split_heads(x, -2)


def test_non_divisible_d_model_raises_value_error_message_has_values():
    x = torch.randn(2, 5, 10)
    with pytest.raises(ValueError) as e:
        split_heads(x, 3)
    msg = str(e.value)
    assert "d_model (10)" in msg and "num_heads (3)" in msg


@pytest.mark.parametrize(
    "B,S,D,H",
    [
        (1, 1, 8, 1),
        (2, 3, 12, 3),
        (4, 7, 32, 8),
        (3, 0, 24, 6),
    ],
)
def test_invariants_random_valid(B, S, D, H):
    x = torch.randn(B, S, D)
    out = split_heads(x, H)
    assert out.shape[0] == B
    assert out.shape[2] == S
    assert (H * (D // H)) == D
    assert x.numel() == out.numel()


def test_grad_propagates():
    x = torch.randn(2, 10, 16, requires_grad=True)
    out = split_heads(x, 4)
    # Do a simple scalar reduction to make grad non-trivial
    loss = out.square().mean()
    loss.backward()
    assert x.grad is not None
    assert x.grad.shape == x.shape


def test_non_contiguous_input():
    # Create a non-contiguous view by transposing and slicing
    base = torch.randn(10, 2, 16)
    x = base.transpose(0, 1)[:, :8, :]  # shape (2, 8, 16), likely non-contiguous
    assert not x.is_contiguous()
    out = split_heads(x, 4)
    assert out.shape == (2, 4, 8, 4)


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_device_preserved_cuda():
    x = torch.randn(2, 10, 16, device="cuda")
    out = split_heads(x, 4)
    assert out.device.type == "cuda"


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_dtype_half_bfloat_on_gpu():
    x_half = torch.randn(2, 10, 16, device="cuda", dtype=torch.float16)
    out_half = split_heads(x_half, 4)
    assert out_half.dtype == torch.float16

    # bfloat16 may not be available on all GPUs; guard with try
    try:
        x_bf16 = torch.randn(2, 10, 16, device="cuda", dtype=torch.bfloat16)
        out_bf16 = split_heads(x_bf16, 4)
        assert out_bf16.dtype == torch.bfloat16
    except RuntimeError:
        pytest.skip("bfloat16 not supported on this device")


###___calculate_attention___###


# ----------------------------
# Helpers
# ----------------------------
def _rand(B=2, H=3, Sq=4, Sk=5, D=8, device="cpu", dtype=torch.float32):
    q = torch.randn(B, H, Sq, D, device=device, dtype=dtype)
    k = torch.randn(B, H, Sk, D, device=device, dtype=dtype)
    v = torch.randn(B, H, Sk, D, device=device, dtype=dtype)
    return q, k, v


# ----------------------------
# A. Basic correctness
# ----------------------------
def test_basic_shapes_and_row_sums():
    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
    q, k, v = _rand(B, H, Sq, Sk, D)
    attn, probs = calculate_attention(q, k, v, mask=None, return_probs=True)
    assert attn.shape == (B, H, Sq, D)
    assert probs.shape == (B, H, Sq, Sk)
    # softmax rows sum to 1
    assert torch.allclose(probs.sum(-1), torch.ones(B, H, Sq), atol=1e-6)


def test_identity_prefers_diagonal():
    # q == k == v as orthogonal-ish basis to encourage diagonal peak
    Sq = Sk = D = 4
    base = torch.eye(D).view(1, 1, Sk, D).expand(1, 1, Sk, D).contiguous()
    q = base.clone()
    k = base.clone()
    v = base.clone()
    _, probs = calculate_attention(q, k, v, mask=None, return_probs=True)
    assert torch.equal(probs[0, 0].argmax(-1), torch.arange(Sq))


# ----------------------------
# B. Mask behavior (boolean & causal)
# ----------------------------
def test_padding_mask_boolean_last_two_keys():
    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
    q, k, v = _rand(B, H, Sq, Sk, D)
    mask = torch.zeros(B, H, Sq, Sk, dtype=torch.bool)
    mask[:, :, :, -2:] = True  # mask last two keys
    _, probs = calculate_attention(q, k, v, mask=mask, return_probs=True)
    assert (probs[:, :, :, -2:] < 1e-6).all()
    # After zeroing masked entries, rows renormalize to ~1
    assert torch.allclose(probs.masked_fill(mask, 0).sum(-1), torch.ones(B, H, Sq), atol=1e-6)


def test_causal_mask_upper_triangle_zero():
    B, H, S, D = 2, 3, 6, 8
    q, k, v = _rand(B, H, S, S, D)
    causal = torch.ones(B, H, S, S, dtype=torch.bool).triu(1)  # True above diagonal
    _, probs = calculate_attention(q, k, v, mask=causal, return_probs=True)
    assert (probs.triu(1) < 1e-6).all()


def test_fully_masked_row_zero_probs_and_output():
    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
    q, k, v = _rand(B, H, Sq, Sk, D)
    mask = torch.zeros(B, H, Sq, Sk, dtype=torch.bool)
    mask[:, :, 1, :] = True  # fully mask a query row
    attn, probs = calculate_attention(q, k, v, mask=mask, return_probs=True)
    assert torch.allclose(probs[:, :, 1, :], torch.zeros(B, H, Sk))
    assert torch.allclose(attn[:, :, 1, :], torch.zeros(B, H, D))


###___extract_all_attention_maps___###


def _make_model_cfg(num_layers: int, layer_norm_style: str | None = None) -> ModelCfg:
    return ModelCfg(
        name="test",
        best_checkpoint_path="/tmp/best.ckpt",
        latest_checkpoint_path="/tmp/latest.ckpt",
        tokenizer="demo",
        d_model=16,
        num_layers=num_layers,
        num_heads=4,
        d_ff=32,
        dropout_rate=0.0,
        max_seq_len=32,
        vocab_size=64,
        pad_id=0,
        bos_id=1,
        eos_id=2,
        layer_norm_style=layer_norm_style,
    )


def _run_attention_maps(
    *,
    num_layers: int,
    layer_norm_style: str | None,
) -> tuple[ModelCfg, BasicEncoderDecoderTransformer, dict, torch.Tensor, torch.Tensor]:
    cfg = _make_model_cfg(num_layers=num_layers, layer_norm_style=layer_norm_style)
    model = BasicEncoderDecoderTransformer(cfg)
    batch, src_len, tgt_len = 2, 5, 4
    src_ids = torch.randint(0, cfg.vocab_size, (batch, src_len), dtype=torch.long)
    tgt_ids = torch.randint(0, cfg.vocab_size, (batch, tgt_len), dtype=torch.long)
    src_pad = torch.zeros(batch, src_len, dtype=torch.bool)
    tgt_pad = torch.zeros(batch, tgt_len, dtype=torch.bool)
    maps = extract_all_attention_maps(model, src_ids, tgt_ids, src_pad, tgt_pad)
    return cfg, model, maps, src_ids, tgt_ids


def _assert_attention_shapes(maps, *, batch: int, heads: int, src_len: int, tgt_len: int) -> None:
    for layer_map in maps["enc_self"]:
        assert layer_map.shape == (batch, heads, src_len, src_len)
        assert torch.isfinite(layer_map).all()
        assert torch.allclose(layer_map.sum(-1), torch.ones(batch, heads, src_len), atol=1e-5)
    for layer_map in maps["dec_self"]:
        assert layer_map.shape == (batch, heads, tgt_len, tgt_len)
        assert torch.isfinite(layer_map).all()
        assert torch.allclose(layer_map.sum(-1), torch.ones(batch, heads, tgt_len), atol=1e-5)
    for layer_map in maps["dec_cross"]:
        assert layer_map.shape == (batch, heads, tgt_len, src_len)
        assert torch.isfinite(layer_map).all()
        assert torch.allclose(layer_map.sum(-1), torch.ones(batch, heads, tgt_len), atol=1e-5)


@pytest.mark.parametrize(
    "layer_norm_style, expected_pre_norm",
    [
        (None, True),  # auto-selects pre-LN when num_layers > 4
        ("post", False),
        ("pre", True),
    ],
)
def test_extract_attention_maps_handles_layer_norm_styles(layer_norm_style, expected_pre_norm):
    cfg, model, maps, src_ids, tgt_ids = _run_attention_maps(
        num_layers=6, layer_norm_style=layer_norm_style
    )
    batch, src_len = src_ids.shape
    tgt_len = tgt_ids.shape[1]
    assert len(maps["enc_self"]) == cfg.num_layers
    assert len(maps["dec_self"]) == cfg.num_layers
    assert len(maps["dec_cross"]) == cfg.num_layers
    _assert_attention_shapes(
        maps,
        batch=batch,
        heads=cfg.num_heads,
        src_len=src_len,
        tgt_len=tgt_len,
    )
    assert all(layer.pre_norm is expected_pre_norm for layer in model.encoder.layers)
    assert all(layer.pre_norm is expected_pre_norm for layer in model.decoder.layers)


# ----------------------------
# C. Numerical stability
# ----------------------------
def test_extreme_logits_no_nans_or_infs():
    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
    q, k, v = _rand(B, H, Sq, Sk, D)
    q = q * 1000
    k = k * 1000
    attn, probs = calculate_attention(q, k, v, mask=None, return_probs=True)
    assert not (torch.isnan(probs).any() or torch.isinf(probs).any())
    assert attn.shape == (B, H, Sq, D)


def test_half_precision_close_to_fp32_or_skip():
    if not torch.cuda.is_available():
        pytest.skip("CUDA not available")
    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
    q32, k32, v32 = _rand(B, H, Sq, Sk, D, device="cuda", dtype=torch.float32)
    q16, k16, v16 = q32.half(), k32.half(), v32.half()
    a16 = calculate_attention(q16, k16, v16, mask=None)
    a32 = calculate_attention(q32, k32, v32, mask=None)
    # looser tolerance for fp16
    assert torch.allclose(a16.float(), a32.float(), atol=5e-2, rtol=5e-2)


# ----------------------------
# D. Edge cases
# ----------------------------
def test_singleton_softmax_is_one():
    q = torch.randn(1, 1, 1, 1)
    k = torch.randn(1, 1, 1, 1)
    v = torch.randn(1, 1, 1, 1)
    _, p = calculate_attention(q, k, v, mask=None, return_probs=True)
    assert torch.allclose(p, torch.ones_like(p), atol=1e-6)


# ----------------------------
# E. Autograd & determinism
# ----------------------------
def test_gradients_flow_no_inplace_breakage():
    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
    q = torch.randn(B, H, Sq, D, requires_grad=True)
    k = torch.randn(B, H, Sk, D, requires_grad=True)
    v = torch.randn(B, H, Sk, D, requires_grad=True)
    out = calculate_attention(q, k, v, mask=None)
    loss = out.pow(2).sum()
    loss.backward()
    for t in (q, k, v):
        assert t.grad is not None and torch.isfinite(t.grad).all()


def test_repeated_calls_consistent_without_dropout():
    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
    q, k, v = _rand(B, H, Sq, Sk, D)
    a1 = calculate_attention(q, k, v, mask=None)
    a2 = calculate_attention(q, k, v, mask=None)
    assert torch.allclose(a1, a2)


# ----------------------------
# F. Parity with PyTorch SDPA (when available)
# ----------------------------
def test_sdpa_parity_if_available():
    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
    q, k, v = _rand(B, H, Sq, Sk, D)
    try:
        sdpa = F.scaled_dot_product_attention(
            q, k, v, dropout_p=0.0, attn_mask=None, is_causal=False
        )
        manual = calculate_attention(q, k, v, mask=None)
        assert torch.allclose(sdpa, manual, atol=1e-5, rtol=1e-4)
    except Exception:
        # Older PyTorch: skip
        pass


# ----------------------------
# G. Performance smoke (sanity only)
# ----------------------------
def test_perf_smoke_runs_reasonably():
    # Not asserting timing; just ensure no OOM or pathological slowdowns
    q, k, v = _rand(B=1, H=8, Sq=1024, Sk=1024, D=64)
    _ = calculate_attention(q, k, v, mask=None, return_probs=False)
    q, k, v = _rand(B=2, H=8, Sq=1536, Sk=1536, D=64)
    _ = calculate_attention(q, k, v, mask=None, return_probs=False)


# ----------------------------
# H. Device behavior
# ----------------------------
def test_mask_on_cpu_qkv_on_gpu_autofix_or_skip():
    if not torch.cuda.is_available():
        pytest.skip("CUDA not available")

    prev = torch.are_deterministic_algorithms_enabled()
    try:
        # Disable determinism *only for this test*
        if prev:
            torch.use_deterministic_algorithms(False)

        q, k, v = _rand(device="cuda")
        mask = torch.zeros(q.shape[0], 1, 1, k.shape[2], dtype=torch.bool)  # CPU
        attn, probs = calculate_attention(q, k, v, mask, return_probs=True)
        assert attn.device.type == "cuda"
        assert probs.device.type == "cuda"
    finally:
        # restore whatever the global setting was
        torch.use_deterministic_algorithms(prev)


def test_qkv_device_mismatch_raises_clear_error():
    if not torch.cuda.is_available():
        pytest.skip("CUDA not available")

    prev = torch.are_deterministic_algorithms_enabled()
    try:
        if prev:
            torch.use_deterministic_algorithms(False)

        q, k, v = _rand(device="cuda")
        k = k.cpu()  # force mismatch
        with pytest.raises(RuntimeError) as ei:
            _ = calculate_attention(q, k, v, mask=None)
        msg = str(ei.value)
        assert (
            "q/k/v must be on the same device" in msg  # our explicit check
            or "Expected all tensors to be on the same device" in msg  # PyTorch matmul
            or "query, key, value must be on the same device" in msg  # updated error wording
        )
    finally:
        torch.use_deterministic_algorithms(prev)


# ----------------------------
# I. Mask broadcastability
# ----------------------------
def test_boolean_mask_broadcast_variants_equivalent():
    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
    q, k, v = _rand(B, H, Sq, Sk, D)
    m_exact = torch.zeros(B, H, Sq, Sk, dtype=torch.bool)
    m_B1_1Sk = torch.zeros(B, 1, 1, Sk, dtype=torch.bool)
    m_11SqSk = torch.zeros(1, 1, Sq, Sk, dtype=torch.bool)

    a_exact, p_exact = calculate_attention(q, k, v, m_exact, return_probs=True)
    a1, p1 = calculate_attention(q, k, v, m_B1_1Sk, return_probs=True)
    a2, p2 = calculate_attention(q, k, v, m_11SqSk, return_probs=True)

    assert torch.allclose(a1, a_exact, atol=1e-6, rtol=1e-5)
    assert torch.allclose(p1, p_exact, atol=1e-6, rtol=1e-5)
    assert torch.allclose(a2, a_exact, atol=1e-6, rtol=1e-5)
    assert torch.allclose(p2, p_exact, atol=1e-6, rtol=1e-5)


def test_non_broadcastable_mask_raises_value_error():
    B, H, Sq, Sk, D = 2, 3, 4, 5, 8
    q, k, v = _rand(B, H, Sq, Sk, D)
    bad = torch.zeros(
        B, H, Sq, dtype=torch.bool
    )  # missing last dim => not broadcastable to (B,H,Sq,Sk)
    with pytest.raises(ValueError) as ei:
        _ = calculate_attention(q, k, v, bad)
    assert "not broadcastable" in str(ei.value)


###___join_heads___###


def test_join_heads_type_error():
    with pytest.raises(TypeError):
        join_heads([1, 2, 3])  # not a tensor


def test_join_heads_dim_error():
    with pytest.raises(ValueError):
        join_heads(torch.randn(2, 3, 4))  # only 3D


def test_join_heads_shape_values():
    x = torch.randn(2, 0, 4, 5)  # zero heads
    y = join_heads(x)
    assert y.shape == (2, 4, 0)


def test_join_heads_roundtrip():
    B, H, T, Dh = 2, 3, 5, 7
    x = torch.randn(B, H, T, Dh)
    y = join_heads(x)
    assert y.shape == (B, T, H * Dh)
    # Check element mapping correctness
    for b in range(B):
        for h in range(H):
            for t in range(T):
                for d in range(Dh):
                    assert torch.allclose(y[b, t, h * Dh + d], x[b, h, t, d])


###___sinusoidal_positional_encoding___###


@pytest.mark.parametrize("seq_len,dim", [(1, 1), (1, 2), (7, 4), (17, 33)])
def test_sinusoidal_shapes_and_dtype(seq_len, dim):
    pe = sinusoidal_positional_encoding(seq_len, dim)
    assert pe.shape == (seq_len, dim)
    assert pe.dtype == torch.float32
    assert pe.requires_grad is False


def test_sinusoidal_matches_reference_small_case():
    seq_len, dim = 4, 6
    pe = sinusoidal_positional_encoding(seq_len, dim)

    ref = torch.zeros_like(pe)
    base = 10_000.0
    for pos in range(seq_len):
        for i in range(0, dim // 2):
            denom = base ** (2 * i / dim)
            ref[pos, 2 * i] = math.sin(pos / denom)
            ref[pos, 2 * i + 1] = math.cos(pos / denom)

    paired = (dim // 2) * 2
    assert torch.allclose(pe[:, :paired], ref[:, :paired], atol=1e-6, rtol=1e-6)


def test_sinusoidal_deterministic():
    a = sinusoidal_positional_encoding(8, 16)
    b = sinusoidal_positional_encoding(8, 16)
    assert torch.allclose(a, b)


def test_sinusoidal_invalid_inputs():
    with pytest.raises(ValueError):
        sinusoidal_positional_encoding(0, 8)
    with pytest.raises(ValueError):
        sinusoidal_positional_encoding(-1, 8)
    with pytest.raises(ValueError):
        sinusoidal_positional_encoding(1, 0)


###___combine_masks___###


def test_both_none_returns_none():
    assert combine_masks(None, None) is None


def test_one_none_returns_other():
    m = torch.tensor([[True, False]])
    assert torch.equal(combine_masks(m, None), m)
    assert torch.equal(combine_masks(None, m), m)


def test_or_combination():
    m1 = torch.tensor([[True, False]])
    m2 = torch.tensor([[False, True]])
    expected = torch.tensor([[True, True]])
    out = combine_masks(m1, m2)
    assert torch.equal(out, expected)


def test_shape_mismatch_error():
    m1 = torch.ones(2, 2, dtype=torch.bool)
    m2 = torch.ones(3, 3, dtype=torch.bool)
    with pytest.raises(ValueError):
        combine_masks(m1, m2)


def test_broadcastable_masks():
    m1 = torch.tensor([[True, False, True]])
    m2 = torch.tensor([[False]])  # broadcast across
    out = combine_masks(m1, m2)
    assert torch.equal(out, m1 | m2)


def test_dtype_conversion():
    # Non-bool masks still should work if they are 0/1 ints
    m1 = torch.tensor([[1, 0]], dtype=torch.int32)
    m2 = torch.tensor([[0, 1]], dtype=torch.int32)
    out = combine_masks(m1.bool(), m2.bool())
    expected = torch.tensor([[True, True]])
    assert torch.equal(out, expected)


###___sample_from_logits___###


def test_greedy_equals_argmax():
    torch.manual_seed(0)
    logits = torch.randn(4, 7)
    ids = sample_from_logits(logits)  # greedy
    ref = torch.argmax(F.softmax(logits, dim=-1), dim=-1)
    assert torch.equal(ids, ref)


def test_temperature_effect_no_sampling():
    # Greedy with temperature should leave argmax unchanged (monotonic transform)
    logits = torch.tensor([[0.1, 1.2, 0.3]], dtype=torch.float32)
    ids_t1 = sample_from_logits(logits, do_sample=False, temperature=1.0)
    ids_t05 = sample_from_logits(logits, do_sample=False, temperature=0.5)
    assert torch.equal(ids_t1, ids_t05)  # argmax unchanged


def test_top_k_limits_candidates():
    logits = torch.tensor([[1.0, 0.9, 0.1, -1.0]], dtype=torch.float32)
    # with top_k=1 greedy must select the max (index 0)
    ids = sample_from_logits(logits, do_sample=False, top_k=1)
    assert torch.equal(ids, torch.tensor([0]))
    # with sampling, and top_k=1, the only candidate is index 0
    gen = torch.Generator().manual_seed(123)
    ids_s = sample_from_logits(logits, do_sample=True, top_k=1, rng=gen)
    assert torch.equal(ids_s, torch.tensor([0]))


def test_top_p_nucleus_filters_tail():
    logits = torch.tensor([[5.0, 4.0, 3.0, -5.0]], dtype=torch.float32)  # probs heavily skewed
    ids = sample_from_logits(logits, do_sample=False, top_p=0.6)  # should keep a minimal head
    assert ids.item() in (0, 1)  # greediest is 0 anyway; check no crash


def test_min_tokens_to_keep_safety():
    logits = torch.tensor([[0.0, 0.0, 0.0, 0.0]], dtype=torch.float32)
    # Even with very small top_p, we keep at least 1 token
    ids = sample_from_logits(logits, do_sample=False, top_p=0.01, min_tokens_to_keep=1)
    assert ids.numel() == 1


def test_allow_and_deny_lists():
    logits = torch.tensor([[0.1, 2.0, 0.3, 3.0]], dtype=torch.float32)  # best is idx=3
    # Allow only {0,2} → best among those is 2
    ids = sample_from_logits(logits, do_sample=False, allowed_tokens=[0, 2])
    assert torch.equal(ids, torch.tensor([2]))
    # Deny {3} → next best is 1
    ids = sample_from_logits(logits, do_sample=False, disallowed_tokens=[3])
    assert torch.equal(ids, torch.tensor([1]))


def test_rng_reproducibility_sampling():
    logits = torch.tensor([[1.0, 1.0, 1.0, 1.0]], dtype=torch.float32)  # uniform
    g1 = torch.Generator().manual_seed(42)
    g2 = torch.Generator().manual_seed(42)
    s1 = sample_from_logits(logits, do_sample=True, rng=g1)
    s2 = sample_from_logits(logits, do_sample=True, rng=g2)
    assert torch.equal(s1, s2)


def test_device_is_preserved_cpu_cuda():
    target_device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
    logits = torch.randn(3, 5, device=target_device)
    out = sample_from_logits(logits)  # greedy by default
    assert out.device == target_device
    assert out.dtype == torch.long


def test_handles_higher_rank_logits_flattening():
    # [B, T, V] -> flattened internally; here we just ensure no crash and correct shape
    logits = torch.randn(2, 3, 7)
    out = sample_from_logits(logits)  # returns [B*T]
    assert out.shape == (2 * 3,)


###___create_causal_mask___###


def test_causal_mask_shape_and_dtype_cpu():
    x = torch.zeros(2, 5, dtype=torch.long)
    mask = create_causal_mask(x, num_heads=3)
    assert mask.shape == (2, 3, 5, 5)
    assert mask.dtype == torch.bool
    assert mask.device == x.device


def test_causal_mask_is_upper_triangular():
    x = torch.zeros(1, 4, dtype=torch.long)
    mask = create_causal_mask(x, num_heads=2)
    tri = mask[0, 0].to(torch.int)
    expected = torch.tensor(
        [
            [0, 1, 1, 1],
            [0, 0, 1, 1],
            [0, 0, 0, 1],
            [0, 0, 0, 0],
        ]
    )
    assert torch.equal(tri, expected)


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_causal_mask_tracks_device():
    x = torch.zeros(1, 3, dtype=torch.long, device="cuda")
    mask = create_causal_mask(x, num_heads=4)
    assert mask.device.type == "cuda"


def test_causal_mask_rejects_invalid_inputs():
    with pytest.raises(ValueError):
        create_causal_mask(torch.zeros(2, 0, dtype=torch.long), num_heads=2)

    with pytest.raises(ValueError):
        create_causal_mask(torch.zeros(2, 3, dtype=torch.long), num_heads=0)