"""
End-to-end validation on meta-llama/Llama-3.2-1B.

Requires: CUDA GPU, HF_TOKEN, and the model downloaded locally.
Run with: pytest -m gpu tests/test_e2e_llama.py -v
"""
import pytest
import torch

from inteprelens import LensAPI

pytestmark = pytest.mark.gpu


@pytest.fixture(scope="module")
def lens(llama_model_name, _authenticated_hf):
    if not torch.cuda.is_available():
        pytest.skip("CUDA required")
    return LensAPI.from_pretrained(
        llama_model_name,
        torch_dtype=torch.bfloat16,
    )


# ---------------------------------------------------------------------------
# CHG fit
# ---------------------------------------------------------------------------

def test_fit_pattern_task(lens):
    # batch_size=4 and num_masks=4 → 16 parallel samples per step,
    # keeping VRAM within bounds alongside the session-scoped shared model.
    results = lens.fit_pattern_task(
        "aba", num_samples=200, num_updates=10, num_reg_updates=5,
        num_masks=4, batch_size=4, verbose=False,
    )
    df = results.head_taxonomy()
    assert df is not None
    assert len(df) == lens.num_layers * lens.num_heads


# ---------------------------------------------------------------------------
# Tracing
# ---------------------------------------------------------------------------

def test_trace_attn_output_shape(lens):
    result = lens.trace(
        texts=["The capital of France is"],
        layers=[0, 1],
        sites=["attn_output"],
    )
    t = result.get("attn_output", layer=0)
    assert t.shape[-1] == lens.model.config.hidden_size


def test_trace_mlp_output_shape(lens):
    result = lens.trace(
        texts=["Hello world"],
        layers=[0],
        sites=["mlp_output"],
    )
    t = result.get("mlp_output", layer=0)
    assert t.shape[-1] == lens.model.config.hidden_size

