import torch

from inteprelens import logits_to_log_probs, logits_to_probs


def test_tracer_captures_block_output_final_norm_and_logits(llama_tracer, llama_adapter):
    layers = [0, llama_adapter.num_layers - 1]
    input_ids = torch.tensor([[1, 2, 3, 4]], device="cuda")

    trace = llama_tracer.trace(
        input_ids,
        layers=layers,
        sites=["block_output", "final_norm", "logits", "lm_head_input"],
    )

    for layer_idx in layers:
        block_output = trace.get("block_output", layer_idx)
        logits = trace.get("logits", layer_idx)
        assert block_output.ndim == 3
        assert logits.ndim == 3
        assert logits.shape[:2] == input_ids.shape

    assert trace.get("final_norm").shape == trace.get("lm_head_input").shape
    assert torch.allclose(trace.get("final_norm"), trace.get("lm_head_input"), atol=1e-3, rtol=1e-3)
    assert trace.final_logits is not None


def test_trace_last_layer_logits_match_model_logits(llama_tracer, llama_adapter):
    last_layer = llama_adapter.num_layers - 1
    input_ids = torch.tensor([[42, 17, 5]], device="cuda")

    trace = llama_tracer.trace(
        input_ids,
        layers=[last_layer],
        sites=["block_output", "logits"],
    )

    traced_logits = trace.get("logits", last_layer)
    assert trace.final_logits is not None
    assert torch.allclose(traced_logits, trace.final_logits, atol=5e-2, rtol=5e-2)


def test_api_trace_uses_text_input_and_cleans_hooks(llama_api, llama_adapter):
    layer = llama_adapter.get_layers()[0]
    o_proj = llama_adapter.get_attention_output_proj(layer)
    final_norm = llama_adapter.get_final_norm()

    o_proj_hooks_before = len(o_proj._forward_pre_hooks)
    final_norm_hooks_before = 0 if final_norm is None else len(final_norm._forward_hooks)

    trace = llama_api.trace(
        texts="Paris is the capital of",
        layers=[0],
        sites=["attn_o_proj_pre", "final_norm", "logits"],
    )

    assert trace.has("attn_o_proj_pre", 0)
    assert trace.has("final_norm")
    assert trace.has("logits", 0)
    assert len(o_proj._forward_pre_hooks) == o_proj_hooks_before
    if final_norm is not None:
        assert len(final_norm._forward_hooks) == final_norm_hooks_before


def test_trace_result_probability_helpers_match_logits(llama_tracer, llama_adapter):
    last_layer = llama_adapter.num_layers - 1
    input_ids = torch.tensor([[42, 17, 5]], device="cuda")

    trace = llama_tracer.trace(
        input_ids,
        layers=[last_layer],
        sites=["block_output", "logits"],
    )

    assert trace.final_logits is not None
    assert torch.allclose(
        trace.final_log_probs(),
        logits_to_log_probs(trace.final_logits),
        atol=1e-6,
        rtol=1e-6,
    )
    assert torch.allclose(
        trace.final_probs(),
        logits_to_probs(trace.final_logits),
        atol=1e-6,
        rtol=1e-6,
    )
    assert torch.allclose(
        trace.layer_log_probs(last_layer),
        logits_to_log_probs(trace.get("logits", last_layer)),
        atol=1e-6,
        rtol=1e-6,
    )
    assert torch.allclose(
        trace.layer_probs(last_layer),
        logits_to_probs(trace.get("logits", last_layer)),
        atol=1e-6,
        rtol=1e-6,
    )
