import torch

from inteprelens import TokenScores, gather_token_scores


def test_api_score_matches_model_logits(llama_api) -> None:
    batch = llama_api.tokenizer(
        ["Paris is the capital of France.", "Hi"],
        return_tensors="pt",
        padding=True,
    )
    input_ids = batch["input_ids"].to("cuda")
    attention_mask = batch["attention_mask"].to("cuda")

    scores = llama_api.score(input_ids=input_ids, attention_mask=attention_mask)

    with torch.no_grad():
        outputs = llama_api.model(input_ids=input_ids, attention_mask=attention_mask)

    assert isinstance(scores, TokenScores)
    assert torch.allclose(scores.logits, outputs.logits, atol=1e-5, rtol=1e-5)


def test_api_score_final_token_helpers_respect_padding(llama_api) -> None:
    texts = [
        "Paris is the capital of France.",
        "Hi",
    ]
    scores = llama_api.score(texts=texts)

    expected_positions = scores.attention_mask.long().sum(dim=-1) - 1
    expected_logits = torch.stack(
        [
            scores.logits[batch_idx, position]
            for batch_idx, position in enumerate(expected_positions.tolist())
        ],
        dim=0,
    )

    final_logits = scores.final_token_logits()
    final_log_probs = scores.final_token_log_probs()
    final_probs = scores.final_token_probs()

    assert torch.allclose(final_logits, expected_logits, atol=1e-5, rtol=1e-5)
    assert torch.allclose(final_log_probs.exp(), final_probs, atol=1e-6, rtol=1e-6)
    # bfloat16 has ~7 mantissa bits; summing 32k vocab probs accumulates ~0.4% error
    assert torch.allclose(final_probs.sum(dim=-1), torch.ones_like(final_probs.sum(dim=-1)), atol=5e-3)


def test_api_score_gather_log_probs_matches_manual_indexing(llama_api) -> None:
    scores = llama_api.score(texts=["The capital of France is"])
    token_ids = scores.input_ids.clone()

    gathered = scores.gather_log_probs(token_ids)
    expected = gather_token_scores(scores.log_probs, token_ids)

    assert torch.allclose(gathered, expected, atol=1e-6, rtol=1e-6)


def test_api_score_applies_temperature_to_logits(llama_api) -> None:
    inputs = llama_api.tokenizer("Paris is the capital of", return_tensors="pt")
    input_ids = inputs["input_ids"].to("cuda")
    attention_mask = inputs["attention_mask"].to("cuda")

    raw_scores = llama_api.score(input_ids=input_ids, attention_mask=attention_mask)
    scaled_scores = llama_api.score(
        input_ids=input_ids,
        attention_mask=attention_mask,
        temperature=2.0,
    )

    assert scaled_scores.temperature == 2.0
    assert torch.allclose(scaled_scores.logits, raw_scores.logits / 2.0, atol=1e-5, rtol=1e-5)
