import torch

from inteprelens import LensAPI
from inteprelens.taxonomy.masks import CHGResults


def test_api_fit_llama(llama_api):
    results = llama_api.fit(
        texts=["What is the capital of France?"],
        targets=["Paris"],
        num_masks=1,
        num_updates=1,
        num_reg_updates=1,
        batch_size=1,
        verbose=False,
    )

    assert isinstance(results, CHGResults)


def test_from_pretrained_moves_model_without_device_map(monkeypatch):
    class DummyTokenizer:
        pad_token = None
        eos_token = "<eos>"

    class DummyModel(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.param = torch.nn.Parameter(torch.zeros(1))
            self.moved_to = None

        def to(self, device):
            self.moved_to = device
            return self

    dummy_model = DummyModel()
    captured = {}

    def fake_tokenizer_from_pretrained(*args, **kwargs):
        return DummyTokenizer()

    def fake_model_from_pretrained(*args, **kwargs):
        captured.update(kwargs)
        return dummy_model

    monkeypatch.setattr("transformers.AutoTokenizer.from_pretrained", fake_tokenizer_from_pretrained)
    monkeypatch.setattr(
        "transformers.AutoModelForCausalLM.from_pretrained",
        fake_model_from_pretrained,
    )
    monkeypatch.setattr("inteprelens.api.get_adapter", lambda model: object())

    analyzer = LensAPI.from_pretrained(
        "dummy/model",
        device="cuda",
        torch_dtype=torch.float32,
    )

    assert isinstance(analyzer, LensAPI)
    assert "device_map" not in captured
    assert captured["torch_dtype"] == torch.float32
    assert dummy_model.moved_to == "cuda"
