from inteprelens.core.trainer import CHGTrainer
from inteprelens.utils.ds import CHGDataset


def _make_dataset(tokenizer, device: str):
    return CHGDataset.from_texts(
        texts=["What is the capital of France?"],
        targets=["Paris"],
        tokenizer=tokenizer,
        device=device,
    )


def _make_trainer(shared_model, llama_tokenizer, llama_adapter, num_masks: int = 10) -> CHGTrainer:
    dataset = _make_dataset(llama_tokenizer, "cuda")
    return CHGTrainer(
        model=shared_model,
        dataset=dataset,
        num_masks=num_masks,
        batch_size=1,
        gradient_accum_steps=1,
        adapter=llama_adapter,
        device="cuda",
    )


def test_trainer_init_creates_masks(shared_model, llama_tokenizer, llama_adapter):
    trainer = _make_trainer(shared_model, llama_tokenizer, llama_adapter, num_masks=10)

    assert trainer.mask_logits.shape == (10, trainer.masked_model.num_layers, trainer.masked_model.num_heads)
    # pos/neg are lazy — not allocated until their stage runs
    assert trainer.mask_logits_pos is None
    assert trainer.mask_logits_neg is None


def test_trainer_lazy_stage_params(shared_model, llama_tokenizer, llama_adapter):
    """pos/neg params allocate only when their stage starts."""
    trainer = _make_trainer(shared_model, llama_tokenizer, llama_adapter, num_masks=2)

    assert trainer.mask_logits_pos is None
    assert trainer.mask_logits_neg is None

    list(trainer.fit_stage("positive", num_updates=1, show_progress=False))
    assert trainer.mask_logits_pos is not None
    assert trainer.mask_logits_neg is None  # neg still not allocated

    list(trainer.fit_stage("negative", num_updates=1, show_progress=False))
    assert trainer.mask_logits_neg is not None


def test_trainer_fit_stage_yields_metrics(shared_model, llama_tokenizer, llama_adapter):
    trainer = _make_trainer(shared_model, llama_tokenizer, llama_adapter, num_masks=10)

    results = list(trainer.fit_stage("none", num_updates=1))
    assert len(results) == 1

    mask, metrics = results[0]
    assert mask.shape == (10, trainer.masked_model.num_layers, trainer.masked_model.num_heads)
    assert metrics["regularization"] == "none"
    assert "nll" in metrics
    assert "L1" in metrics


def test_trainer_fit_runs_all_stages(shared_model, llama_tokenizer, llama_adapter):
    trainer = _make_trainer(shared_model, llama_tokenizer, llama_adapter, num_masks=10)

    metrics = [m for _, m in trainer.fit(num_updates=1, num_reg_updates=1, verbose=False)]
    assert [m["regularization"] for m in metrics] == ["none", "positive", "negative"]
