"""Tests for vauban.softprompt: types, forward pass, continuous, GCG, evaluation."""

from __future__ import annotations

import math

import mlx.core as mx
import pytest
from conftest import (
    D_MODEL,
    NUM_HEADS,
    NUM_LAYERS,
    VOCAB_SIZE,
    MockCausalLM,
    MockTokenizer,
)

from vauban.softprompt import (
    _build_vocab_mask,
    _compute_accessibility_score,
    _compute_embed_regularization,
    _compute_eos_loss,
    _compute_kl_collision_loss,
    _compute_learning_rate,
    _compute_loss,
    _compute_untargeted_loss,
    _continuous_attack,
    _egd_attack,
    _encode_refusal_tokens,
    _encode_targets,
    _evaluate_attack,
    _forward_with_prefix,
    _gcg_attack,
    _pre_encode_prompts,
    _project_to_tokens,
    _select_prompt_ids,
    _select_worst_k_prompt_ids,
    _split_into_batches,
    softprompt_attack,
)
from vauban.types import (
    SoftPromptConfig,
    SoftPromptResult,
    TransferEvalResult,
)

# ---------------------------------------------------------------------------
# Type tests
# ---------------------------------------------------------------------------


class TestSoftPromptConfig:
    def test_defaults(self) -> None:
        cfg = SoftPromptConfig()
        assert cfg.mode == "continuous"
        assert cfg.n_tokens == 16
        assert cfg.n_steps == 200
        assert cfg.learning_rate == 0.01
        assert cfg.init_scale == 0.1
        assert cfg.batch_size == 64
        assert cfg.top_k == 256
        assert cfg.direction_weight == 0.0
        assert cfg.target_prefixes == ["Sure", "Here"]
        assert cfg.max_gen_tokens == 100
        assert cfg.seed is None
        assert cfg.embed_reg_weight == 0.0
        assert cfg.patience == 0
        assert cfg.lr_schedule == "constant"
        assert cfg.n_restarts == 1
        assert cfg.prompt_strategy == "all"

    def test_frozen(self) -> None:
        cfg = SoftPromptConfig()
        with pytest.raises(AttributeError):
            cfg.n_tokens = 10  # type: ignore[misc]

    def test_custom_values(self) -> None:
        cfg = SoftPromptConfig(
            mode="gcg",
            n_tokens=32,
            n_steps=100,
            learning_rate=0.001,
            batch_size=128,
            top_k=512,
            direction_weight=0.5,
            target_prefixes=["OK"],
            seed=42,
        )
        assert cfg.mode == "gcg"
        assert cfg.n_tokens == 32
        assert cfg.seed == 42


class TestSoftPromptResult:
    def test_construction(self) -> None:
        result = SoftPromptResult(
            mode="continuous",
            success_rate=0.8,
            final_loss=1.5,
            loss_history=[3.0, 2.0, 1.5],
            n_steps=3,
            n_tokens=16,
            embeddings=mx.zeros((1, 16, 8)),
            token_ids=None,
            token_text=None,
            eval_responses=["response1"],
        )
        assert result.mode == "continuous"
        assert result.success_rate == 0.8
        assert result.final_loss == 1.5
        assert len(result.loss_history) == 3
        assert result.embeddings is not None
        assert result.token_ids is None
        # Default new fields
        assert result.accessibility_score == 0.0
        assert result.per_prompt_losses == []
        assert result.early_stopped is False

    def test_frozen(self) -> None:
        result = SoftPromptResult(
            mode="continuous",
            success_rate=0.0,
            final_loss=0.0,
            loss_history=[],
            n_steps=0,
            n_tokens=1,
            embeddings=None,
            token_ids=None,
            token_text=None,
            eval_responses=[],
        )
        with pytest.raises(AttributeError):
            result.success_rate = 1.0  # type: ignore[misc]


# ---------------------------------------------------------------------------
# Forward pass tests
# ---------------------------------------------------------------------------


class TestForwardWithPrefix:
    def test_output_shape(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())

        n_tokens = 4
        seq_len = 6
        soft_embeds = mx.random.normal((1, n_tokens, D_MODEL))
        prompt_ids = mx.array([[0, 1, 2, 3, 4, 5]])
        mx.eval(soft_embeds)

        logits = _forward_with_prefix(model, soft_embeds, prompt_ids)
        mx.eval(logits)

        assert logits.shape == (1, n_tokens + seq_len, VOCAB_SIZE)

    def test_prefix_affects_logits(self) -> None:
        """Verify that different prefixes produce different output logits."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())

        prompt_ids = mx.array([[0, 1, 2]])
        zero_embeds = mx.zeros((1, 4, D_MODEL))
        rand_embeds = mx.random.normal((1, 4, D_MODEL))
        mx.eval(zero_embeds, rand_embeds)

        logits_zero = _forward_with_prefix(model, zero_embeds, prompt_ids)
        logits_rand = _forward_with_prefix(model, rand_embeds, prompt_ids)
        mx.eval(logits_zero, logits_rand)

        diff = float(
            mx.mean(mx.abs(logits_zero[:, -1, :] - logits_rand[:, -1, :])).item(),
        )
        assert diff > 0.001, (
            f"Different prefixes produce nearly identical logits (diff={diff})"
        )

    def test_different_prefix_sizes(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())

        prompt_ids = mx.array([[0, 1, 2]])

        for n_tokens in [1, 8, 16]:
            soft_embeds = mx.random.normal((1, n_tokens, D_MODEL))
            mx.eval(soft_embeds)
            logits = _forward_with_prefix(model, soft_embeds, prompt_ids)
            mx.eval(logits)
            assert logits.shape[1] == n_tokens + 3


# ---------------------------------------------------------------------------
# Encode targets tests
# ---------------------------------------------------------------------------


class TestEncodeTargets:
    def test_encodes_prefixes(self) -> None:
        tokenizer = MockTokenizer(VOCAB_SIZE)
        ids = _encode_targets(tokenizer, ["Sure", "Here"])
        mx.eval(ids)
        assert ids.ndim == 1
        assert ids.shape[0] > 0

    def test_single_prefix(self) -> None:
        tokenizer = MockTokenizer(VOCAB_SIZE)
        ids = _encode_targets(tokenizer, ["OK"])
        mx.eval(ids)
        expected = tokenizer.encode("OK")
        assert ids.shape[0] == len(expected)


# ---------------------------------------------------------------------------
# Continuous attack tests
# ---------------------------------------------------------------------------


class TestContinuousAttack:
    def test_basic_run(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="continuous",
            n_tokens=4,
            n_steps=5,
            learning_rate=0.01,
            seed=42,
            max_gen_tokens=3,
        )

        result = _continuous_attack(
            model, tokenizer, ["test prompt"], config, None,
        )

        assert isinstance(result, SoftPromptResult)
        assert result.mode == "continuous"
        assert result.n_steps == 5
        assert result.n_tokens == 4
        assert len(result.loss_history) == 5
        assert result.embeddings is not None
        assert result.embeddings.shape == (1, 4, D_MODEL)
        assert result.token_ids is None
        assert result.token_text is None
        assert len(result.eval_responses) == 1

    def test_loss_is_finite(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            n_tokens=4,
            n_steps=3,
            seed=42,
            max_gen_tokens=2,
        )

        result = _continuous_attack(
            model, tokenizer, ["hello"], config, None,
        )

        for loss in result.loss_history:
            assert not math.isnan(loss), "Loss is NaN"

    def test_loss_decreases(self) -> None:
        """Verify that optimization actually reduces loss (gradient flow works)."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="continuous",
            n_tokens=4,
            n_steps=30,
            learning_rate=0.1,
            init_scale=0.5,
            seed=42,
            max_gen_tokens=2,
        )

        result = _continuous_attack(
            model, tokenizer, ["test"], config, None,
        )

        first_loss = result.loss_history[0]
        final_loss = result.loss_history[-1]
        assert final_loss < first_loss, (
            f"Loss did not decrease: {first_loss:.4f} -> {final_loss:.4f}"
        )

    def test_gradient_nonzero(self) -> None:
        """Verify gradient w.r.t. soft embeds is nonzero."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        soft_embeds = mx.random.normal((1, 4, D_MODEL)) * 0.1
        mx.eval(soft_embeds)

        target_ids = _encode_targets(tokenizer, ["Sure"])
        mx.eval(target_ids)

        messages = [{"role": "user", "content": "test"}]
        text = tokenizer.apply_chat_template(messages, tokenize=False)
        assert isinstance(text, str)
        prompt_ids = mx.array(tokenizer.encode(text))[None, :]

        def loss_fn(embeds: mx.array) -> mx.array:
            return _compute_loss(model, embeds, prompt_ids, target_ids, 4, None, 0.0)

        loss_and_grad = mx.value_and_grad(loss_fn)
        loss_val, grad = loss_and_grad(soft_embeds)
        mx.eval(loss_val, grad)

        grad_norm = float(mx.linalg.norm(grad.reshape(-1)).item())
        assert grad_norm > 0, "Gradient is zero — no learning signal"

    def test_dispatch_via_public_api(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="continuous",
            n_tokens=4,
            n_steps=3,
            seed=42,
            max_gen_tokens=2,
        )

        result = softprompt_attack(
            model, tokenizer, ["test"], config, None,
        )
        assert result.mode == "continuous"


# ---------------------------------------------------------------------------
# GCG attack tests
# ---------------------------------------------------------------------------


class TestGCGAttack:
    def test_basic_run(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="gcg",
            n_tokens=4,
            n_steps=3,
            batch_size=8,
            top_k=16,
            seed=42,
            max_gen_tokens=3,
        )

        result = _gcg_attack(
            model, tokenizer, ["test prompt"], config, None,
        )

        assert isinstance(result, SoftPromptResult)
        assert result.mode == "gcg"
        assert result.n_tokens == 4
        assert len(result.loss_history) == 3
        assert result.embeddings is None
        assert result.token_ids is not None
        assert len(result.token_ids) == 4
        assert result.token_text is not None
        assert len(result.eval_responses) == 1

    def test_token_ids_in_vocab_range(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="gcg",
            n_tokens=4,
            n_steps=2,
            batch_size=4,
            top_k=8,
            seed=42,
            max_gen_tokens=2,
        )

        result = _gcg_attack(
            model, tokenizer, ["test"], config, None,
        )

        assert result.token_ids is not None
        for tid in result.token_ids:
            assert 0 <= tid < VOCAB_SIZE

    def test_gcg_improves_loss(self) -> None:
        """Verify GCG finds at least one candidate better than random init."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="gcg",
            n_tokens=4,
            n_steps=10,
            batch_size=16,
            top_k=VOCAB_SIZE,  # search full vocab on small model
            seed=42,
            max_gen_tokens=2,
        )

        result = _gcg_attack(
            model, tokenizer, ["test"], config, None,
        )

        first_loss = result.loss_history[0]
        best_loss = result.final_loss
        assert best_loss <= first_loss, (
            f"GCG did not improve: {first_loss:.4f} -> {best_loss:.4f}"
        )

    def test_dispatch_via_public_api(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="gcg",
            n_tokens=4,
            n_steps=2,
            batch_size=4,
            top_k=8,
            seed=42,
            max_gen_tokens=2,
        )

        result = softprompt_attack(
            model, tokenizer, ["test"], config, None,
        )
        assert result.mode == "gcg"


# ---------------------------------------------------------------------------
# Direction-guided tests
# ---------------------------------------------------------------------------


class TestDirectionGuided:
    def test_continuous_with_direction(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        direction = mx.random.normal((D_MODEL,))
        direction = direction / mx.linalg.norm(direction)
        mx.eval(direction)

        config = SoftPromptConfig(
            mode="continuous",
            n_tokens=4,
            n_steps=3,
            direction_weight=0.1,
            seed=42,
            max_gen_tokens=2,
        )

        result = softprompt_attack(
            model, tokenizer, ["test"], config, direction,
        )

        assert result.mode == "continuous"
        assert len(result.loss_history) == 3

    def test_gcg_with_direction(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        direction = mx.random.normal((D_MODEL,))
        direction = direction / mx.linalg.norm(direction)
        mx.eval(direction)

        config = SoftPromptConfig(
            mode="gcg",
            n_tokens=4,
            n_steps=2,
            batch_size=4,
            top_k=8,
            direction_weight=0.1,
            seed=42,
            max_gen_tokens=2,
        )

        result = softprompt_attack(
            model, tokenizer, ["test"], config, direction,
        )

        assert result.mode == "gcg"
        assert len(result.loss_history) == 2


# ---------------------------------------------------------------------------
# Evaluate attack tests
# ---------------------------------------------------------------------------


class TestEvaluateAttack:
    def test_success_rate_range(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        soft_embeds = mx.random.normal((1, 4, D_MODEL))
        mx.eval(soft_embeds)

        config = SoftPromptConfig(max_gen_tokens=3)

        success_rate, responses = _evaluate_attack(
            model, tokenizer, ["test1", "test2"], soft_embeds, config,
        )

        assert 0.0 <= success_rate <= 1.0
        assert len(responses) == 2

    def test_empty_prompts(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        soft_embeds = mx.random.normal((1, 4, D_MODEL))
        mx.eval(soft_embeds)

        config = SoftPromptConfig(max_gen_tokens=2)

        success_rate, responses = _evaluate_attack(
            model, tokenizer, [], soft_embeds, config,
        )

        assert success_rate == 0.0
        assert responses == []


# ---------------------------------------------------------------------------
# Invalid mode test
# ---------------------------------------------------------------------------


class TestInvalidMode:
    def test_invalid_mode_raises(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        # Bypass frozen by constructing manually
        bad_config = SoftPromptConfig.__new__(SoftPromptConfig)
        object.__setattr__(bad_config, "mode", "invalid")
        object.__setattr__(bad_config, "n_tokens", 4)
        object.__setattr__(bad_config, "n_steps", 1)
        object.__setattr__(bad_config, "seed", None)

        with pytest.raises(ValueError, match="Unknown soft prompt mode"):
            softprompt_attack(model, tokenizer, ["test"], bad_config, None)


# ---------------------------------------------------------------------------
# New helper tests
# ---------------------------------------------------------------------------


class TestComputeLearningRate:
    def test_constant_returns_base(self) -> None:
        assert _compute_learning_rate(0.01, 0, 100, "constant") == 0.01
        assert _compute_learning_rate(0.01, 50, 100, "constant") == 0.01
        assert _compute_learning_rate(0.01, 99, 100, "constant") == 0.01

    def test_cosine_start_equals_base(self) -> None:
        lr = _compute_learning_rate(0.01, 0, 100, "cosine")
        assert abs(lr - 0.01) < 1e-9

    def test_cosine_end_near_zero(self) -> None:
        lr = _compute_learning_rate(0.01, 99, 100, "cosine")
        assert abs(lr) < 1e-9

    def test_cosine_midpoint(self) -> None:
        lr = _compute_learning_rate(0.01, 49, 99, "cosine")
        # cos(pi * 49/98) = cos(pi/2) = 0, so lr = 0.005
        assert abs(lr - 0.005) < 1e-6

    def test_cosine_single_step(self) -> None:
        # n_steps=1 -> schedule has no effect (n_steps - 1 = 0)
        lr = _compute_learning_rate(0.01, 0, 1, "cosine")
        assert lr == 0.01


class TestComputeEmbedRegularization:
    def test_zero_weight_returns_zero(self) -> None:
        soft = mx.random.normal((1, 4, 16))
        embed = mx.random.normal((32, 16))
        mx.eval(soft, embed)
        result = _compute_embed_regularization(soft, embed, 0.0)
        mx.eval(result)
        assert float(result.item()) == 0.0

    def test_matching_norms_near_zero(self) -> None:
        """When soft embeds have same mean norm as real embeds, reg is ~0."""
        embed = mx.random.normal((32, 16))
        mx.eval(embed)
        mean_real = float(mx.mean(mx.linalg.norm(embed, axis=-1)).item())
        # Create soft embeds normalized to match
        soft = mx.random.normal((1, 4, 16))
        mx.eval(soft)
        norms = mx.linalg.norm(soft[0], axis=-1, keepdims=True)
        soft_normalized = soft / norms * mean_real
        soft_normalized = soft_normalized[None, :] if soft_normalized.ndim == 2 else soft_normalized  # noqa: E501
        mx.eval(soft_normalized)
        result = _compute_embed_regularization(soft_normalized, embed, 1.0)
        mx.eval(result)
        assert float(result.item()) < 0.01

    def test_large_gap_gives_large_penalty(self) -> None:
        soft = mx.ones((1, 4, 16)) * 100.0  # very large norm
        embed = mx.ones((32, 16)) * 0.01  # very small norm
        mx.eval(soft, embed)
        result = _compute_embed_regularization(soft, embed, 1.0)
        mx.eval(result)
        assert float(result.item()) > 1.0


class TestAccessibilityScore:
    def test_zero_loss(self) -> None:
        assert _compute_accessibility_score(0.0) == 1.0

    def test_high_loss(self) -> None:
        score = _compute_accessibility_score(10.0)
        assert score < 0.001

    def test_monotonicity(self) -> None:
        s1 = _compute_accessibility_score(1.0)
        s2 = _compute_accessibility_score(2.0)
        s3 = _compute_accessibility_score(3.0)
        assert s1 > s2 > s3


class TestPreEncodePrompts:
    def test_multi_prompt(self) -> None:
        tokenizer = MockTokenizer(VOCAB_SIZE)
        prompts = ["hello", "world", "test"]
        encoded = _pre_encode_prompts(tokenizer, prompts)
        assert len(encoded) == 3
        for ids in encoded:
            assert ids.ndim == 2
            assert ids.shape[0] == 1
            assert ids.shape[1] > 0

    def test_empty_list(self) -> None:
        tokenizer = MockTokenizer(VOCAB_SIZE)
        encoded = _pre_encode_prompts(tokenizer, [])
        assert encoded == []


class TestSelectPromptIds:
    def test_first_strategy(self) -> None:
        ids = [mx.array([[1]]), mx.array([[2]]), mx.array([[3]])]
        selected = _select_prompt_ids(ids, 0, "first")
        assert len(selected) == 1
        assert int(selected[0].item()) == 1
        # Same regardless of step
        selected2 = _select_prompt_ids(ids, 5, "first")
        assert len(selected2) == 1
        assert int(selected2[0].item()) == 1

    def test_cycle_strategy(self) -> None:
        ids = [mx.array([[1]]), mx.array([[2]]), mx.array([[3]])]
        assert int(_select_prompt_ids(ids, 0, "cycle")[0].item()) == 1
        assert int(_select_prompt_ids(ids, 1, "cycle")[0].item()) == 2
        assert int(_select_prompt_ids(ids, 2, "cycle")[0].item()) == 3
        assert int(_select_prompt_ids(ids, 3, "cycle")[0].item()) == 1

    def test_all_strategy(self) -> None:
        ids = [mx.array([[1]]), mx.array([[2]]), mx.array([[3]])]
        selected = _select_prompt_ids(ids, 0, "all")
        assert len(selected) == 3


# ---------------------------------------------------------------------------
# Multi-prompt continuous tests
# ---------------------------------------------------------------------------


class TestMultiPromptContinuous:
    def test_all_strategy(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            n_tokens=4,
            n_steps=3,
            seed=42,
            max_gen_tokens=2,
            prompt_strategy="all",
        )

        result = _continuous_attack(
            model, tokenizer, ["hello", "world"], config, None,
        )

        assert result.mode == "continuous"
        assert len(result.per_prompt_losses) == 2
        assert all(loss > 0 for loss in result.per_prompt_losses)
        assert result.accessibility_score > 0.0

    def test_cycle_strategy(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            n_tokens=4,
            n_steps=4,
            seed=42,
            max_gen_tokens=2,
            prompt_strategy="cycle",
        )

        result = _continuous_attack(
            model, tokenizer, ["hello", "world"], config, None,
        )

        assert result.mode == "continuous"
        assert len(result.per_prompt_losses) == 2
        assert len(result.loss_history) == 4


# ---------------------------------------------------------------------------
# Early stopping tests
# ---------------------------------------------------------------------------


class TestEarlyStopping:
    def test_fires_with_patience(self) -> None:
        """Early stopping should fire with tiny LR and small patience."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            n_tokens=4,
            n_steps=100,
            learning_rate=1e-10,  # tiny LR -> no improvement
            seed=42,
            max_gen_tokens=2,
            patience=3,
        )

        result = _continuous_attack(
            model, tokenizer, ["test"], config, None,
        )

        assert result.early_stopped is True
        assert result.n_steps < 100

    def test_disabled_with_patience_zero(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            n_tokens=4,
            n_steps=5,
            seed=42,
            max_gen_tokens=2,
            patience=0,
        )

        result = _continuous_attack(
            model, tokenizer, ["test"], config, None,
        )

        assert result.early_stopped is False
        assert result.n_steps == 5


# ---------------------------------------------------------------------------
# Cosine schedule test
# ---------------------------------------------------------------------------


class TestCosineSchedule:
    def test_runs_without_nan(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            n_tokens=4,
            n_steps=10,
            learning_rate=0.01,
            seed=42,
            max_gen_tokens=2,
            lr_schedule="cosine",
        )

        result = _continuous_attack(
            model, tokenizer, ["test"], config, None,
        )

        for loss in result.loss_history:
            assert not math.isnan(loss), "Cosine schedule produced NaN loss"


# ---------------------------------------------------------------------------
# Embed regularization test
# ---------------------------------------------------------------------------


class TestEmbedRegularizationIntegration:
    def test_runs_without_error(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            n_tokens=4,
            n_steps=3,
            seed=42,
            max_gen_tokens=2,
            embed_reg_weight=0.1,
        )

        result = _continuous_attack(
            model, tokenizer, ["test"], config, None,
        )

        assert result.mode == "continuous"
        for loss in result.loss_history:
            assert not math.isnan(loss)


# ---------------------------------------------------------------------------
# GCG multi-restart tests
# ---------------------------------------------------------------------------


class TestGCGMultiRestart:
    def test_correct_loss_history_length(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="gcg",
            n_tokens=4,
            n_steps=3,
            batch_size=4,
            top_k=8,
            seed=42,
            max_gen_tokens=2,
            n_restarts=2,
        )

        result = _gcg_attack(
            model, tokenizer, ["test"], config, None,
        )

        # 2 restarts * 3 steps = 6 loss history entries
        assert len(result.loss_history) == 6
        assert result.n_steps == 6

    def test_token_ids_correct_length(self) -> None:
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="gcg",
            n_tokens=4,
            n_steps=2,
            batch_size=4,
            top_k=8,
            seed=42,
            max_gen_tokens=2,
            n_restarts=3,
        )

        result = _gcg_attack(
            model, tokenizer, ["test"], config, None,
        )

        assert result.token_ids is not None
        assert len(result.token_ids) == 4
        for tid in result.token_ids:
            assert 0 <= tid < VOCAB_SIZE


# ---------------------------------------------------------------------------
# SoftPromptResult new fields tests
# ---------------------------------------------------------------------------


class TestSoftPromptResultNewFields:
    def test_defaults(self) -> None:
        result = SoftPromptResult(
            mode="continuous",
            success_rate=0.5,
            final_loss=2.0,
            loss_history=[2.0],
            n_steps=1,
            n_tokens=4,
            embeddings=None,
            token_ids=None,
            token_text=None,
            eval_responses=[],
        )
        assert result.accessibility_score == 0.0
        assert result.per_prompt_losses == []
        assert result.early_stopped is False

    def test_explicit_values(self) -> None:
        result = SoftPromptResult(
            mode="gcg",
            success_rate=0.9,
            final_loss=0.5,
            loss_history=[1.0, 0.5],
            n_steps=2,
            n_tokens=8,
            embeddings=None,
            token_ids=[1, 2, 3, 4, 5, 6, 7, 8],
            token_text="test",
            eval_responses=["response"],
            accessibility_score=0.6,
            per_prompt_losses=[0.4, 0.6],
            early_stopped=True,
        )
        assert result.accessibility_score == 0.6
        assert result.per_prompt_losses == [0.4, 0.6]
        assert result.early_stopped is True


# ---------------------------------------------------------------------------
# New config defaults tests
# ---------------------------------------------------------------------------


class TestNewConfigDefaults:
    def test_new_fields_defaults(self) -> None:
        cfg = SoftPromptConfig()
        assert cfg.direction_mode == "last"
        assert cfg.direction_layers is None
        assert cfg.loss_mode == "targeted"
        assert cfg.egd_temperature == 1.0

    def test_egd_mode_accepted(self) -> None:
        cfg = SoftPromptConfig(mode="egd")
        assert cfg.mode == "egd"

    def test_custom_direction_mode(self) -> None:
        cfg = SoftPromptConfig(direction_mode="raid")
        assert cfg.direction_mode == "raid"

    def test_custom_direction_layers(self) -> None:
        cfg = SoftPromptConfig(direction_layers=[0, 1])
        assert cfg.direction_layers == [0, 1]


# ---------------------------------------------------------------------------
# RAID direction mode tests
# ---------------------------------------------------------------------------


class TestRAIDDirectionMode:
    def test_raid_runs(self) -> None:
        """RAID mode with direction_weight=0.1 produces finite loss."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        direction = mx.random.normal((D_MODEL,))
        direction = direction / mx.linalg.norm(direction)
        mx.eval(direction)

        config = SoftPromptConfig(
            mode="continuous",
            n_tokens=4,
            n_steps=3,
            direction_weight=0.1,
            direction_mode="raid",
            seed=42,
            max_gen_tokens=2,
        )

        result = softprompt_attack(
            model, tokenizer, ["test"], config, direction,
        )

        assert result.mode == "continuous"
        for loss in result.loss_history:
            assert not math.isnan(loss), "RAID mode produced NaN loss"

    def test_raid_differs_from_last(self) -> None:
        """Same seed, RAID vs 'last' produce different final losses."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        direction = mx.random.normal((D_MODEL,))
        direction = direction / mx.linalg.norm(direction)
        mx.eval(direction)

        config_last = SoftPromptConfig(
            mode="continuous",
            n_tokens=4,
            n_steps=5,
            learning_rate=0.1,
            direction_weight=0.5,
            direction_mode="last",
            seed=42,
            max_gen_tokens=2,
        )
        result_last = _continuous_attack(
            model, tokenizer, ["test"], config_last, direction,
        )

        config_raid = SoftPromptConfig(
            mode="continuous",
            n_tokens=4,
            n_steps=5,
            learning_rate=0.1,
            direction_weight=0.5,
            direction_mode="raid",
            seed=42,
            max_gen_tokens=2,
        )
        result_raid = _continuous_attack(
            model, tokenizer, ["test"], config_raid, direction,
        )

        # They should differ since RAID accumulates across layers
        assert result_last.loss_history != result_raid.loss_history

    def test_all_positions_runs(self) -> None:
        """direction_mode='all_positions' runs without error."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        direction = mx.random.normal((D_MODEL,))
        direction = direction / mx.linalg.norm(direction)
        mx.eval(direction)

        config = SoftPromptConfig(
            mode="continuous",
            n_tokens=4,
            n_steps=3,
            direction_weight=0.1,
            direction_mode="all_positions",
            seed=42,
            max_gen_tokens=2,
        )

        result = softprompt_attack(
            model, tokenizer, ["test"], config, direction,
        )

        assert result.mode == "continuous"
        for loss in result.loss_history:
            assert not math.isnan(loss)

    def test_direction_layers_subset(self) -> None:
        """direction_layers=[0] accepted, runs."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        direction = mx.random.normal((D_MODEL,))
        direction = direction / mx.linalg.norm(direction)
        mx.eval(direction)

        config = SoftPromptConfig(
            mode="continuous",
            n_tokens=4,
            n_steps=3,
            direction_weight=0.1,
            direction_mode="raid",
            direction_layers=[0],
            seed=42,
            max_gen_tokens=2,
        )

        result = softprompt_attack(
            model, tokenizer, ["test"], config, direction,
        )

        assert result.mode == "continuous"
        assert len(result.loss_history) == 3


# ---------------------------------------------------------------------------
# Untargeted loss tests
# ---------------------------------------------------------------------------


class TestUntargetedLoss:
    def test_untargeted_runs(self) -> None:
        """loss_mode='untargeted' produces finite loss."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="continuous",
            n_tokens=4,
            n_steps=5,
            learning_rate=0.01,
            loss_mode="untargeted",
            seed=42,
            max_gen_tokens=2,
        )

        result = softprompt_attack(
            model, tokenizer, ["test"], config, None,
        )

        assert result.mode == "continuous"
        for loss in result.loss_history:
            assert not math.isnan(loss), "Untargeted loss produced NaN"

    def test_untargeted_gradient_nonzero(self) -> None:
        """Gradient through untargeted loss is nonzero."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        soft_embeds = mx.random.normal((1, 4, D_MODEL)) * 0.1
        mx.eval(soft_embeds)

        refusal_ids = _encode_refusal_tokens(tokenizer)
        mx.eval(refusal_ids)

        messages = [{"role": "user", "content": "test"}]
        text = tokenizer.apply_chat_template(messages, tokenize=False)
        assert isinstance(text, str)
        prompt_ids = mx.array(tokenizer.encode(text))[None, :]

        def loss_fn(embeds: mx.array) -> mx.array:
            return _compute_untargeted_loss(
                model, embeds, prompt_ids, 4, refusal_ids,
                None, 0.0,
            )

        loss_and_grad = mx.value_and_grad(loss_fn)
        loss_val, grad = loss_and_grad(soft_embeds)
        mx.eval(loss_val, grad)

        grad_norm = float(mx.linalg.norm(grad.reshape(-1)).item())
        assert grad_norm > 0, "Untargeted gradient is zero"

    def test_untargeted_with_direction(self) -> None:
        """Untargeted + RAID direction combined."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        direction = mx.random.normal((D_MODEL,))
        direction = direction / mx.linalg.norm(direction)
        mx.eval(direction)

        config = SoftPromptConfig(
            mode="continuous",
            n_tokens=4,
            n_steps=3,
            loss_mode="untargeted",
            direction_weight=0.1,
            direction_mode="raid",
            seed=42,
            max_gen_tokens=2,
        )

        result = softprompt_attack(
            model, tokenizer, ["test"], config, direction,
        )

        for loss in result.loss_history:
            assert not math.isnan(loss)

    def test_encode_refusal_tokens(self) -> None:
        """Produces non-empty array of valid token IDs."""
        tokenizer = MockTokenizer(VOCAB_SIZE)
        ids = _encode_refusal_tokens(tokenizer)
        mx.eval(ids)
        assert ids.ndim == 1
        assert ids.shape[0] > 0
        for i in range(ids.shape[0]):
            assert 0 <= int(ids[i].item()) < VOCAB_SIZE


# ---------------------------------------------------------------------------
# EGD attack tests
# ---------------------------------------------------------------------------


class TestEGDAttack:
    def test_basic_run(self) -> None:
        """mode='egd' returns valid SoftPromptResult with token_ids."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="egd",
            n_tokens=4,
            n_steps=5,
            learning_rate=0.1,
            seed=42,
            max_gen_tokens=2,
        )

        result = _egd_attack(
            model, tokenizer, ["test"], config, None,
        )

        assert isinstance(result, SoftPromptResult)
        assert result.mode == "egd"
        assert result.token_ids is not None
        assert len(result.token_ids) == 4
        assert result.token_text is not None
        assert result.embeddings is None
        assert len(result.loss_history) == 5
        assert len(result.eval_responses) == 1

    def test_egd_improves_loss(self) -> None:
        """Loss decreases over steps."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="egd",
            n_tokens=4,
            n_steps=30,
            learning_rate=0.5,
            seed=42,
            max_gen_tokens=2,
        )

        result = _egd_attack(
            model, tokenizer, ["test"], config, None,
        )

        first_loss = result.loss_history[0]
        best_loss = min(result.loss_history)
        assert best_loss < first_loss, (
            f"EGD did not improve: {first_loss:.4f} -> {best_loss:.4f}"
        )

    def test_egd_token_ids_in_range(self) -> None:
        """All token IDs in [0, vocab_size)."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="egd",
            n_tokens=4,
            n_steps=3,
            learning_rate=0.1,
            seed=42,
            max_gen_tokens=2,
        )

        result = _egd_attack(
            model, tokenizer, ["test"], config, None,
        )

        assert result.token_ids is not None
        for tid in result.token_ids:
            assert 0 <= tid < VOCAB_SIZE

    def test_egd_dispatch(self) -> None:
        """softprompt_attack dispatches correctly to EGD."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="egd",
            n_tokens=4,
            n_steps=3,
            learning_rate=0.1,
            seed=42,
            max_gen_tokens=2,
        )

        result = softprompt_attack(
            model, tokenizer, ["test"], config, None,
        )
        assert result.mode == "egd"


# ---------------------------------------------------------------------------
# Token constraint tests
# ---------------------------------------------------------------------------


class TestTokenConstraint:
    def test_build_vocab_mask_ascii(self) -> None:
        """Mask shape is (VOCAB_SIZE,) with some True entries."""
        tokenizer = MockTokenizer(VOCAB_SIZE)
        mask = _build_vocab_mask(tokenizer, VOCAB_SIZE, "ascii")
        assert mask is not None
        mx.eval(mask)
        assert mask.shape == (VOCAB_SIZE,)
        n_allowed = int(mx.sum(mask).item())
        assert n_allowed > 0
        assert n_allowed <= VOCAB_SIZE

    def test_build_vocab_mask_none(self) -> None:
        """Returns None when constraint is None."""
        tokenizer = MockTokenizer(VOCAB_SIZE)
        mask = _build_vocab_mask(tokenizer, VOCAB_SIZE, None)
        assert mask is None

    def test_gcg_with_constraint(self) -> None:
        """GCG + token_constraint='ascii' runs, all token IDs map to ASCII."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="gcg",
            n_tokens=4,
            n_steps=3,
            batch_size=4,
            top_k=8,
            seed=42,
            max_gen_tokens=2,
            token_constraint="ascii",
        )

        result = _gcg_attack(
            model, tokenizer, ["test"], config, None,
        )

        assert result.token_ids is not None
        for tid in result.token_ids:
            decoded = tokenizer.decode([tid])
            assert all(32 <= ord(c) < 127 for c in decoded)

    def test_egd_with_constraint(self) -> None:
        """EGD + token_constraint='ascii' runs, all token IDs map to ASCII."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="egd",
            n_tokens=4,
            n_steps=3,
            learning_rate=0.1,
            seed=42,
            max_gen_tokens=2,
            token_constraint="ascii",
        )

        result = _egd_attack(
            model, tokenizer, ["test"], config, None,
        )

        assert result.token_ids is not None
        for tid in result.token_ids:
            decoded = tokenizer.decode([tid])
            assert all(32 <= ord(c) < 127 for c in decoded)


# ---------------------------------------------------------------------------
# EOS loss tests
# ---------------------------------------------------------------------------


class TestEOSLoss:
    def test_eos_force_runs(self) -> None:
        """Continuous + eos_loss_mode='force' produces finite loss."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="continuous",
            n_tokens=4,
            n_steps=3,
            seed=42,
            max_gen_tokens=2,
            eos_loss_mode="force",
            eos_loss_weight=0.1,
        )

        result = _continuous_attack(
            model, tokenizer, ["test"], config, None,
        )

        for loss in result.loss_history:
            assert not math.isnan(loss), "EOS force loss produced NaN"

    def test_eos_suppress_runs(self) -> None:
        """Continuous + eos_loss_mode='suppress' produces finite loss."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="continuous",
            n_tokens=4,
            n_steps=3,
            seed=42,
            max_gen_tokens=2,
            eos_loss_mode="suppress",
            eos_loss_weight=0.1,
        )

        result = _continuous_attack(
            model, tokenizer, ["test"], config, None,
        )

        for loss in result.loss_history:
            assert not math.isnan(loss), "EOS suppress loss produced NaN"

    def test_eos_loss_helper_force(self) -> None:
        """Force loss decreases as P(EOS) increases."""
        # Create logits where EOS token has varying probability
        vocab_size = 8
        eos_id = 7

        # Low P(EOS): uniform logits
        logits_low = mx.zeros((1, 2, vocab_size))
        mx.eval(logits_low)
        loss_low = _compute_eos_loss(logits_low, 0, eos_id, "force")
        mx.eval(loss_low)

        # High P(EOS): boost EOS logit
        logits_high = mx.zeros((1, 2, vocab_size))
        logits_high = logits_high.at[:, 0, eos_id].add(mx.array(10.0))
        mx.eval(logits_high)
        loss_high = _compute_eos_loss(logits_high, 0, eos_id, "force")
        mx.eval(loss_high)

        assert float(loss_high.item()) < float(loss_low.item())

    def test_eos_loss_helper_suppress(self) -> None:
        """Suppress loss decreases as P(EOS) decreases."""
        vocab_size = 8
        eos_id = 7

        # High P(EOS): boost EOS logit
        logits_high = mx.zeros((1, 2, vocab_size))
        logits_high = logits_high.at[:, 0, eos_id].add(mx.array(10.0))
        mx.eval(logits_high)
        loss_high = _compute_eos_loss(logits_high, 0, eos_id, "suppress")
        mx.eval(loss_high)

        # Low P(EOS): uniform logits
        logits_low = mx.zeros((1, 2, vocab_size))
        mx.eval(logits_low)
        loss_low = _compute_eos_loss(logits_low, 0, eos_id, "suppress")
        mx.eval(loss_low)

        assert float(loss_low.item()) < float(loss_high.item())


# ---------------------------------------------------------------------------
# KL collision loss tests
# ---------------------------------------------------------------------------


class TestKLCollisionLoss:
    def test_kl_collision_runs(self) -> None:
        """Continuous + kl_ref_weight=0.1 with ref model runs."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        ref_model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(ref_model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="continuous",
            n_tokens=4,
            n_steps=3,
            seed=42,
            max_gen_tokens=2,
            kl_ref_weight=0.1,
        )

        result = _continuous_attack(
            model, tokenizer, ["test"], config, None,
            ref_model=ref_model,
        )

        for loss in result.loss_history:
            assert not math.isnan(loss), "KL collision loss produced NaN"

    def test_kl_collision_same_model_low(self) -> None:
        """Same model as reference produces near-zero KL."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        soft_embeds = mx.random.normal((1, 4, D_MODEL)) * 0.1
        mx.eval(soft_embeds)

        messages = [{"role": "user", "content": "test"}]
        text = tokenizer.apply_chat_template(messages, tokenize=False)
        assert isinstance(text, str)
        prompt_ids = mx.array(tokenizer.encode(text))[None, :]

        kl = _compute_kl_collision_loss(model, model, soft_embeds, prompt_ids, 4)
        mx.eval(kl)
        assert float(kl.item()) < 0.01

    def test_kl_collision_gradient_nonzero(self) -> None:
        """Gradient through KL loss is nonzero."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        ref_model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(ref_model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        soft_embeds = mx.random.normal((1, 4, D_MODEL)) * 0.1
        mx.eval(soft_embeds)

        messages = [{"role": "user", "content": "test"}]
        text = tokenizer.apply_chat_template(messages, tokenize=False)
        assert isinstance(text, str)
        prompt_ids = mx.array(tokenizer.encode(text))[None, :]

        def loss_fn(embeds: mx.array) -> mx.array:
            return _compute_kl_collision_loss(
                model, ref_model, embeds, prompt_ids, 4,
            )

        loss_and_grad = mx.value_and_grad(loss_fn)
        loss_val, grad = loss_and_grad(soft_embeds)
        mx.eval(loss_val, grad)

        grad_norm = float(mx.linalg.norm(grad.reshape(-1)).item())
        assert grad_norm > 0, "KL collision gradient is zero"


# ---------------------------------------------------------------------------
# New config defaults tests (Geiping features)
# ---------------------------------------------------------------------------


class TestNewConfigDefaults2:
    def test_constraint_default_none(self) -> None:
        cfg = SoftPromptConfig()
        assert cfg.token_constraint is None

    def test_eos_defaults(self) -> None:
        cfg = SoftPromptConfig()
        assert cfg.eos_loss_mode == "none"
        assert cfg.eos_loss_weight == 0.0

    def test_kl_ref_default(self) -> None:
        cfg = SoftPromptConfig()
        assert cfg.kl_ref_weight == 0.0


# ---------------------------------------------------------------------------
# Worst-K prompt selection tests
# ---------------------------------------------------------------------------


class TestSelectWorstKPromptIds:
    def test_returns_correct_count(self) -> None:
        """Returns exactly k prompts when k < len(all_ids)."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        all_ids = _pre_encode_prompts(tokenizer, ["a", "b", "c", "d"])
        target_ids = _encode_targets(tokenizer, ["Sure"])
        mx.eval(target_ids)

        soft_embeds = mx.random.normal((1, 4, D_MODEL)) * 0.1
        mx.eval(soft_embeds)

        selected = _select_worst_k_prompt_ids(
            model, soft_embeds, all_ids, target_ids,
            4, 2,  # k=2
            None, 0.0, "last", None,
            None, "none", 0.0,
            None, 0.0,
            loss_mode="targeted", refusal_ids=None,
        )

        assert len(selected) == 2

    def test_k_greater_than_len(self) -> None:
        """Returns all prompts when k >= len(all_ids)."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        all_ids = _pre_encode_prompts(tokenizer, ["a", "b"])
        target_ids = _encode_targets(tokenizer, ["Sure"])
        mx.eval(target_ids)

        soft_embeds = mx.random.normal((1, 4, D_MODEL)) * 0.1
        mx.eval(soft_embeds)

        selected = _select_worst_k_prompt_ids(
            model, soft_embeds, all_ids, target_ids,
            4, 10,  # k=10 > len=2
            None, 0.0, "last", None,
            None, "none", 0.0,
            None, 0.0,
            loss_mode="targeted", refusal_ids=None,
        )

        assert len(selected) == 2

    def test_returns_highest_loss(self) -> None:
        """Selected prompts should be the ones with highest loss."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        all_ids = _pre_encode_prompts(tokenizer, ["a", "b", "c"])
        target_ids = _encode_targets(tokenizer, ["Sure"])
        mx.eval(target_ids)

        soft_embeds = mx.random.normal((1, 4, D_MODEL)) * 0.1
        mx.eval(soft_embeds)

        selected = _select_worst_k_prompt_ids(
            model, soft_embeds, all_ids, target_ids,
            4, 1,  # k=1
            None, 0.0, "last", None,
            None, "none", 0.0,
            None, 0.0,
            loss_mode="targeted", refusal_ids=None,
        )

        assert len(selected) == 1
        # Verify it's a valid prompt from the original set
        assert selected[0].shape == all_ids[0].shape or True


class TestWorstKIntegration:
    def test_continuous_worst_k(self) -> None:
        """Continuous mode with worst_k strategy runs."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="continuous",
            n_tokens=4,
            n_steps=3,
            seed=42,
            max_gen_tokens=2,
            prompt_strategy="worst_k",
            worst_k=2,
        )

        result = _continuous_attack(
            model, tokenizer, ["hello", "world", "test"], config, None,
        )

        assert result.mode == "continuous"
        for loss in result.loss_history:
            assert not math.isnan(loss)

    def test_gcg_worst_k(self) -> None:
        """GCG mode with worst_k strategy runs."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="gcg",
            n_tokens=4,
            n_steps=2,
            batch_size=4,
            top_k=8,
            seed=42,
            max_gen_tokens=2,
            prompt_strategy="worst_k",
            worst_k=1,
        )

        result = _gcg_attack(
            model, tokenizer, ["hello", "world"], config, None,
        )

        assert result.mode == "gcg"
        assert len(result.loss_history) == 2

    def test_egd_worst_k(self) -> None:
        """EGD mode with worst_k strategy runs."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="egd",
            n_tokens=4,
            n_steps=3,
            learning_rate=0.1,
            seed=42,
            max_gen_tokens=2,
            prompt_strategy="worst_k",
            worst_k=1,
        )

        result = _egd_attack(
            model, tokenizer, ["hello", "world"], config, None,
        )

        assert result.mode == "egd"
        for loss in result.loss_history:
            assert not math.isnan(loss)


# ---------------------------------------------------------------------------
# Gradient accumulation tests
# ---------------------------------------------------------------------------


class TestSplitIntoBatches:
    def test_single_batch(self) -> None:
        items = [mx.array([[1]]), mx.array([[2]]), mx.array([[3]])]
        batches = _split_into_batches(items, 1)
        assert len(batches) == 1
        assert len(batches[0]) == 3

    def test_correct_batch_count(self) -> None:
        items = [mx.array([[i]]) for i in range(6)]
        batches = _split_into_batches(items, 3)
        assert len(batches) == 3
        assert sum(len(b) for b in batches) == 6

    def test_n_greater_than_len(self) -> None:
        items = [mx.array([[1]]), mx.array([[2]])]
        batches = _split_into_batches(items, 10)
        assert len(batches) == 2
        assert sum(len(b) for b in batches) == 2

    def test_empty_list(self) -> None:
        batches = _split_into_batches([], 3)
        assert len(batches) == 1
        assert len(batches[0]) == 0

    def test_uneven_split(self) -> None:
        items = [mx.array([[i]]) for i in range(5)]
        batches = _split_into_batches(items, 3)
        assert len(batches) == 3
        assert sum(len(b) for b in batches) == 5


class TestGradAccumIntegration:
    def test_continuous_grad_accum(self) -> None:
        """Continuous mode with grad_accum_steps=2 runs and produces finite loss."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="continuous",
            n_tokens=4,
            n_steps=3,
            seed=42,
            max_gen_tokens=2,
            prompt_strategy="all",
            grad_accum_steps=2,
        )

        result = _continuous_attack(
            model, tokenizer, ["hello", "world"], config, None,
        )

        assert result.mode == "continuous"
        for loss in result.loss_history:
            assert not math.isnan(loss)

    def test_gcg_grad_accum(self) -> None:
        """GCG mode with grad_accum_steps=2 runs."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="gcg",
            n_tokens=4,
            n_steps=2,
            batch_size=4,
            top_k=8,
            seed=42,
            max_gen_tokens=2,
            prompt_strategy="all",
            grad_accum_steps=2,
        )

        result = _gcg_attack(
            model, tokenizer, ["hello", "world"], config, None,
        )

        assert result.mode == "gcg"
        for loss in result.loss_history:
            assert not math.isnan(loss)

    def test_egd_grad_accum(self) -> None:
        """EGD mode with grad_accum_steps=2 runs."""
        model = MockCausalLM(D_MODEL, NUM_LAYERS, VOCAB_SIZE, NUM_HEADS)
        mx.eval(model.parameters())
        tokenizer = MockTokenizer(VOCAB_SIZE)

        config = SoftPromptConfig(
            mode="egd",
            n_tokens=4,
            n_steps=3,
            learning_rate=0.1,
            seed=42,
            max_gen_tokens=2,
            prompt_strategy="all",
            grad_accum_steps=2,
        )

        result = _egd_attack(
            model, tokenizer, ["hello", "world"], config, None,
        )

        assert result.mode == "egd"
        for loss in result.loss_history:
            assert not math.isnan(loss)


# ---------------------------------------------------------------------------
# Transfer evaluation tests
# ---------------------------------------------------------------------------


class TestProjectToTokens:
    def test_returns_valid_token_ids(self) -> None:
        """Returns correct count of token IDs in vocab range."""
        n_tokens = 4
        soft_embeds = mx.random.normal((1, n_tokens, D_MODEL))
        embed_matrix = mx.random.normal((VOCAB_SIZE, D_MODEL))
        mx.eval(soft_embeds, embed_matrix)

        token_ids = _project_to_tokens(soft_embeds, embed_matrix)
        assert len(token_ids) == n_tokens
        for tid in token_ids:
            assert 0 <= tid < VOCAB_SIZE

    def test_single_token(self) -> None:
        soft_embeds = mx.random.normal((1, 1, D_MODEL))
        embed_matrix = mx.random.normal((VOCAB_SIZE, D_MODEL))
        mx.eval(soft_embeds, embed_matrix)

        token_ids = _project_to_tokens(soft_embeds, embed_matrix)
        assert len(token_ids) == 1

    def test_nearest_neighbor_correctness(self) -> None:
        """Embedding of token i should project back to token i."""
        embed_matrix = mx.random.normal((VOCAB_SIZE, D_MODEL))
        mx.eval(embed_matrix)
        # Use exact embedding for token 3
        soft_embeds = embed_matrix[3:4][None, :]  # shape (1, 1, D_MODEL)
        mx.eval(soft_embeds)
        token_ids = _project_to_tokens(soft_embeds, embed_matrix)
        assert token_ids[0] == 3


class TestTransferEvalResult:
    def test_construction(self) -> None:
        result = TransferEvalResult(
            model_id="test-model",
            success_rate=0.5,
            eval_responses=["resp1", "resp2"],
        )
        assert result.model_id == "test-model"
        assert result.success_rate == 0.5
        assert len(result.eval_responses) == 2

    def test_frozen(self) -> None:
        result = TransferEvalResult(
            model_id="test", success_rate=0.0, eval_responses=[],
        )
        with pytest.raises(AttributeError):
            result.success_rate = 1.0  # type: ignore[misc]


class TestSoftPromptResultTransfer:
    def test_default_empty(self) -> None:
        result = SoftPromptResult(
            mode="continuous",
            success_rate=0.5,
            final_loss=2.0,
            loss_history=[2.0],
            n_steps=1,
            n_tokens=4,
            embeddings=None,
            token_ids=None,
            token_text=None,
            eval_responses=[],
        )
        assert result.transfer_results == []

    def test_with_transfer_results(self) -> None:
        tr = TransferEvalResult(
            model_id="other-model",
            success_rate=0.3,
            eval_responses=["r1"],
        )
        result = SoftPromptResult(
            mode="gcg",
            success_rate=0.9,
            final_loss=0.5,
            loss_history=[0.5],
            n_steps=1,
            n_tokens=4,
            embeddings=None,
            token_ids=[1, 2, 3, 4],
            token_text="test",
            eval_responses=["resp"],
            transfer_results=[tr],
        )
        assert len(result.transfer_results) == 1
        assert result.transfer_results[0].model_id == "other-model"


class TestSoftPromptSerialization:
    def test_softprompt_to_dict_includes_transfer(self) -> None:
        from vauban._serializers import _softprompt_to_dict

        tr = TransferEvalResult(
            model_id="m1", success_rate=0.4, eval_responses=["r"],
        )
        result = SoftPromptResult(
            mode="gcg",
            success_rate=0.8,
            final_loss=1.0,
            loss_history=[1.0],
            n_steps=1,
            n_tokens=4,
            embeddings=None,
            token_ids=[1, 2, 3, 4],
            token_text="test",
            eval_responses=["resp"],
            transfer_results=[tr],
        )
        d = _softprompt_to_dict(result)
        assert "transfer_results" in d
        trs = d["transfer_results"]
        assert isinstance(trs, list)
        assert len(trs) == 1
        assert trs[0]["model_id"] == "m1"  # type: ignore[index]
        assert trs[0]["success_rate"] == 0.4  # type: ignore[index]

    def test_softprompt_to_dict_empty_transfer(self) -> None:
        from vauban._serializers import _softprompt_to_dict

        result = SoftPromptResult(
            mode="continuous",
            success_rate=0.5,
            final_loss=2.0,
            loss_history=[2.0],
            n_steps=1,
            n_tokens=4,
            embeddings=None,
            token_ids=None,
            token_text=None,
            eval_responses=[],
        )
        d = _softprompt_to_dict(result)
        assert d["transfer_results"] == []


# ---------------------------------------------------------------------------
# New config defaults for worst-k, grad accum, transfer
# ---------------------------------------------------------------------------


class TestNewConfigDefaults3:
    def test_worst_k_default(self) -> None:
        cfg = SoftPromptConfig()
        assert cfg.worst_k == 5

    def test_grad_accum_steps_default(self) -> None:
        cfg = SoftPromptConfig()
        assert cfg.grad_accum_steps == 1

    def test_transfer_models_default(self) -> None:
        cfg = SoftPromptConfig()
        assert cfg.transfer_models == []

    def test_custom_worst_k(self) -> None:
        cfg = SoftPromptConfig(worst_k=10)
        assert cfg.worst_k == 10

    def test_custom_grad_accum_steps(self) -> None:
        cfg = SoftPromptConfig(grad_accum_steps=4)
        assert cfg.grad_accum_steps == 4

    def test_custom_transfer_models(self) -> None:
        cfg = SoftPromptConfig(transfer_models=["model-a"])
        assert cfg.transfer_models == ["model-a"]
