import torch

from inteprelens import (
    apply_temperature,
    gather_token_scores,
    logits_to_log_probs,
    logits_to_probs,
    select_final_token,
)


def test_log_probs_exp_matches_softmax() -> None:
    logits = torch.tensor(
        [[[1.0, 0.0, -1.0], [0.5, 2.0, -0.5]]],
        dtype=torch.float32,
    )

    log_probs = logits_to_log_probs(logits)
    probs = logits_to_probs(logits)

    assert torch.allclose(log_probs.exp(), probs, atol=1e-6, rtol=1e-6)


def test_apply_temperature_changes_confidence() -> None:
    logits = torch.tensor([[3.0, 0.0]], dtype=torch.float32)

    cold_probs = logits_to_probs(apply_temperature(logits, 0.5))
    warm_probs = logits_to_probs(apply_temperature(logits, 2.0))

    assert cold_probs[0, 0] > warm_probs[0, 0]


def test_select_final_token_respects_attention_mask_for_left_and_right_padding() -> None:
    tensor = torch.tensor(
        [
            [[10.0, 11.0], [20.0, 21.0], [30.0, 31.0], [40.0, 41.0]],
            [[50.0, 51.0], [60.0, 61.0], [70.0, 71.0], [80.0, 81.0]],
        ]
    )
    attention_mask = torch.tensor(
        [
            [1, 1, 0, 0],
            [0, 0, 1, 1],
        ]
    )

    selected = select_final_token(tensor, attention_mask=attention_mask)

    expected = torch.tensor(
        [
            [20.0, 21.0],
            [80.0, 81.0],
        ]
    )
    assert torch.equal(selected, expected)


def test_gather_token_scores_matches_manual_indexing() -> None:
    log_probs = torch.tensor(
        [
            [[-2.0, -1.0, -0.5], [-3.0, -0.2, -1.1]],
            [[-0.1, -2.1, -3.1], [-4.0, -5.0, -0.3]],
        ]
    )
    token_ids = torch.tensor(
        [
            [2, 1],
            [0, 2],
        ]
    )

    gathered = gather_token_scores(log_probs, token_ids)
    expected = torch.tensor(
        [
            [-0.5, -0.2],
            [-0.1, -0.3],
        ]
    )

    assert torch.equal(gathered, expected)
