import pytest
import torch
from transformers import PreTrainedTokenizerBase

from deception_detection.activations import Activations
from deception_detection.sae import get_sae_acts
from deception_detection.tokenized_data import TokenizedDataset
from deception_detection.types import Dialogue, Message


@pytest.fixture
def tokenized_dataset(gemma_tokenizer: PreTrainedTokenizerBase) -> TokenizedDataset:
    dialogue = Dialogue(
        [
            Message("user", "How are you?", False),
            Message("assistant", "Very well!", True),
        ]
    )
    return TokenizedDataset.from_dialogue_list(
        [dialogue], gemma_tokenizer, padding={"gemma": {"left": 0, "right": 0}}
    )


@pytest.fixture
def random_activations(tokenized_dataset: TokenizedDataset) -> Activations:
    d_model = 3584  # gemma-9b
    acts = Activations.__new__(Activations)
    batch_size, seq_len = tokenized_dataset.tokens.shape
    acts.all_acts = torch.randn(batch_size, seq_len, 1, d_model)
    acts.layers = [31]
    acts.tokenized_dataset = tokenized_dataset
    return acts


@pytest.mark.gpu
@pytest.mark.parametrize("latent_whitelist", [None, [0, 1, 2, 3]])
def test_get_sae_acts(
    random_activations: Activations,
    latent_whitelist: list[int] | None,
):
    sae_acts = get_sae_acts(random_activations, latent_whitelist, batch_size=5)

    assert isinstance(sae_acts, Activations)
    assert sae_acts.layers == random_activations.layers
    assert sae_acts.tokenized_dataset == random_activations.tokenized_dataset

    expected_latents = 131072 if latent_whitelist is None else len(latent_whitelist)
    batch_size, seq_len = random_activations.all_acts.shape[:2]
    assert sae_acts.all_acts.shape == (batch_size, seq_len, 1, expected_latents)
