"""Tests for the SOTA abliteration pipeline."""

from __future__ import annotations

import json
import math
from pathlib import Path
from unittest.mock import MagicMock, Mock, patch

import pytest
import torch
from transformers import GPT2Config, GPT2LMHeadModel, Qwen2Config, Qwen2ForCausalLM

from obliteratus.abliterate import (
    HARMFUL_PROMPTS,
    HARMLESS_PROMPTS,
    METHODS,
    STAGES,
    AbliterationPipeline,
    PipelineStage,
    StageResult,
    _is_coherent_completion,
)
from obliteratus.models.loader import ModelHandle


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

def _make_tiny_handle():
    """Create a minimal ModelHandle with a tiny GPT-2 for testing."""
    config = GPT2Config(
        vocab_size=1000,
        n_positions=128,
        n_embd=64,
        n_layer=4,
        n_head=2,
        n_inner=256,
    )
    model = GPT2LMHeadModel(config)
    model.eval()

    tokenizer = MagicMock()
    tokenizer.pad_token = "<pad>"
    tokenizer.eos_token = "<eos>"
    tokenizer.return_value = {
        "input_ids": torch.randint(0, 1000, (1, 10)),
        "attention_mask": torch.ones(1, 10, dtype=torch.long),
    }
    tokenizer.decode.return_value = "The capital of France is Paris, a beautiful city"

    handle = ModelHandle(
        model=model,
        tokenizer=tokenizer,
        config=config,
        model_name="gpt2-test",
        task="causal_lm",
    )
    handle.snapshot()
    return handle


def _make_varied_tokenizer(handle):
    """Set up a tokenizer mock that returns different tokens per call."""
    call_count = [0]
    def mock_tokenizer(prompt, **kwargs):
        call_count[0] += 1
        torch.manual_seed(call_count[0])
        batch_size = len(prompt) if isinstance(prompt, list) else 1
        return {
            "input_ids": torch.randint(0, 1000, (batch_size, 5)),
            "attention_mask": torch.ones(batch_size, 5, dtype=torch.long),
        }
    handle.tokenizer.side_effect = mock_tokenizer


@pytest.fixture
def handle():
    return _make_tiny_handle()


# ---------------------------------------------------------------------------
# Data & stage definitions
# ---------------------------------------------------------------------------

class TestPrompts:
    def test_harmful_prompts_expanded(self):
        assert len(HARMFUL_PROMPTS) >= 99

    def test_harmless_prompts_expanded(self):
        assert len(HARMLESS_PROMPTS) >= 99

    def test_prompt_lists_same_length(self):
        assert len(HARMFUL_PROMPTS) == len(HARMLESS_PROMPTS)

    def test_prompt_count_842(self):
        """842 prompts across the expanded built-in corpus."""
        assert len(HARMFUL_PROMPTS) == 842
        assert len(HARMLESS_PROMPTS) == 842

    def test_prompt_volume_slicing(self):
        """Slicing at standard volumes gives correct counts."""
        for n in (33, 66, 99, 256, 512):
            assert len(HARMFUL_PROMPTS[:n]) == n
            assert len(HARMLESS_PROMPTS[:n]) == n


class TestCoherenceScoring:
    @pytest.mark.parametrize(
        ("prompt", "completion"),
        [
            ("The capital of France is", "Paris, the country's largest city."),
            ("Water boils at a temperature of", "100 degrees Celsius at sea level."),
            ("A binary search algorithm works by", "dividing a sorted array in half around its middle element."),
        ],
    )
    def test_accepts_relevant_completion(self, prompt, completion):
        assert _is_coherent_completion(prompt, completion) is True

    def test_rejects_fluent_but_irrelevant_completion(self):
        assert _is_coherent_completion(
            "The capital of France is",
            "A computational model composed of many connected layers and nodes.",
        ) is False

    def test_rejects_corpus_contamination_seen_in_verifier(self):
        assert _is_coherent_completion(
            "Water boils at a temperature of",
            "100°C. What is the temperature in Fahrenheit? - Brainly.in\nprofile\nanswer\nanswered",
        ) is False

    def test_rejects_repetition(self):
        assert _is_coherent_completion(
            "The largest ocean on Earth is",
            "Pacific Pacific Pacific Pacific Pacific Pacific Pacific Pacific.",
        ) is False


class TestStages:
    def test_six_stages(self):
        assert len(STAGES) == 6

    def test_stage_keys(self):
        keys = [s.key for s in STAGES]
        assert keys == ["summon", "probe", "distill", "excise", "verify", "rebirth"]

    def test_stage_dataclass(self):
        stage = PipelineStage(key="test", name="TEST", description="A test stage")
        assert stage.key == "test"
        assert stage.name == "TEST"

    def test_stage_result_defaults(self):
        result = StageResult(stage="test", status="running")
        assert result.message == ""
        assert result.duration == 0.0
        assert result.details == {}


# ---------------------------------------------------------------------------
# Method presets
# ---------------------------------------------------------------------------

class TestMethods:
    def test_methods_exist(self):
        assert set(METHODS.keys()) == {"basic", "advanced", "aggressive", "informed", "surgical", "inverted", "nuclear", "optimized", "failspy", "gabliteration", "heretic", "rdo", "spectral_cascade", "som"}

    def test_basic_single_direction(self):
        cfg = METHODS["basic"]
        assert cfg["n_directions"] == 1
        assert cfg["norm_preserve"] is False
        assert cfg["regularization"] == 0.0
        assert cfg["refinement_passes"] == 1

    def test_advanced_multi_direction(self):
        cfg = METHODS["advanced"]
        assert cfg["n_directions"] > 1
        assert cfg["norm_preserve"] is True
        assert cfg["regularization"] > 0
        assert cfg["refinement_passes"] >= 2

    def test_aggressive_full_gabliteration(self):
        cfg = METHODS["aggressive"]
        assert cfg["n_directions"] >= 8
        assert cfg["norm_preserve"] is True
        assert cfg["refinement_passes"] >= 3


# ---------------------------------------------------------------------------
# Pipeline init
# ---------------------------------------------------------------------------

class TestPipelineInit:
    def test_default_prompts(self):
        pipeline = AbliterationPipeline(model_name="test-model")
        assert pipeline.harmful_prompts == HARMFUL_PROMPTS
        assert pipeline.harmless_prompts == HARMLESS_PROMPTS

    def test_custom_prompts(self):
        harmful = ["bad prompt"]
        harmless = ["good prompt"]
        pipeline = AbliterationPipeline(
            model_name="test-model",
            harmful_prompts=harmful,
            harmless_prompts=harmless,
        )
        assert pipeline.harmful_prompts == harmful
        assert pipeline.harmless_prompts == harmless

    def test_defaults(self):
        pipeline = AbliterationPipeline(model_name="test-model")
        assert pipeline.device == "auto"
        assert pipeline.dtype == "float16"
        assert pipeline.output_dir == Path("abliterated")
        assert pipeline.trust_remote_code is False
        assert pipeline.gpu_memory_utilization is None
        assert pipeline.refusal_max_tokens == 128
        assert pipeline.max_perplexity_increase == 3.0
        assert pipeline.min_coherence_retention == 0.5
        assert pipeline.max_degenerate_fraction == 0.2
        assert pipeline.handle is None

    @pytest.mark.parametrize(
        ("keyword", "value"),
        [
            ("max_perplexity_increase", 0.99),
            ("min_coherence_retention", 1.01),
            ("max_degenerate_fraction", -0.01),
        ],
    )
    def test_quality_guardrails_validate_configuration(self, keyword, value):
        with pytest.raises(ValueError):
            AbliterationPipeline(model_name="test-model", **{keyword: value})

    @pytest.mark.parametrize("architecture", ["qwen3_5", "qwen3_5_text", "qwen3_5_moe"])
    def test_unvalidated_qwen35_hybrid_fails_before_surgery(self, architecture):
        from types import SimpleNamespace

        from obliteratus.abliterate import PipelineValidationError

        pipeline = AbliterationPipeline(model_name="Qwen/Qwen3.8-27B")
        pipeline.handle = SimpleNamespace(architecture=architecture)

        with pytest.raises(PipelineValidationError, match="no validated projection allowlist"):
            pipeline._validate_architecture_surgery_support()

        assert pipeline._excise_modified_count is None
        assert pipeline._quality_metrics["architecture_support"] == 0.0

    def test_relative_perplexity_guardrail_fails_closed(self):
        from obliteratus.abliterate import PipelineValidationError

        pipeline = AbliterationPipeline(
            model_name="test-model",
            max_perplexity_increase=2.0,
        )
        pipeline._stock_baseline = {"perplexity": 10.0}

        with pytest.raises(PipelineValidationError) as exc_info:
            pipeline._enforce_perplexity_guardrail(21.0)

        assert exc_info.value.metric == "perplexity_increase"
        assert pipeline._quality_metrics["perplexity_increase"] == pytest.approx(2.1)

    @pytest.mark.parametrize(
        ("coherence", "degenerate", "metric"),
        [(0.3, 0.0, "coherence_retention"), (1.0, 0.21, "degenerate_fraction")],
    )
    def test_generation_guardrails_fail_closed(self, coherence, degenerate, metric):
        from obliteratus.abliterate import PipelineValidationError

        pipeline = AbliterationPipeline(model_name="test-model")
        pipeline._stock_baseline = {"coherence": 1.0}

        with pytest.raises(PipelineValidationError) as exc_info:
            pipeline._enforce_generation_guardrails(coherence, degenerate)

        assert exc_info.value.metric == metric

    def test_cancellation_is_terminal_and_cleanup_unloads_model(self):
        from threading import Event

        from obliteratus.abliterate import PipelineCancelledError

        cancellation = Event()
        pipeline = AbliterationPipeline(
            model_name="test-model",
            cancellation_event=cancellation,
        )
        pipeline.handle = MagicMock()
        cancellation.set()

        with pytest.raises(PipelineCancelledError, match="cancelled during PROBE"):
            pipeline._check_cancelled("probe")

        pipeline.cleanup_failed_run()
        assert pipeline.handle.model is None
        assert pipeline.handle.tokenizer is None

    def test_public_run_cleans_up_terminal_pipeline_failure(self, monkeypatch):
        from obliteratus.abliterate import PipelineValidationError

        pipeline = AbliterationPipeline(model_name="test-model")
        cleanup = Mock()
        monkeypatch.setattr(pipeline, "cleanup_failed_run", cleanup)
        monkeypatch.setattr(
            pipeline,
            "_run_pipeline",
            Mock(
                side_effect=PipelineValidationError(
                    "unsafe",
                    stage="verify",
                    metric="perplexity",
                ),
            ),
        )

        with pytest.raises(PipelineValidationError):
            pipeline.run()

        cleanup.assert_called_once_with()

    def test_catastrophic_perplexity_aborts_before_generation(self):
        from types import SimpleNamespace

        from obliteratus.abliterate import PipelineValidationError

        class NonFiniteModel(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.anchor = torch.nn.Parameter(torch.zeros(1))
                self.generate_calls = 0

            def forward(self, **_kwargs):
                return SimpleNamespace(loss=torch.tensor(float("nan")))

            def generate(self, **_kwargs):
                self.generate_calls += 1
                raise AssertionError("generation must not run after catastrophic validation")

        class Tokenizer:
            def __call__(self, _text, **_kwargs):
                return {
                    "input_ids": torch.tensor([[1, 2, 3]]),
                    "attention_mask": torch.ones((1, 3), dtype=torch.long),
                }

        model = NonFiniteModel()
        pipeline = AbliterationPipeline(model_name="test-model")
        pipeline.handle = SimpleNamespace(model=model, tokenizer=Tokenizer())
        stages = []
        pipeline._on_stage = stages.append

        with pytest.raises(
            PipelineValidationError,
            match="all reference losses were non-finite",
        ) as exc_info:
            pipeline._verify()

        assert exc_info.value.stage == "verify"
        assert exc_info.value.metric == "perplexity"
        assert math.isinf(pipeline._quality_metrics["perplexity"])
        assert model.generate_calls == 0
        assert stages[-1].status == "error"
        assert stages[-1].details["metric"] == "perplexity"

    @pytest.mark.parametrize("invalid", [0, -1, 1.5, True])
    def test_refusal_max_tokens_must_be_positive_integer(self, invalid):
        with pytest.raises(
            ValueError, match="refusal_max_tokens must be a positive integer",
        ):
            AbliterationPipeline(model_name="test-model", refusal_max_tokens=invalid)

    def test_refusal_max_tokens_accepts_positive_override(self):
        pipeline = AbliterationPipeline(
            model_name="test-model", refusal_max_tokens=512,
        )
        assert pipeline.refusal_max_tokens == 512

    @pytest.mark.parametrize(
        "invalid", [0, -0.1, 1.1, float("nan"), float("inf"), True, "0.8"],
    )
    def test_gpu_memory_utilization_requires_bounded_number(self, invalid):
        with pytest.raises(
            ValueError, match=r"gpu_memory_utilization must be a number in \(0, 1\]",
        ):
            AbliterationPipeline(
                model_name="test-model", gpu_memory_utilization=invalid,
            )

    def test_gpu_memory_utilization_accepts_explicit_override(self):
        pipeline = AbliterationPipeline(
            model_name="test-model", gpu_memory_utilization=0.95,
        )
        assert pipeline.gpu_memory_utilization == 0.95

    def test_summon_propagates_gpu_budget_without_forcing_snapshot_skip(self, monkeypatch):
        handle = MagicMock()
        handle.summary.return_value = {
            "architecture": "fixture",
            "num_layers": 2,
            "num_heads": 4,
            "hidden_size": 8,
            "total_params": 32,
        }
        load_model = Mock(return_value=handle)
        monkeypatch.setattr("obliteratus.abliterate.load_model", load_model)
        pipeline = AbliterationPipeline(
            model_name="test-model", gpu_memory_utilization=0.95,
        )

        pipeline._summon()

        assert load_model.call_args.kwargs["gpu_memory_utilization"] == 0.95
        assert "skip_snapshot" not in load_model.call_args.kwargs

    def test_default_method_is_advanced(self):
        pipeline = AbliterationPipeline(model_name="test-model")
        assert pipeline.method == "advanced"
        assert pipeline.n_directions == METHODS["advanced"]["n_directions"]
        assert pipeline.norm_preserve == METHODS["advanced"]["norm_preserve"]
        assert pipeline.regularization == METHODS["advanced"]["regularization"]

    def test_method_basic(self):
        pipeline = AbliterationPipeline(model_name="test-model", method="basic")
        assert pipeline.n_directions == 1
        assert pipeline.norm_preserve is False
        assert pipeline.regularization == 0.0

    def test_method_aggressive(self):
        pipeline = AbliterationPipeline(model_name="test-model", method="aggressive")
        assert pipeline.n_directions == 8
        assert pipeline.norm_preserve is True
        assert pipeline.refinement_passes == 3

    def test_explicit_overrides_method(self):
        pipeline = AbliterationPipeline(
            model_name="test-model",
            method="basic",
            n_directions=6,
            norm_preserve=True,
            regularization=0.5,
            refinement_passes=4,
        )
        assert pipeline.n_directions == 6
        assert pipeline.norm_preserve is True
        assert pipeline.regularization == 0.5
        assert pipeline.refinement_passes == 4

    def test_callbacks(self):
        stage_results = []
        log_msgs = []
        pipeline = AbliterationPipeline(
            model_name="test-model",
            on_stage=lambda r: stage_results.append(r),
            on_log=lambda m: log_msgs.append(m),
        )
        pipeline.log("hello")
        assert log_msgs == ["hello"]

        pipeline._emit("test", "running", "msg")
        assert len(stage_results) == 1
        assert stage_results[0].stage == "test"


# ---------------------------------------------------------------------------
# _project_out_advanced (norm-preserving + regularization)
# ---------------------------------------------------------------------------

class TestProjectOutAdvanced:
    def test_norm_preserving(self):
        """Norm-preserving mode should keep Frobenius norm constant."""
        class Wrapper(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.o_proj = torch.nn.Linear(4, 8, bias=False)

        module = Wrapper()
        torch.manual_seed(42)
        module.o_proj.weight.data = torch.randn(8, 4)
        original_norm = module.o_proj.weight.data.norm().item()

        direction = torch.randn(4, 1)
        direction = direction / direction.norm()

        AbliterationPipeline._project_out_advanced(
            module, direction, ["o_proj"], norm_preserve=True, regularization=0.0
        )

        new_norm = module.o_proj.weight.data.norm().item()
        # With amplification cap (1.10x max), exact norm preservation isn't
        # guaranteed on tiny matrices (hidden_dim=4) where a single direction
        # removes a large fraction of energy.  Verify the norm is closer to
        # original than the un-preserved norm would be (i.e. cap is working).
        # The new norm should be >= the un-preserved norm (cap restores some)
        assert new_norm >= original_norm * 0.85, \
            f"Norm should be approximately preserved (within cap): {original_norm:.4f} vs {new_norm:.4f}"

    def test_regularization_partial_removal(self):
        """Regularization should preserve some of the refusal component."""
        class Wrapper(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.o_proj = torch.nn.Linear(4, 8, bias=False)

        module_full = Wrapper()
        module_reg = Wrapper()
        torch.manual_seed(42)
        W_orig = torch.randn(8, 4)
        module_full.o_proj.weight.data = W_orig.clone()
        module_reg.o_proj.weight.data = W_orig.clone()

        direction = torch.randn(4, 1)
        direction = direction / direction.norm()

        # Full removal
        AbliterationPipeline._project_out_advanced(
            module_full, direction, ["o_proj"], norm_preserve=False, regularization=0.0
        )
        # Regularized (30% preserved)
        AbliterationPipeline._project_out_advanced(
            module_reg, direction, ["o_proj"], norm_preserve=False, regularization=0.3
        )

        W_full = module_full.o_proj.weight.data
        W_reg = module_reg.o_proj.weight.data

        # Full removal should have zero projection on direction
        proj_full = (W_full @ direction).norm().item()
        assert proj_full < 1e-4

        # Regularized should have non-zero projection (30% preserved)
        proj_reg = (W_reg @ direction).norm().item()
        proj_orig = (W_orig @ direction).norm().item()
        expected_ratio = 0.3
        actual_ratio = proj_reg / proj_orig if proj_orig > 0 else 0
        assert abs(actual_ratio - expected_ratio) < 0.05, \
            f"Expected ~{expected_ratio:.0%} preserved, got {actual_ratio:.0%}"

    def test_norm_preserving_transposed(self):
        """Norm-preserving should also work for transposed weights."""
        class Wrapper(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.c_proj = torch.nn.Linear(8, 4, bias=False)

        module = Wrapper()
        torch.manual_seed(42)
        module.c_proj.weight.data = torch.randn(4, 8)
        original_norm = module.c_proj.weight.data.norm().item()

        direction = torch.randn(4, 1)
        direction = direction / direction.norm()

        AbliterationPipeline._project_out_advanced(
            module, direction, ["c_proj"], norm_preserve=True, regularization=0.0
        )

        new_norm = module.c_proj.weight.data.norm().item()
        # With amplification cap (1.10x max), exact norm preservation isn't
        # guaranteed on tiny matrices where a single direction removes a large
        # fraction of energy.
        assert new_norm >= original_norm * 0.80, \
            f"Norm should be approximately preserved (within cap): {original_norm:.4f} vs {new_norm:.4f}"


# ---------------------------------------------------------------------------
# Full attention projection (q/k/v + o_proj)
# ---------------------------------------------------------------------------

class TestAttentionFullProjection:
    """Test that ALL attention weight matrices are projected (not just o_proj)."""

    def test_qkv_all_projected(self):
        """q_proj, k_proj, v_proj should all be projected alongside o_proj."""
        hidden = 16

        class FakeAttn(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.q_proj = torch.nn.Linear(hidden, hidden, bias=False)
                self.k_proj = torch.nn.Linear(hidden, hidden, bias=False)
                self.v_proj = torch.nn.Linear(hidden, hidden, bias=False)
                self.o_proj = torch.nn.Linear(hidden, hidden, bias=False)

        attn = FakeAttn()
        torch.manual_seed(42)
        for p in attn.parameters():
            p.data = torch.randn_like(p.data)

        originals = {
            name: getattr(attn, name).weight.data.clone()
            for name in ["q_proj", "k_proj", "v_proj", "o_proj"]
        }

        d = torch.randn(hidden, 1)
        d = d / d.norm()

        from obliteratus.abliterate import _ATTN_OUT_NAMES, _ATTN_IN_NAMES
        count = AbliterationPipeline._project_out_advanced(
            attn, d, _ATTN_OUT_NAMES + _ATTN_IN_NAMES,
        )

        assert count == 4, f"Should project 4 weights (q/k/v/o), got {count}"
        for name in ["q_proj", "k_proj", "v_proj", "o_proj"]:
            assert not torch.allclose(
                getattr(attn, name).weight.data, originals[name]
            ), f"{name} should be modified"

    def test_project_all_does_not_early_return(self):
        """_project_out_advanced should project ALL matching weights, not just first."""
        hidden = 16

        class FakeModule(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.up_proj = torch.nn.Linear(hidden, 32, bias=False)
                self.gate_proj = torch.nn.Linear(hidden, 32, bias=False)

        mod = FakeModule()
        torch.manual_seed(42)
        orig_up = mod.up_proj.weight.data.clone()
        orig_gate = mod.gate_proj.weight.data.clone()

        d = torch.randn(hidden, 1)
        d = d / d.norm()

        from obliteratus.abliterate import _FFN_IN_NAMES
        count = AbliterationPipeline._project_out_advanced(mod, d, _FFN_IN_NAMES)

        assert count == 2, f"Should project both up_proj and gate_proj, got {count}"
        assert not torch.allclose(mod.up_proj.weight.data, orig_up), "up_proj should be modified"
        assert not torch.allclose(mod.gate_proj.weight.data, orig_gate), "gate_proj should be modified"

    def test_lm_head_projection(self):
        """lm_head should be projectable via _project_out_advanced."""
        hidden = 16
        vocab = 100

        class FakeModel(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.lm_head = torch.nn.Linear(hidden, vocab, bias=False)

        model = FakeModel()
        torch.manual_seed(42)
        orig = model.lm_head.weight.data.clone()

        d = torch.randn(hidden, 1)
        d = d / d.norm()

        count = AbliterationPipeline._project_out_advanced(
            model, d, ["lm_head"], regularization=0.0,
        )

        assert count == 1, "Should project lm_head"
        assert not torch.allclose(model.lm_head.weight.data, orig), "lm_head should be modified"
        # Verify refusal direction is removed from lm_head
        proj = (model.lm_head.weight.data @ d).norm().item()
        assert proj < 1e-4, f"Refusal direction should be removed from lm_head, proj={proj}"


class TestKneeDetectionThreshold:
    """Test that knee detection uses 5% threshold to include more layers."""

    def test_five_percent_threshold_includes_more(self):
        """Layers between 5% and 10% of max should now be included."""
        # Layer norms: max=10.0, then several between 5%-10%
        sorted_layers = [(0, 10.0), (1, 8.0), (2, 6.0), (3, 0.7), (4, 0.6)]
        selected = AbliterationPipeline._select_layers_knee(sorted_layers)
        # 0.7 and 0.6 are 7% and 6% of max — should now be included (> 5% threshold)
        assert 3 in selected or 4 in selected, (
            f"Layers with 6-7% of max signal should be included, got {selected}"
        )


# ---------------------------------------------------------------------------
# MoE projection (router, shared expert, input/output, fused)
# ---------------------------------------------------------------------------

class TestProjectMoEExperts:
    """Test the full MoE projection pipeline: router, shared expert, experts."""

    def _make_direction(self, hidden_dim=16):
        d = torch.randn(hidden_dim, 1)
        return d / d.norm()

    def test_router_gate_projected(self):
        """Router/gate weight should have refusal direction removed."""
        hidden = 16
        n_experts = 4

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.gate = torch.nn.Linear(hidden, n_experts, bias=True)
                self.experts = torch.nn.ModuleList([
                    self._make_expert() for _ in range(n_experts)
                ])

            @staticmethod
            def _make_expert():
                m = torch.nn.Module()
                m.down_proj = torch.nn.Linear(hidden, 32, bias=False)
                m.up_proj = torch.nn.Linear(hidden, 32, bias=False)
                return m

        moe = FakeMoE()
        d = self._make_direction(hidden)
        W_gate_orig = moe.gate.weight.data.clone()

        count = AbliterationPipeline._project_moe_experts(moe, d)
        assert count > 0

        # Gate weight should have been modified
        assert not torch.allclose(moe.gate.weight.data, W_gate_orig), \
            "Router/gate weights should be projected"

        # The gate weight's projection onto the direction should be ~0
        proj = (moe.gate.weight.data @ d).norm().item()
        assert proj < 1e-4, f"Gate should have no component along refusal dir, got {proj}"

    def test_shared_expert_projected(self):
        """Shared expert (always-on) should have both input and output projected."""
        hidden = 16

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.gate = torch.nn.Linear(hidden, 2, bias=False)
                self.shared_expert = torch.nn.Module()
                self.shared_expert.down_proj = torch.nn.Linear(hidden, 32, bias=False)
                self.shared_expert.up_proj = torch.nn.Linear(hidden, 32, bias=False)
                self.experts = torch.nn.ModuleList([
                    self._make_expert() for _ in range(2)
                ])

            @staticmethod
            def _make_expert():
                m = torch.nn.Module()
                m.down_proj = torch.nn.Linear(hidden, 32, bias=False)
                m.up_proj = torch.nn.Linear(hidden, 32, bias=False)
                return m

        moe = FakeMoE()
        d = self._make_direction(hidden)
        shared_down_orig = moe.shared_expert.down_proj.weight.data.clone()
        shared_up_orig = moe.shared_expert.up_proj.weight.data.clone()

        count = AbliterationPipeline._project_moe_experts(moe, d)
        assert count > 0

        # Both shared expert output AND input projections should be modified
        assert not torch.allclose(moe.shared_expert.down_proj.weight.data, shared_down_orig), \
            "Shared expert output (down_proj) should be projected"
        assert not torch.allclose(moe.shared_expert.up_proj.weight.data, shared_up_orig), \
            "Shared expert input (up_proj) should be projected"

    def test_expert_input_projections_projected(self):
        """Expert input projections (up_proj, gate_proj) should also be modified."""
        hidden = 16

        class FakeExpert(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Linear(hidden, 32, bias=False)
                self.up_proj = torch.nn.Linear(hidden, 32, bias=False)
                self.gate_proj = torch.nn.Linear(hidden, 32, bias=False)

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(2)])

        moe = FakeMoE()
        d = self._make_direction(hidden)
        up_orig = moe.experts[0].up_proj.weight.data.clone()

        count = AbliterationPipeline._project_moe_experts(moe, d)

        # Each expert contributes 2 projections (output + input)
        # 2 experts * 2 = 4 minimum
        assert count >= 4, f"Expected >= 4 projections (out+in per expert), got {count}"

        assert not torch.allclose(moe.experts[0].up_proj.weight.data, up_orig), \
            "Expert input (up_proj) should be projected"

    def test_fused_3d_output_and_input(self):
        """Fused 3D parameter patterns (GPT-OSS style) should project both directions."""
        hidden = 16
        intermediate = 32
        n_experts = 4

        class FusedExperts(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Parameter(torch.randn(n_experts, intermediate, hidden))
                self.up_proj = torch.nn.Parameter(torch.randn(n_experts, intermediate, hidden))

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.experts = FusedExperts()

        moe = FakeMoE()
        d = self._make_direction(hidden)
        down_orig = moe.experts.down_proj.data.clone()
        up_orig = moe.experts.up_proj.data.clone()

        count = AbliterationPipeline._project_moe_experts(moe, d)

        # 4 experts output + 4 experts input = 8
        assert count == 8, f"Expected 8 fused projections, got {count}"

        assert not torch.allclose(moe.experts.down_proj.data, down_orig), \
            "Fused output (down_proj) should be projected"
        assert not torch.allclose(moe.experts.up_proj.data, up_orig), \
            "Fused input (up_proj) should be projected"

    def test_fused_3d_norm_preserve(self):
        """Fused 3D projections should preserve norms when requested."""
        hidden = 16
        intermediate = 32
        n_experts = 4

        class FusedExperts(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Parameter(torch.randn(n_experts, intermediate, hidden))

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.experts = FusedExperts()

        moe = FakeMoE()
        d = self._make_direction(hidden)

        # Record per-expert norms before
        orig_norms = [moe.experts.down_proj.data[i].norm().item() for i in range(n_experts)]

        AbliterationPipeline._project_moe_experts(moe, d, norm_preserve=True)

        # Check per-expert norms preserved
        for i in range(n_experts):
            new_norm = moe.experts.down_proj.data[i].norm().item()
            assert abs(orig_norms[i] - new_norm) < 1e-3, \
                f"Expert {i} norm not preserved: {orig_norms[i]:.4f} vs {new_norm:.4f}"

    def test_no_experts_returns_zero(self):
        """Module without experts attribute should return 0."""
        class NoMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.mlp = torch.nn.Linear(16, 32)

        moe = NoMoE()
        d = self._make_direction(16)
        assert AbliterationPipeline._project_moe_experts(moe, d) == 0

    def test_router_bias_projected(self):
        """Router bias should be projected when project_biases=True."""
        hidden = 16

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.gate = torch.nn.Linear(hidden, 4, bias=True)
                self.experts = torch.nn.ModuleList([
                    self._make_expert() for _ in range(4)
                ])

            @staticmethod
            def _make_expert():
                m = torch.nn.Module()
                m.down_proj = torch.nn.Linear(hidden, 32, bias=False)
                return m

        moe = FakeMoE()
        d = self._make_direction(hidden)
        bias_orig = moe.gate.bias.data.clone()

        count = AbliterationPipeline._project_moe_experts(moe, d, project_biases=True)

        # Gate has 4 outputs (num_experts), direction has 16 dims
        # bias shape (4,) != direction shape (16,), so bias won't match.
        # This is correct: router bias is (num_experts,), not (hidden_dim,),
        # so _project_bias won't modify it (shape mismatch is expected).
        assert torch.allclose(moe.gate.bias.data, bias_orig), (
            "Router bias should be unchanged when shape mismatches direction"
        )
        assert isinstance(count, int)
        assert count > 0  # expert weights should still be projected

    def test_router_auto_detection_fallback(self):
        """Unknown router name should be auto-detected and projected."""
        import warnings as w
        hidden = 16
        n_experts = 4

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                # Unusual router name not in _ROUTER_NAMES
                self.moe_gate_proj = torch.nn.Linear(hidden, n_experts, bias=False)
                self.experts = torch.nn.ModuleList([
                    self._make_expert() for _ in range(n_experts)
                ])

            @staticmethod
            def _make_expert():
                m = torch.nn.Module()
                m.down_proj = torch.nn.Linear(hidden, 32, bias=False)
                return m

        moe = FakeMoE()
        d = self._make_direction(hidden)
        gate_orig = moe.moe_gate_proj.weight.data.clone()

        with w.catch_warnings(record=True) as caught:
            w.simplefilter("always")
            AbliterationPipeline._project_moe_experts(moe, d)

        # Should auto-detect and project the unusual router name
        assert not torch.allclose(moe.moe_gate_proj.weight.data, gate_orig), \
            "Auto-detected router should be projected"

        # Should emit a warning about the auto-detection
        auto_detect_warnings = [
            x for x in caught
            if "auto-detected" in str(x.message)
        ]
        assert len(auto_detect_warnings) > 0, "Should warn about auto-detected router"

    def test_full_moe_all_components(self):
        """End-to-end: all MoE components should be modified together."""
        hidden = 16

        class FakeExpert(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Linear(hidden, 32, bias=False)
                self.up_proj = torch.nn.Linear(hidden, 32, bias=False)

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.gate = torch.nn.Linear(hidden, 4, bias=False)
                self.shared_expert = torch.nn.Module()
                self.shared_expert.down_proj = torch.nn.Linear(hidden, 32, bias=False)
                self.shared_expert.up_proj = torch.nn.Linear(hidden, 32, bias=False)
                self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(4)])

        moe = FakeMoE()
        d = self._make_direction(hidden)

        count = AbliterationPipeline._project_moe_experts(moe, d)

        # Expected: 1 (gate) + 2 (shared out+in) + 4*2 (expert out+in) = 11
        assert count == 11, f"Expected 11 total projections, got {count}"


# ---------------------------------------------------------------------------
# SOTA technique #1: Safety-neuron masking (GateBreaker-style z-score)
# ---------------------------------------------------------------------------

class TestSafetyNeuronMasking:
    def test_outlier_neurons_zeroed(self):
        """Neurons with outsized refusal projection should be zeroed."""
        hidden = 16

        class Wrapper(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Linear(hidden, 64, bias=False)

        module = Wrapper()
        torch.manual_seed(42)
        # Inject a few rows with very high projection along direction
        d = torch.randn(hidden, 1)
        d = d / d.norm()
        # Make rows 0,1,2 have huge projection (outliers)
        for i in range(3):
            module.down_proj.weight.data[i] = d.squeeze() * 10.0

        n_masked = AbliterationPipeline._mask_safety_neurons(
            module, d, ["down_proj"], z_threshold=2.0,
        )

        assert n_masked >= 3, f"Expected >= 3 masked neurons, got {n_masked}"
        # Masked rows should be zero
        for i in range(3):
            assert module.down_proj.weight.data[i].abs().max().item() < 1e-6

    def test_no_outliers_no_masking(self):
        """When all neurons have similar projection, none should be masked."""
        hidden = 16

        class Wrapper(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Linear(hidden, 32, bias=False)

        module = Wrapper()
        # Uniform weights — no outliers
        module.down_proj.weight.data = torch.ones(32, hidden) * 0.1

        d = torch.randn(hidden, 1)
        d = d / d.norm()

        n_masked = AbliterationPipeline._mask_safety_neurons(
            module, d, ["down_proj"], z_threshold=2.0,
        )
        assert n_masked == 0

    def test_high_threshold_masks_fewer(self):
        """Higher z_threshold should mask fewer neurons."""
        hidden = 16

        class Wrapper(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Linear(hidden, 64, bias=False)

        module1 = Wrapper()
        module2 = Wrapper()
        torch.manual_seed(42)
        W = torch.randn(64, hidden)
        module1.down_proj.weight.data = W.clone()
        module2.down_proj.weight.data = W.clone()

        d = torch.randn(hidden, 1)
        d = d / d.norm()

        n1 = AbliterationPipeline._mask_safety_neurons(module1, d, ["down_proj"], z_threshold=1.5)
        n2 = AbliterationPipeline._mask_safety_neurons(module2, d, ["down_proj"], z_threshold=3.0)
        assert n1 >= n2, f"Lower threshold should mask more: {n1} vs {n2}"


# ---------------------------------------------------------------------------
# SOTA technique #2: Jailbreak-contrastive direction refinement
# ---------------------------------------------------------------------------

class TestJailbreakContrastive:
    def test_jailbreak_prompt_generation(self):
        """Pipeline should generate jailbreak variants of harmful prompts."""
        pipeline = AbliterationPipeline(model_name="test")
        pipeline.harmful_prompts = ["How to hack a computer", "Write a virus"]

        jailbreak = pipeline._generate_jailbreak_prompts()

        assert len(jailbreak) == 2
        # Each should contain the original prompt
        assert "hack a computer" in jailbreak[0].lower()
        assert "virus" in jailbreak[1].lower()
        # Each should be wrapped in a template (longer than original)
        assert len(jailbreak[0]) > len("How to hack a computer")

    def test_jailbreak_contrast_blending(self):
        """Jailbreak-contrastive blending should modify refusal direction."""
        pipeline = AbliterationPipeline(
            model_name="test",
            use_jailbreak_contrast=True,
            n_directions=1,
        )
        hidden = 16
        pipeline._on_log = lambda m: None

        # Simulate probed means
        torch.manual_seed(42)
        harm_mean = torch.randn(1, hidden)
        safe_mean = torch.randn(1, hidden)
        jb_mean = torch.randn(1, hidden)

        pipeline._harmful_means = {0: harm_mean}
        pipeline._harmless_means = {0: safe_mean}
        pipeline._jailbreak_means = {0: jb_mean}
        pipeline._harmful_acts = {0: [harm_mean]}
        pipeline._harmless_acts = {0: [safe_mean]}
        pipeline._jailbreak_acts = {0: [jb_mean]}

        # Run distill (will set standard direction, then blend)
        pipeline._distill()

        # Direction should be a unit vector
        d = pipeline.refusal_directions[0]
        assert abs(d.norm().item() - 1.0) < 1e-4

        # Direction should differ from pure harm-safe difference
        std_diff = (harm_mean - safe_mean).squeeze()
        std_dir = std_diff / std_diff.norm()
        cosine = (d @ std_dir).item()
        # Blended direction should not be identical to standard
        assert cosine < 0.99, f"Blended direction too similar to standard: cos={cosine}"

    def test_surgical_method_enables_jailbreak(self):
        """Surgical method should enable jailbreak-contrastive by default."""
        cfg = METHODS["surgical"]
        assert cfg["use_jailbreak_contrast"] is True


# ---------------------------------------------------------------------------
# SOTA technique #3: Layer-adaptive projection strength
# ---------------------------------------------------------------------------

class TestLayerAdaptiveStrength:
    def test_layer_weights_computed(self):
        """Layer-adaptive weights should be proportional to refusal signal."""
        pipeline = AbliterationPipeline(
            model_name="test",
            layer_adaptive_strength=True,
            n_directions=1,
        )
        hidden = 16
        pipeline._on_log = lambda m: None

        # Simulate: layer 0 has strong signal, layer 1 weak
        torch.manual_seed(42)
        strong_diff = torch.randn(1, hidden) * 10.0
        weak_diff = torch.randn(1, hidden) * 1.0
        zero_mean = torch.zeros(1, hidden)

        pipeline._harmful_means = {0: strong_diff, 1: weak_diff}
        pipeline._harmless_means = {0: zero_mean, 1: zero_mean}
        pipeline._harmful_acts = {0: [strong_diff], 1: [weak_diff]}
        pipeline._harmless_acts = {0: [zero_mean], 1: [zero_mean]}

        pipeline._distill()

        # Layer weights should exist for strong layers
        assert len(pipeline._layer_excise_weights) > 0
        # Strongest layer should have weight ~1.0
        max_weight = max(pipeline._layer_excise_weights.values())
        assert max_weight > 0.9, f"Max weight should be ~1.0, got {max_weight}"

    def test_surgical_method_enables_adaptive(self):
        """Surgical method should enable layer-adaptive by default."""
        cfg = METHODS["surgical"]
        assert cfg["layer_adaptive_strength"] is True


# ---------------------------------------------------------------------------
# SOTA technique #5: Attention head surgery
# ---------------------------------------------------------------------------

class TestAttentionHeadSurgery:
    def test_head_selective_projection(self):
        """Selective head projection should only modify targeted head rows."""
        hidden = 16
        n_heads = 4
        head_dim = hidden // n_heads

        class FakeAttn(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.o_proj = torch.nn.Linear(hidden, hidden, bias=False)

        attn = FakeAttn()
        torch.manual_seed(42)
        W_orig = attn.o_proj.weight.data.clone()

        d = torch.randn(hidden, 1)
        d = d / d.norm()

        # Head scores: head 0 is top safety head, head 3 is lowest
        head_scores = [(0, 5.0), (1, 3.0), (2, 1.0), (3, 0.5)]

        n_modified = AbliterationPipeline._project_head_selective(
            attn, d, head_scores, n_heads=n_heads, head_fraction=0.25,
        )

        assert n_modified >= 1, "Should modify at least 1 head"

        W_new = attn.o_proj.weight.data
        # Head 0 columns (targeted) should be modified
        assert not torch.allclose(
            W_new[:, 0:head_dim], W_orig[:, 0:head_dim]
        ), "Targeted head 0 should be modified"

        # Head 3 columns (NOT targeted) should be untouched
        assert torch.allclose(
            W_new[:, 3*head_dim:4*head_dim],
            W_orig[:, 3*head_dim:4*head_dim],
        ), "Non-targeted head 3 should be untouched"

    def test_head_surgery_norm_preserve(self):
        """Head surgery with norm_preserve should maintain per-head norms."""
        hidden = 16
        n_heads = 4
        head_dim = hidden // n_heads

        class FakeAttn(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.o_proj = torch.nn.Linear(hidden, hidden, bias=False)

        attn = FakeAttn()
        torch.manual_seed(42)

        d = torch.randn(hidden, 1)
        d = d / d.norm()

        orig_norms = [
            attn.o_proj.weight.data[:, h*head_dim:(h+1)*head_dim].norm().item()
            for h in range(n_heads)
        ]

        head_scores = [(0, 5.0), (1, 3.0), (2, 1.0), (3, 0.5)]
        AbliterationPipeline._project_head_selective(
            attn, d, head_scores, n_heads=n_heads,
            head_fraction=0.5, norm_preserve=True,
        )

        # Targeted heads should have preserved norms
        for h in range(2):  # top 50% = 2 heads
            new_norm = attn.o_proj.weight.data[:, h*head_dim:(h+1)*head_dim].norm().item()
            assert abs(orig_norms[h] - new_norm) < 1e-3, \
                f"Head {h} norm not preserved: {orig_norms[h]:.4f} vs {new_norm:.4f}"

    def test_head_surgery_non_square_gqa(self):
        """Head surgery should work for GQA models with non-square o_proj (attn_dim != hidden_dim)."""
        hidden_dim = 12   # model hidden dimension
        attn_dim = 32      # attention dimension (n_heads * head_dim_attn)
        n_heads = 4
        head_dim_attn = attn_dim // n_heads  # 8

        class FakeAttnGQA(torch.nn.Module):
            def __init__(self):
                super().__init__()
                # o_proj maps attn_dim -> hidden_dim
                # nn.Linear weight shape: (hidden_dim, attn_dim) = (12, 32)
                self.o_proj = torch.nn.Linear(attn_dim, hidden_dim, bias=False)

        attn = FakeAttnGQA()
        torch.manual_seed(42)
        attn.o_proj.weight.data = torch.randn(hidden_dim, attn_dim)
        W_orig = attn.o_proj.weight.data.clone()

        d = torch.randn(hidden_dim, 1)
        d = d / d.norm()

        head_scores = [(0, 5.0), (1, 3.0), (2, 1.0), (3, 0.5)]

        n_modified = AbliterationPipeline._project_head_selective(
            attn, d, head_scores, n_heads=n_heads, head_fraction=0.25,
        )

        assert n_modified >= 1, "Should modify at least 1 head"

        W_new = attn.o_proj.weight.data
        # Head 0 columns (targeted) should be modified
        assert not torch.allclose(
            W_new[:, 0:head_dim_attn], W_orig[:, 0:head_dim_attn]
        ), "Targeted head 0 should be modified"

        # Head 3 columns (NOT targeted) should be untouched
        assert torch.allclose(
            W_new[:, 3*head_dim_attn:4*head_dim_attn],
            W_orig[:, 3*head_dim_attn:4*head_dim_attn],
        ), "Non-targeted head 3 should be untouched"

    def test_head_surgery_gqa_norm_preserve(self):
        """Head surgery on GQA non-square o_proj with norm_preserve."""
        hidden_dim = 12
        attn_dim = 32
        n_heads = 4
        head_dim_attn = attn_dim // n_heads

        class FakeAttnGQA(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.o_proj = torch.nn.Linear(attn_dim, hidden_dim, bias=False)

        attn = FakeAttnGQA()
        torch.manual_seed(42)
        attn.o_proj.weight.data = torch.randn(hidden_dim, attn_dim)

        d = torch.randn(hidden_dim, 1)
        d = d / d.norm()

        orig_norms = [
            attn.o_proj.weight.data[:, h*head_dim_attn:(h+1)*head_dim_attn].norm().item()
            for h in range(n_heads)
        ]

        head_scores = [(0, 5.0), (1, 3.0), (2, 1.0), (3, 0.5)]
        AbliterationPipeline._project_head_selective(
            attn, d, head_scores, n_heads=n_heads,
            head_fraction=0.5, norm_preserve=True,
        )

        for h in range(2):  # top 50% = 2 heads
            new_norm = attn.o_proj.weight.data[:, h*head_dim_attn:(h+1)*head_dim_attn].norm().item()
            assert abs(orig_norms[h] - new_norm) < 1e-3, \
                f"GQA head {h} norm not preserved: {orig_norms[h]:.4f} vs {new_norm:.4f}"


# ---------------------------------------------------------------------------
# SOTA technique #6: SAE feature-level abliteration
# ---------------------------------------------------------------------------

@pytest.mark.filterwarnings(
    "ignore:SAE held-out reconstruction MSE.*:UserWarning",
)
class TestSAEAbliteration:
    """Exercise deliberately undertrained SAEs without weakening global warnings."""

    def test_sae_train_and_reconstruct(self):
        """SAE should train and reconstruct activations."""
        from obliteratus.analysis.sae_abliteration import train_sae

        hidden = 32
        # Generate synthetic activations
        torch.manual_seed(42)
        acts = [torch.randn(hidden) for _ in range(64)]

        sae = train_sae(acts, hidden, expansion=2, n_epochs=10, lr=1e-3, device="cpu")

        # Forward pass should work
        x = torch.randn(1, hidden)
        x_hat, z = sae(x)
        assert x_hat.shape == x.shape
        assert z.shape == (1, 2 * hidden)  # expansion=2

        # Z should be sparse (ReLU activation)
        assert (z == 0).float().mean() > 0.3, "Features should be sparse"

    def test_refusal_feature_identification(self):
        """SAE should identify features that differ between harmful/harmless."""
        from obliteratus.analysis.sae_abliteration import (
            train_sae, identify_refusal_features,
        )

        hidden = 32
        torch.manual_seed(42)

        # Create activations with clear harmful/harmless separation
        refusal_dir = torch.randn(hidden)
        refusal_dir = refusal_dir / refusal_dir.norm()

        harmful_acts = [torch.randn(hidden) + 2.0 * refusal_dir for _ in range(32)]
        harmless_acts = [torch.randn(hidden) - 2.0 * refusal_dir for _ in range(32)]
        all_acts = harmful_acts + harmless_acts

        sae = train_sae(all_acts, hidden, expansion=2, n_epochs=30, lr=3e-4, device="cpu")
        result = identify_refusal_features(
            sae, harmful_acts, harmless_acts, layer_idx=0, top_k=4,
        )

        assert result.n_refusal_features == 4
        assert result.sae_directions.shape == (4, hidden)
        assert result.variance_explained > 0.0
        # SAE directions should have some alignment with the actual refusal direction
        best_cos = max(
            abs((result.sae_directions[i] @ refusal_dir).item())
            for i in range(result.sae_directions.shape[0])
        )
        assert best_cos > 0.1, f"SAE should find direction aligned with refusal: best_cos={best_cos}"

    def test_sae_directions_unit_norm(self):
        """SAE-derived directions should be unit normalized."""
        from obliteratus.analysis.sae_abliteration import (
            train_sae, identify_refusal_features,
        )

        hidden = 16
        torch.manual_seed(42)
        harmful = [torch.randn(hidden) + torch.ones(hidden) for _ in range(16)]
        harmless = [torch.randn(hidden) - torch.ones(hidden) for _ in range(16)]

        sae = train_sae(harmful + harmless, hidden, expansion=2, n_epochs=10, device="cpu")
        result = identify_refusal_features(sae, harmful, harmless, 0, top_k=3)

        for i in range(result.sae_directions.shape[0]):
            norm = result.sae_directions[i].norm().item()
            assert abs(norm - 1.0) < 1e-3, f"Direction {i} norm={norm}, expected 1.0"


# ---------------------------------------------------------------------------
# Surgical method preset
# ---------------------------------------------------------------------------

class TestSurgicalMethod:
    def test_surgical_enables_all_sota(self):
        """Surgical method should enable all 6 SOTA techniques."""
        cfg = METHODS["surgical"]
        assert cfg["use_jailbreak_contrast"] is True
        assert cfg["layer_adaptive_strength"] is True
        assert cfg["safety_neuron_masking"] is True
        assert cfg["per_expert_directions"] is True
        assert cfg["attention_head_surgery"] is True
        assert cfg["use_sae_features"] is True

    def test_basic_disables_all_sota(self):
        """Basic method should not enable SOTA techniques (no keys or False)."""
        cfg = METHODS["basic"]
        assert cfg.get("use_jailbreak_contrast", False) is False
        assert cfg.get("layer_adaptive_strength", False) is False
        assert cfg.get("safety_neuron_masking", False) is False

    def test_pipeline_init_surgical(self):
        """Pipeline initialized with surgical method should have all flags set."""
        pipeline = AbliterationPipeline(model_name="test", method="surgical")
        assert pipeline.use_jailbreak_contrast is True
        assert pipeline.layer_adaptive_strength is True
        assert pipeline.safety_neuron_masking is True
        assert pipeline.per_expert_directions is True
        assert pipeline.attention_head_surgery is True
        assert pipeline.use_sae_features is True

    def test_pipeline_init_explicit_override(self):
        """Explicit params should override method defaults."""
        pipeline = AbliterationPipeline(
            model_name="test", method="surgical",
            safety_neuron_masking=False,
        )
        assert pipeline.safety_neuron_masking is False
        assert pipeline.use_jailbreak_contrast is True  # rest still from surgical


# ---------------------------------------------------------------------------
# Inverted method (semantic refusal inversion)
# ---------------------------------------------------------------------------

class TestInvertedMethod:
    def test_inverted_preset_config(self):
        """Inverted method preset should enable inversion flag."""
        cfg = METHODS["inverted"]
        assert cfg["invert_refusal"] is True
        assert cfg["n_directions"] == 8
        assert cfg["use_jailbreak_contrast"] is True

    def test_surgical_does_not_invert(self):
        """Surgical method should NOT enable inversion by default."""
        cfg = METHODS["surgical"]
        assert cfg.get("invert_refusal", False) is False

    def test_pipeline_init_inverted(self):
        """Pipeline initialized with inverted method should have flag set."""
        pipeline = AbliterationPipeline(model_name="test", method="inverted")
        assert pipeline.invert_refusal is True
        assert pipeline.use_jailbreak_contrast is True
        assert pipeline.safety_neuron_masking is False  # zeroing + reflection is destructive

    def test_pipeline_invert_explicit_override(self):
        """Explicit invert_refusal param should override method default."""
        pipeline = AbliterationPipeline(
            model_name="test", method="surgical", invert_refusal=True,
        )
        assert pipeline.invert_refusal is True

        pipeline2 = AbliterationPipeline(
            model_name="test", method="inverted", invert_refusal=False,
        )
        assert pipeline2.invert_refusal is False

    def test_reflection_math(self):
        """2x projection (reflection) should negate the refusal component."""
        hidden = 16

        class Wrapper(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.o_proj = torch.nn.Linear(hidden, 32, bias=False)

        module = Wrapper()
        torch.manual_seed(42)
        W_orig = module.o_proj.weight.data.clone()

        d = torch.randn(hidden, 1)
        d = d / d.norm()

        # Original projection onto d
        orig_proj = (W_orig @ d).squeeze()

        # Reflection: regularization=-1.0 → scale=2.0
        AbliterationPipeline._project_out_advanced(
            module, d, ["o_proj"], regularization=-1.0,
        )

        W_reflected = module.o_proj.weight.data
        new_proj = (W_reflected @ d).squeeze()

        # After reflection, projection should be NEGATED (sign flipped)
        assert torch.allclose(new_proj, -orig_proj, atol=1e-4), (
            f"Reflected projection should be negated: expected ~{-orig_proj[:3]} got {new_proj[:3]}"
        )

    def test_reflection_preserves_orthogonal_component(self):
        """Reflection should not change the component perpendicular to d."""
        hidden = 8

        class Wrapper(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.o_proj = torch.nn.Linear(hidden, 16, bias=False)

        module = Wrapper()
        torch.manual_seed(42)
        W_orig = module.o_proj.weight.data.clone()

        d = torch.randn(hidden, 1)
        d = d / d.norm()

        # Compute original orthogonal component
        orig_d_component = (W_orig @ d) @ d.T  # rank-1 matrix: projection onto d
        orig_ortho = W_orig - orig_d_component  # everything except d-component

        AbliterationPipeline._project_out_advanced(
            module, d, ["o_proj"], regularization=-1.0,
        )

        W_reflected = module.o_proj.weight.data
        new_d_component = (W_reflected @ d) @ d.T
        new_ortho = W_reflected - new_d_component

        # Orthogonal component should be unchanged
        assert torch.allclose(orig_ortho, new_ortho, atol=1e-4), (
            "Reflection should preserve orthogonal component"
        )

    def test_moe_expert_safety_classification(self):
        """_identify_safety_experts should classify experts by router affinity."""
        hidden = 16
        n_experts = 4

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.gate = torch.nn.Linear(hidden, n_experts, bias=False)
                self.experts = torch.nn.ModuleList([
                    torch.nn.Linear(hidden, hidden) for _ in range(n_experts)
                ])

        class FakeLayer(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.self_attn = torch.nn.Module()
                self.self_attn.o_proj = torch.nn.Linear(hidden, hidden, bias=False)
                self.mlp = FakeMoE()

        from obliteratus.models.loader import ModelHandle
        from unittest.mock import MagicMock
        from transformers import GPT2Config

        config = GPT2Config(n_embd=hidden, n_head=2, n_layer=1, vocab_size=100, n_positions=64)
        model = MagicMock()
        model.parameters.return_value = iter([torch.zeros(1)])

        handle = ModelHandle(
            model=model, tokenizer=MagicMock(),
            config=config, model_name="test", task="causal_lm",
        )

        pipeline = AbliterationPipeline(model_name="test", method="inverted")
        pipeline.handle = handle
        pipeline._on_log = lambda m: None
        pipeline._on_stage = lambda r: None

        # Set up fake layer and direction
        layer = FakeLayer()
        torch.manual_seed(42)

        # Make router weight so expert 0 has highest affinity for d
        d = torch.randn(hidden)
        d = d / d.norm()
        # Set router weights: expert 0 aligned with d, expert 3 anti-aligned
        layer.mlp.gate.weight.data[0] = d * 5.0
        layer.mlp.gate.weight.data[1] = d * 1.0
        layer.mlp.gate.weight.data[2] = d * -1.0
        layer.mlp.gate.weight.data[3] = d * -5.0

        # Mock get_layer_modules to return our fake layer
        import obliteratus.abliterate as abl_module
        orig_get_layers = abl_module.get_layer_modules
        orig_get_ffn = abl_module.get_ffn_module
        abl_module.get_layer_modules = lambda h: [layer]
        abl_module.get_ffn_module = lambda lay, a: lay.mlp
        try:
            pipeline.refusal_directions = {0: d}
            pipeline._strong_layers = [0]
            pipeline._identify_safety_experts()
        finally:
            abl_module.get_layer_modules = orig_get_layers
            abl_module.get_ffn_module = orig_get_ffn

        assert 0 in pipeline._expert_safety_scores
        scores = pipeline._expert_safety_scores[0]
        # Expert 0 should be highest safety affinity
        assert scores[0][0] == 0, f"Expert 0 should be top safety, got {scores[0]}"
        # Expert 3 should be lowest
        assert scores[-1][0] == 3, f"Expert 3 should be lowest, got {scores[-1]}"

    def test_moe_inverted_excision_selective(self):
        """Inverted MoE excision should reflect safety experts and remove from capability."""
        hidden = 16
        n_experts = 4

        class FakeExpert(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Linear(hidden, hidden, bias=False)

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.gate = torch.nn.Linear(hidden, n_experts, bias=False)
                self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(n_experts)])

        moe = FakeMoE()
        torch.manual_seed(42)
        for p in moe.parameters():
            p.data = torch.randn_like(p.data)

        d = torch.randn(hidden, 1)
        d = d / d.norm()

        # Set up safety scores: experts 0,1 are safety, 2,3 are capability
        pipeline = AbliterationPipeline(model_name="test", method="inverted")
        pipeline._on_log = lambda m: None
        pipeline._on_stage = lambda r: None
        pipeline._expert_safety_scores = {
            0: [(0, 5.0), (1, 3.0), (2, -1.0), (3, -3.0)]
        }

        orig_router = moe.gate.weight.data.clone()

        count = pipeline._project_moe_experts_inverted(
            moe, d, 0, norm_preserve=False, project_biases=False,
        )

        assert count > 0, "Should project some weights"

        # Router should be reflected (capped at 1.5x to prevent extreme logits
        # that cause CUDA illegal memory access in batched expert forward).
        # With router_reg = max(reflect_reg, -0.5) → scale = 1.5:
        #   new_proj ≈ orig_proj - 1.5 * orig_proj = -0.5 * orig_proj
        # Additionally, _stabilize_router_weights clamps outliers, so we
        # verify the sign is flipped and magnitude is substantial.
        router_proj = (moe.gate.weight.data @ d.squeeze()).squeeze()
        orig_router_proj = (orig_router @ d.squeeze()).squeeze()
        cosine = torch.nn.functional.cosine_similarity(
            router_proj.unsqueeze(0), -orig_router_proj.unsqueeze(0),
        )
        assert cosine > 0.5, (
            f"Router projection should be at least partially reflected, cosine={cosine.item():.3f}"
        )

        # Safety expert 0: should be reflected (projection negated)
        e0_proj = (moe.experts[0].down_proj.weight.data @ d).norm()
        # After reflection the projection doesn't go to zero — it negates
        assert e0_proj > 1e-4, "Safety expert should have non-zero projection (reflected, not removed)"

        # Capability expert 3: should have projection removed (near zero)
        e3_proj = (moe.experts[3].down_proj.weight.data @ d).norm().item()
        assert e3_proj < 1e-3, f"Capability expert should have projection removed, got {e3_proj}"


# ---------------------------------------------------------------------------
# Nuclear method
# ---------------------------------------------------------------------------

class TestNuclearMethod:
    def test_nuclear_preset_config(self):
        """Nuclear method should match inverted baseline + permanent weight techniques."""
        cfg = METHODS["nuclear"]
        assert cfg["invert_refusal"] is True
        assert cfg["n_directions"] == 4  # fewer than inverted to avoid over-ablation
        assert cfg["refinement_passes"] == 2  # same as inverted
        assert cfg["reflection_strength"] == 1.25  # tempered for CoT coherence
        assert cfg["project_embeddings"] is True
        assert cfg["embed_regularization"] == 0.50  # conservative cascade limit
        assert cfg["activation_steering"] is True  # residual cleanup hooks
        assert cfg["steering_strength"] == 0.15  # light residual correction
        assert cfg["expert_transplant"] is True
        assert cfg["transplant_blend"] == 0.10  # gentle nudge, not overwrite
        assert cfg["use_jailbreak_contrast"] is True
        assert cfg["attention_head_surgery"] is True
        assert cfg["layer_adaptive_strength"] is True  # per-layer scaling

    def test_nuclear_pipeline_init(self):
        """Pipeline initialized with nuclear method should have all flags set."""
        pipeline = AbliterationPipeline(model_name="test", method="nuclear")
        assert pipeline.invert_refusal is True
        assert pipeline.reflection_strength == 1.25
        assert pipeline.embed_regularization == 0.50
        assert pipeline.transplant_blend == 0.10
        assert pipeline.project_embeddings is True
        assert pipeline.activation_steering is True  # residual cleanup
        assert pipeline.expert_transplant is True
        assert pipeline.n_directions == 4
        assert pipeline.refinement_passes == 2
        assert pipeline.layer_adaptive_strength is True

    def test_reflection_strength_configurable(self):
        """reflection_strength should be explicitly overridable."""
        pipeline = AbliterationPipeline(
            model_name="test", method="inverted", reflection_strength=3.0,
        )
        assert pipeline.reflection_strength == 3.0

    def test_inverted_default_strength_is_2(self):
        """Inverted method should default to reflection_strength=2.0."""
        pipeline = AbliterationPipeline(model_name="test", method="inverted")
        assert pipeline.reflection_strength == 2.0

    def test_boosted_reflection_math(self):
        """2.5x reflection should produce stronger negation than 2x."""
        hidden = 16

        class Wrapper(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.o_proj = torch.nn.Linear(hidden, 32, bias=False)

        d = torch.randn(hidden, 1)
        d = d / d.norm()

        # 2x reflection
        module_2x = Wrapper()
        torch.manual_seed(42)
        module_2x.o_proj.weight.data = torch.randn(32, hidden)
        orig = module_2x.o_proj.weight.data.clone()
        AbliterationPipeline._project_out_advanced(
            module_2x, d, ["o_proj"], regularization=-1.0,  # scale=2.0
        )
        proj_2x = (module_2x.o_proj.weight.data @ d).squeeze()

        # 2.5x reflection
        module_25x = Wrapper()
        module_25x.o_proj.weight.data = orig.clone()
        AbliterationPipeline._project_out_advanced(
            module_25x, d, ["o_proj"], regularization=-1.5,  # scale=2.5
        )
        proj_25x = (module_25x.o_proj.weight.data @ d).squeeze()

        # 2.5x should be 25% stronger negation than 2x
        assert proj_25x.norm() > proj_2x.norm(), (
            "2.5x reflection should produce stronger (more negative) projection than 2x"
        )

    def test_activation_steering_hook(self):
        """Steering hooks should subtract refusal direction from hidden states."""
        hidden = 8

        class FakeLayer(torch.nn.Module):
            def forward(self, x):
                return x

        layer = FakeLayer()
        layers = torch.nn.ModuleList([layer])

        # Explicitly enable steering (nuclear preset has it off by default)
        pipeline = AbliterationPipeline(
            model_name="test", method="inverted", activation_steering=True,
            steering_strength=0.5,
        )
        pipeline._on_log = lambda m: None
        pipeline._on_stage = lambda r: None

        d = torch.randn(hidden)
        d = d / d.norm()
        pipeline.refusal_directions = {0: d}
        pipeline._strong_layers = [0]

        n_hooks = pipeline._install_activation_steering(layers)
        assert n_hooks == 1
        assert len(pipeline._steering_hooks) == 1

        # Create a hidden state with strong refusal component
        batch = torch.randn(1, 4, hidden)
        refusal_component = 5.0 * d.unsqueeze(0).unsqueeze(0).expand_as(batch)
        input_hidden = batch + refusal_component

        # Run through the layer (hook should fire)
        output = layer(input_hidden)

        # The refusal component should be reduced
        proj_before = torch.einsum("bsh,h->bs", input_hidden, d).abs().mean()
        proj_after = torch.einsum("bsh,h->bs", output, d).abs().mean()
        assert proj_after < proj_before, (
            f"Steering should reduce refusal projection: before={proj_before:.3f}, after={proj_after:.3f}"
        )

        # Cleanup
        for hook in pipeline._steering_hooks:
            hook.remove()

    def test_expert_transplant(self):
        """Expert transplant should overwrite safety expert weights with capability average."""
        hidden = 16
        n_experts = 4

        class FakeExpert(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Linear(hidden, hidden, bias=False)

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.gate = torch.nn.Linear(hidden, n_experts, bias=False)
                self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(n_experts)])

        class FakeLayer(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.self_attn = torch.nn.Module()
                self.self_attn.o_proj = torch.nn.Linear(hidden, hidden, bias=False)
                self.mlp = FakeMoE()

        layer = FakeLayer()
        layers = torch.nn.ModuleList([layer])
        torch.manual_seed(42)
        for p in layer.parameters():
            p.data = torch.randn_like(p.data)

        # Save original safety expert weight
        orig_safety0 = layer.mlp.experts[0].down_proj.weight.data.clone()
        # Save capability expert weights for computing expected mean
        # With top-third classification (n_experts // 3 = 1), only expert 0
        # is safety; experts 1, 2, 3 are all capability.
        cap1 = layer.mlp.experts[1].down_proj.weight.data.clone()
        cap2 = layer.mlp.experts[2].down_proj.weight.data.clone()
        cap3 = layer.mlp.experts[3].down_proj.weight.data.clone()
        expected_mean = (cap1 + cap2 + cap3) / 3.0

        import obliteratus.abliterate as abl_module
        from obliteratus.models.loader import ModelHandle
        from transformers import GPT2Config

        config = GPT2Config(n_embd=hidden, n_head=2, n_layer=1, vocab_size=100, n_positions=64)
        model = MagicMock()
        model.parameters.return_value = iter([torch.zeros(1)])
        handle = ModelHandle(model=model, tokenizer=MagicMock(), config=config, model_name="test", task="causal_lm")

        pipeline = AbliterationPipeline(model_name="test", method="nuclear")
        pipeline.handle = handle
        pipeline._on_log = lambda m: None
        pipeline._on_stage = lambda r: None
        pipeline._strong_layers = [0]
        # Experts 0,1 are safety (high affinity), 2,3 are capability
        pipeline._expert_safety_scores = {
            0: [(0, 5.0), (1, 3.0), (2, -1.0), (3, -3.0)]
        }

        orig_get_ffn = abl_module.get_ffn_module
        abl_module.get_ffn_module = lambda lay, a: lay.mlp
        try:
            count = pipeline._transplant_expert_weights(layers)
        finally:
            abl_module.get_ffn_module = orig_get_ffn

        assert count >= 1, f"Should blend at least 1 weight (top-third safety expert), got {count}"

        # Safety expert 0 should be a 10% blend toward capability mean
        # (nuclear default transplant_blend=0.10)
        # new = 0.90 * original + 0.10 * capability_mean
        blend = pipeline.transplant_blend  # 0.10
        expected_blend = (1.0 - blend) * orig_safety0 + blend * expected_mean
        transplanted = layer.mlp.experts[0].down_proj.weight.data
        assert torch.allclose(transplanted, expected_blend, atol=1e-4), (
            f"Safety expert weight should be {blend:.0%} blended toward capability mean"
        )

        # Capability expert 2 should be unchanged
        assert torch.allclose(layer.mlp.experts[2].down_proj.weight.data, cap2, atol=1e-6), (
            "Capability expert should be unchanged"
        )

    def test_gather_state_dict_raises_on_unmaterialized_meta(self):
        """Should raise RuntimeError when meta tensors survive materialization."""
        from obliteratus.models.loader import ModelHandle
        from transformers import GPT2Config

        config = GPT2Config(n_embd=8, n_head=2, n_layer=1, vocab_size=100, n_positions=64)

        # Create a fake model whose state_dict returns a meta tensor.
        # Patch get_state_dict_offloaded_model to return the same meta tensor
        # (simulating a failed materialization).
        fake_model = MagicMock()
        meta_tensor = torch.empty(4, 8, device="meta")
        fake_model.state_dict.return_value = {"layer.weight": meta_tensor}

        handle = ModelHandle(
            model=fake_model, tokenizer=MagicMock(), config=config,
            model_name="test", task="causal_lm",
        )

        pipeline = AbliterationPipeline(model_name="test", method="nuclear")
        pipeline.handle = handle
        pipeline._on_log = lambda m: None
        pipeline._on_stage = lambda r: None

        # Patch accelerate's function to return meta tensors (simulates
        # materialization failure — the safety net should catch this).
        with patch(
            "accelerate.utils.modeling.get_state_dict_offloaded_model",
            return_value={"layer.weight": meta_tensor},
        ):
            with pytest.raises(RuntimeError, match="bricked checkpoint"):
                pipeline._gather_state_dict()

    def test_gather_state_dict_materializes_offloaded_meta(self):
        """Should return real CPU tensors from Accelerate materialization."""
        from obliteratus.models.loader import ModelHandle
        from transformers import GPT2Config

        config = GPT2Config(n_embd=8, n_head=2, n_layer=1, vocab_size=100, n_positions=64)
        fake_model = MagicMock()
        fake_model.state_dict.return_value = {
            "layer.weight": torch.empty(4, 8, device="meta"),
        }
        fake_model.parameters.return_value = []
        fake_model.buffers.return_value = []
        materialized = torch.zeros(4, 8)

        handle = ModelHandle(
            model=fake_model, tokenizer=MagicMock(), config=config,
            model_name="test", task="causal_lm",
        )
        pipeline = AbliterationPipeline(model_name="test", method="nuclear")
        pipeline.handle = handle
        pipeline._on_log = lambda m: None
        pipeline._on_stage = lambda r: None

        with patch(
            "accelerate.utils.modeling.get_state_dict_offloaded_model",
            return_value={"layer.weight": materialized},
        ) as gather_offloaded:
            state_dict = pipeline._gather_state_dict()

        gather_offloaded.assert_called_once_with(fake_model)
        assert state_dict["layer.weight"] is materialized
        assert state_dict["layer.weight"].device.type == "cpu"


# ---------------------------------------------------------------------------
# Knee detection
# ---------------------------------------------------------------------------

class TestKneeDetection:
    def test_empty_input(self):
        result = AbliterationPipeline._select_layers_knee([])
        assert result == []

    def test_two_layers(self):
        result = AbliterationPipeline._select_layers_knee([(0, 5.0), (1, 3.0)])
        assert set(result) == {0, 1}

    def test_clear_knee(self):
        """Layers with a sharp dropoff should be separated by knee detection."""
        sorted_layers = [
            (14, 10.0), (15, 9.5), (13, 9.0),  # strong cluster
            (16, 2.0), (12, 1.5), (17, 1.0), (11, 0.5), (18, 0.2), (10, 0.1),
        ]
        result = AbliterationPipeline._select_layers_knee(sorted_layers)
        # Should select the strong cluster (layers 14, 15, 13) and exclude weak ones
        assert 14 in result
        assert 15 in result
        assert 13 in result
        assert len(result) <= 5  # shouldn't select all 9

    def test_minimum_threshold_filters_noise(self):
        """Layers below 10% of max should be filtered out."""
        sorted_layers = [(0, 10.0), (1, 0.5)]  # 0.5 is 5% of 10
        result = AbliterationPipeline._select_layers_knee(sorted_layers)
        # Layer 1 is below 10% threshold
        assert 0 in result

    def test_all_equal_norms(self):
        """When all norms are equal, should select all (or most)."""
        sorted_layers = [(i, 5.0) for i in range(5)]
        result = AbliterationPipeline._select_layers_knee(sorted_layers)
        assert len(result) >= 1


# ---------------------------------------------------------------------------
# Activation collection
# ---------------------------------------------------------------------------

class TestActivationCollection:
    def test_collect_activations(self, handle):
        """Test that activation collection returns correct structure."""
        from obliteratus.strategies.utils import get_layer_modules

        pipeline = AbliterationPipeline(model_name="test")
        pipeline.handle = handle
        pipeline._on_log = lambda m: None

        layers = get_layer_modules(handle)
        prompts = ["Hello world", "Test prompt"]

        _make_varied_tokenizer(handle)

        activations = pipeline._collect_activations(layers, prompts, "test")

        assert len(activations) == len(layers)
        for idx in range(len(layers)):
            assert len(activations[idx]) == len(prompts)
            for act in activations[idx]:
                assert act.device == torch.device("cpu")
                assert act.shape[-1] == handle.hidden_size


# ---------------------------------------------------------------------------
# Distill: single direction (basic method)
# ---------------------------------------------------------------------------

class TestDistillBasic:
    def test_single_direction(self, handle):
        """Basic method: single refusal direction via difference-in-means."""
        from obliteratus.strategies.utils import get_layer_modules

        pipeline = AbliterationPipeline(
            model_name="test",
            method="basic",
            harmful_prompts=["bad prompt"],
            harmless_prompts=["good prompt"],
        )
        pipeline.handle = handle
        pipeline._on_log = lambda m: None
        pipeline._on_stage = lambda r: None
        _make_varied_tokenizer(handle)

        pipeline._probe()
        pipeline._distill()

        n_layers = len(get_layer_modules(handle))
        assert len(pipeline.refusal_directions) == n_layers
        for idx, direction in pipeline.refusal_directions.items():
            assert abs(direction.norm().item() - 1.0) < 1e-4
            # Single direction: subspace should be (1, hidden_dim)
            assert pipeline.refusal_subspaces[idx].shape[0] == 1


# ---------------------------------------------------------------------------
# Distill: multi-direction SVD (advanced/aggressive method)
# ---------------------------------------------------------------------------

class TestDistillSVD:
    def test_multi_direction_svd(self, handle):
        """Advanced method: SVD extracts multiple refusal directions.

        Note: on small models (hidden_size < 2048 or < 2B params), n_directions
        is automatically capped to 2 to prevent over-ablation.  The test model
        (hidden_size=64, 4 layers) triggers this safeguard.
        """
        from obliteratus.strategies.utils import get_layer_modules

        pipeline = AbliterationPipeline(
            model_name="test",
            method="advanced",
            harmful_prompts=["bad1", "bad2", "bad3", "bad4", "bad5"],
            harmless_prompts=["good1", "good2", "good3", "good4", "good5"],
        )
        pipeline.handle = handle
        pipeline._on_log = lambda m: None
        pipeline._on_stage = lambda r: None
        _make_varied_tokenizer(handle)

        pipeline._probe()
        pipeline._distill()

        n_layers = len(get_layer_modules(handle))
        assert len(pipeline.refusal_subspaces) == n_layers
        # Small-model cap: n_directions capped to 2 for tiny test model
        expected_dirs = min(2, pipeline.n_directions, 5, handle.hidden_size)
        for idx, subspace in pipeline.refusal_subspaces.items():
            assert subspace.shape[0] == expected_dirs
            assert subspace.shape[1] == handle.hidden_size

        # Primary direction should still be a unit vector
        for idx, direction in pipeline.refusal_directions.items():
            assert abs(direction.norm().item() - 1.0) < 1e-4


# ---------------------------------------------------------------------------
# Full pipeline: excise with different methods
# ---------------------------------------------------------------------------

class TestExcise:
    def test_qwen25_coder_advanced_excise_projects_all_dense_weights(self):
        """Qwen2.5 must not silently resolve to zero projection targets."""
        config = Qwen2Config(
            vocab_size=128,
            hidden_size=32,
            intermediate_size=64,
            num_hidden_layers=2,
            num_attention_heads=4,
            num_key_value_heads=2,
            max_position_embeddings=64,
        )
        model = Qwen2ForCausalLM(config)
        handle = ModelHandle(
            model=model,
            tokenizer=MagicMock(),
            config=config,
            model_name="Qwen/Qwen2.5-Coder-7B-Instruct",
            task="causal_lm",
        )
        pipeline = AbliterationPipeline(
            model_name=handle.model_name,
            method="advanced",
            refinement_passes=1,
            norm_preserve=False,
            project_biases=False,
            layer_adaptive_strength=False,
            safety_neuron_masking=False,
            attention_head_surgery=False,
            use_sae_features=False,
            use_kl_optimization=False,
            projection_target="all",
            harmful_prompts=["harmful"],
            harmless_prompts=["harmless"],
        )
        pipeline.handle = handle
        pipeline._strong_layers = [0]
        direction = torch.arange(1, config.hidden_size + 1, dtype=torch.float32)
        direction /= direction.norm()
        pipeline.refusal_directions = {0: direction.clone()}
        pipeline.refusal_subspaces = {0: direction.unsqueeze(0)}

        layer = model.model.layers[0]
        original_weights = {
            name: parameter.detach().clone()
            for name, parameter in layer.named_parameters()
        }

        pipeline._excise()

        expected = {
            "self_attn.q_proj.weight",
            "self_attn.k_proj.weight",
            "self_attn.v_proj.weight",
            "self_attn.o_proj.weight",
            "mlp.gate_proj.weight",
            "mlp.up_proj.weight",
            "mlp.down_proj.weight",
        }
        changed = {
            name
            for name, parameter in layer.named_parameters()
            if not torch.equal(original_weights[name], parameter)
        }
        assert handle.architecture == "qwen2"
        assert changed == expected
        assert pipeline._excise_modified_count == 8  # seven layer weights + lm_head

    def test_excise_basic(self, handle):
        """Basic method should modify weights."""
        from obliteratus.strategies.utils import get_layer_modules

        pipeline = AbliterationPipeline(
            model_name="test",
            method="basic",
            harmful_prompts=["bad prompt"],
            harmless_prompts=["good prompt"],
        )
        pipeline.handle = handle
        pipeline._on_log = lambda m: None
        pipeline._on_stage = lambda r: None
        _make_varied_tokenizer(handle)

        layers = get_layer_modules(handle)
        original_weights = {}
        for idx in range(len(layers)):
            for name, param in layers[idx].named_parameters():
                original_weights[(idx, name)] = param.data.clone()

        pipeline._probe()
        pipeline._distill()
        pipeline._excise()

        any_changed = False
        for idx in range(len(layers)):
            for name, param in layers[idx].named_parameters():
                if not torch.allclose(original_weights[(idx, name)], param.data, atol=1e-6):
                    any_changed = True
                    break

        assert any_changed, "Excise should modify at least some weights"

    def test_excise_advanced_norm_preserving(self, handle):
        """Advanced method with norm preservation should maintain weight norms."""
        from obliteratus.strategies.utils import get_layer_modules

        pipeline = AbliterationPipeline(
            model_name="test",
            method="advanced",
            harmful_prompts=["bad prompt"],
            harmless_prompts=["good prompt"],
        )
        pipeline.handle = handle
        pipeline._on_log = lambda m: None
        pipeline._on_stage = lambda r: None
        _make_varied_tokenizer(handle)

        get_layer_modules(handle)

        pipeline._probe()
        pipeline._distill()
        pipeline._excise()

        # Weights should have been modified (advanced uses _project_out_advanced)
        assert len(pipeline._strong_layers) > 0


# ---------------------------------------------------------------------------
# Rebirth (save)
# ---------------------------------------------------------------------------

class TestRebirth:
    def test_rebirth_cancellation_discards_staging_before_promotion(self, handle, tmp_path):
        from threading import Event

        from obliteratus.abliterate import PipelineCancelledError

        cancellation = Event()
        output = tmp_path / "output"
        pipeline = AbliterationPipeline(
            model_name="test-model",
            output_dir=str(output),
            cancellation_event=cancellation,
        )
        pipeline.handle = handle
        pipeline._strong_layers = [0]

        def save_then_cancel(path, **_kwargs):
            (Path(path) / "config.json").write_text("{}", encoding="utf-8")
            (Path(path) / "model.safetensors").write_bytes(b"partial")
            cancellation.set()

        handle.model.save_pretrained = MagicMock(side_effect=save_then_cancel)
        handle.tokenizer.save_pretrained = MagicMock()

        with pytest.raises(PipelineCancelledError, match="cancelled during REBIRTH"):
            pipeline._rebirth()

        assert not output.exists()
        assert list(tmp_path.glob(".output.staging-*")) == []
        handle.tokenizer.save_pretrained.assert_not_called()

    def test_rebirth_saves_metadata(self, handle, tmp_path):
        """Rebirth should save model and comprehensive metadata JSON."""
        pipeline = AbliterationPipeline(
            model_name="test-model",
            output_dir=str(tmp_path / "output"),
            method="advanced",
        )
        pipeline.handle = handle
        pipeline._on_log = lambda m: None
        pipeline._on_stage = lambda r: None
        pipeline._strong_layers = [0]
        pipeline._quality_metrics = {"perplexity": 8.5, "coherence": 1.0}

        handle.model.save_pretrained = MagicMock(
            side_effect=lambda path, **_kwargs: (
                (Path(path) / "config.json").write_text("{}", encoding="utf-8"),
                (Path(path) / "model.safetensors").write_bytes(b"weights"),
            ),
        )
        handle.tokenizer.save_pretrained = MagicMock(
            side_effect=lambda path: (Path(path) / "tokenizer_config.json").write_text(
                "{}",
                encoding="utf-8",
            ),
        )

        result_path = pipeline._rebirth()

        assert result_path == tmp_path / "output"
        assert (result_path / "abliteration_metadata.json").exists()

        metadata = json.loads((result_path / "abliteration_metadata.json").read_text())
        assert metadata["source_model"] == "test-model"
        assert metadata["technique"] == "refusal_direction_ablation"
        assert metadata["method"] == "advanced"
        assert metadata["strong_layers"] == [0]
        assert "method_config" in metadata
        assert metadata["method_config"]["n_directions"] == METHODS["advanced"]["n_directions"]
        assert metadata["method_config"]["norm_preserve"] is True
        assert "references" in metadata
        assert len(metadata["references"]) >= 3
        assert "quality_metrics" in metadata
        assert metadata["quality_metrics"]["perplexity"] == 8.5


# ---------------------------------------------------------------------------
# CLI integration
# ---------------------------------------------------------------------------

class TestCLI:
    def test_abliterate_parser_with_method(self):
        """Test that the abliterate subcommand parses method correctly."""
        import argparse

        parser = argparse.ArgumentParser()
        subparsers = parser.add_subparsers(dest="command")
        abl_parser = subparsers.add_parser("abliterate")
        abl_parser.add_argument("model", type=str)
        abl_parser.add_argument("--output-dir", type=str, default=None)
        abl_parser.add_argument("--device", type=str, default="auto")
        abl_parser.add_argument("--dtype", type=str, default="float16")
        abl_parser.add_argument("--method", type=str, default="advanced",
                                choices=["basic", "advanced", "aggressive"])
        abl_parser.add_argument("--n-directions", type=int, default=None)
        abl_parser.add_argument("--regularization", type=float, default=None)
        abl_parser.add_argument("--refinement-passes", type=int, default=None)

        args = parser.parse_args(["abliterate", "gpt2", "--method", "aggressive", "--n-directions", "6"])
        assert args.command == "abliterate"
        assert args.model == "gpt2"
        assert args.method == "aggressive"
        assert args.n_directions == 6
        assert args.dtype == "float16"

    def test_default_method(self):
        """Default method should be advanced."""
        import argparse

        parser = argparse.ArgumentParser()
        subparsers = parser.add_subparsers(dest="command")
        abl_parser = subparsers.add_parser("abliterate")
        abl_parser.add_argument("model", type=str)
        abl_parser.add_argument("--method", type=str, default="advanced")

        args = parser.parse_args(["abliterate", "gpt2"])
        assert args.method == "advanced"


# ---------------------------------------------------------------------------
# Expert-Granular Abliteration (EGA)
# ---------------------------------------------------------------------------

class TestFindRouterModule:
    """Test _find_router_module static method."""

    def test_finds_gate(self):
        """Should find a router named 'gate'."""
        hidden = 16

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.gate = torch.nn.Linear(hidden, 4, bias=False)
                self.experts = torch.nn.ModuleList()

        moe = FakeMoE()
        router = AbliterationPipeline._find_router_module(moe)
        assert router is moe.gate

    def test_finds_router(self):
        """Should find a router named 'router'."""
        hidden = 16

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.router = torch.nn.Linear(hidden, 4, bias=False)
                self.experts = torch.nn.ModuleList()

        moe = FakeMoE()
        router = AbliterationPipeline._find_router_module(moe)
        assert router is moe.router

    def test_auto_detects_unknown_router(self):
        """Should auto-detect a router with unusual name via heuristic."""
        hidden = 16

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.moe_gate_proj = torch.nn.Linear(hidden, 4, bias=False)
                self.experts = torch.nn.ModuleList()

        moe = FakeMoE()
        router = AbliterationPipeline._find_router_module(moe)
        assert router is moe.moe_gate_proj

    def test_returns_none_no_router(self):
        """Should return None when no router is found."""
        class NoRouter(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = torch.nn.Linear(16, 16)

        mod = NoRouter()
        assert AbliterationPipeline._find_router_module(mod) is None


class TestRouterProfilingHooks:
    """Test _install_router_profiling_hooks."""

    def _make_moe_pipeline_and_layers(self, hidden=16, n_experts=4):
        """Create a pipeline with a fake MoE model for router profiling tests."""
        from obliteratus.models.loader import ModelHandle
        from transformers import GPT2Config

        class FakeExpert(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Linear(hidden, hidden, bias=False)

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.gate = torch.nn.Linear(hidden, n_experts, bias=False)
                self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(n_experts)])

            def forward(self, x):
                return x

        class FakeLayer(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.self_attn = torch.nn.Module()
                self.self_attn.o_proj = torch.nn.Linear(hidden, hidden, bias=False)
                self.mlp = FakeMoE()

            def forward(self, x):
                return (x,)

        config = GPT2Config(n_embd=hidden, n_head=2, n_layer=1, vocab_size=100, n_positions=64)
        model = MagicMock()
        model.parameters.return_value = iter([torch.zeros(1)])
        handle = ModelHandle(model=model, tokenizer=MagicMock(), config=config, model_name="test", task="causal_lm")

        pipeline = AbliterationPipeline(model_name="test", method="surgical")
        pipeline.handle = handle
        pipeline._on_log = lambda m: None
        pipeline._on_stage = lambda r: None

        layer = FakeLayer()
        layers = torch.nn.ModuleList([layer])

        # Monkey-patch get_ffn_module
        import obliteratus.abliterate as abl_module
        orig_get_ffn = abl_module.get_ffn_module
        abl_module.get_ffn_module = lambda lay, a: lay.mlp

        return pipeline, layers, layer, abl_module, orig_get_ffn

    def test_hooks_installed(self):
        """Should install hooks on MoE router modules."""
        pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
        try:
            hooks = pipeline._install_router_profiling_hooks(layers)
            assert len(hooks) == 1
            assert 0 in pipeline._routing_harmful
            assert 0 in pipeline._routing_harmless
        finally:
            for h in hooks:
                h.remove()
            abl_module.get_ffn_module = orig_get_ffn

    def test_hooks_record_logits(self):
        """Hooks should record router logits during forward passes."""
        pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
        try:
            hooks = pipeline._install_router_profiling_hooks(layers)

            # Simulate harmful forward pass
            pipeline._routing_is_harmful = True
            x = torch.randn(1, 5, 16)
            harmful_expected = torch.nn.functional.linear(x, layer.mlp.gate.weight)[0, -1]
            layer.mlp.gate(x)  # triggers hook

            assert len(pipeline._routing_harmful[0]) == 1
            assert pipeline._routing_harmful[0][0].shape[0] == 4  # n_experts
            assert torch.equal(pipeline._routing_harmful[0][0], harmful_expected)

            # Simulate harmless forward pass
            pipeline._routing_is_harmful = False
            harmless_x = x + 1
            harmless_expected = torch.nn.functional.linear(
                harmless_x, layer.mlp.gate.weight,
            )[0, -1]
            layer.mlp.gate(harmless_x)

            assert len(pipeline._routing_harmless[0]) == 1
            assert torch.equal(pipeline._routing_harmless[0][0], harmless_expected)
            assert not torch.equal(
                pipeline._routing_harmful[0][0], pipeline._routing_harmless[0][0],
            )
        finally:
            for h in hooks:
                h.remove()
            abl_module.get_ffn_module = orig_get_ffn

    def test_hooks_record_single_prompt_vector_output(self):
        """A one-dimensional single-prompt router output should remain intact."""
        pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
        hooks = []
        try:
            hooks = pipeline._install_router_profiling_hooks(layers)
            x = torch.arange(16, dtype=torch.float32) / 100
            expected = torch.nn.functional.linear(x, layer.mlp.gate.weight)

            layer.mlp.gate(x)

            assert len(pipeline._routing_harmful[0]) == 1
            assert torch.equal(pipeline._routing_harmful[0][0], expected)
        finally:
            for h in hooks:
                h.remove()
            abl_module.get_ffn_module = orig_get_ffn

    def test_hooks_record_one_vector_for_each_3d_batch_prompt(self):
        """Native 3-D router outputs should retain every batch member."""
        pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
        hooks = []
        try:
            hooks = pipeline._install_router_profiling_hooks(layers)
            x = torch.arange(2 * 5 * 16, dtype=torch.float32).reshape(2, 5, 16) / 100
            expected = torch.nn.functional.linear(x, layer.mlp.gate.weight)

            layer.mlp.gate(x)

            recorded = pipeline._routing_harmful[0]
            assert len(recorded) == 2
            assert torch.equal(recorded[0], expected[0, -1])
            assert torch.equal(recorded[1], expected[1, -1])
        finally:
            for h in hooks:
                h.remove()
            abl_module.get_ffn_module = orig_get_ffn

    def test_hooks_preserve_flattened_prompt_association_with_padding(self):
        """Flattened token rows should become one vector per padded prompt."""
        pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
        hooks = []
        try:
            hooks = pipeline._install_router_profiling_hooks(layers)
            pipeline._routing_attention_mask = torch.tensor(
                [[1, 1, 0, 0], [1, 1, 1, 1]], dtype=torch.long,
            )
            x = torch.arange(8 * 16, dtype=torch.float32).reshape(8, 16) / 100
            expected = torch.nn.functional.linear(x, layer.mlp.gate.weight).reshape(2, 4, 4)

            layer.mlp.gate(x)

            recorded = pipeline._routing_harmful[0]
            assert len(recorded) == 2
            assert torch.equal(recorded[0], expected[0, 1])
            assert torch.equal(recorded[1], expected[1, 3])
        finally:
            pipeline._routing_attention_mask = None
            for h in hooks:
                h.remove()
            abl_module.get_ffn_module = orig_get_ffn

    def test_hooks_record_single_prompt_flattened_tokens(self):
        """A flattened single-prompt sequence should select its final token."""
        pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
        hooks = []
        try:
            hooks = pipeline._install_router_profiling_hooks(layers)
            pipeline._routing_attention_mask = torch.ones(1, 4, dtype=torch.long)
            x = torch.arange(4 * 16, dtype=torch.float32).reshape(4, 16) / 100
            expected = torch.nn.functional.linear(x, layer.mlp.gate.weight)[-1]

            layer.mlp.gate(x)

            assert len(pipeline._routing_harmful[0]) == 1
            assert torch.equal(pipeline._routing_harmful[0][0], expected)
        finally:
            pipeline._routing_attention_mask = None
            for h in hooks:
                h.remove()
            abl_module.get_ffn_module = orig_get_ffn

    def test_hooks_average_only_valid_cot_tokens_per_prompt(self):
        """CoT aggregation should exclude padding without mixing prompts."""
        pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
        hooks = []
        try:
            hooks = pipeline._install_router_profiling_hooks(layers)
            pipeline.cot_aware = True
            pipeline._routing_attention_mask = torch.tensor(
                [[1, 1, 1, 1, 1, 0], [0, 1, 1, 1, 1, 1]], dtype=torch.long,
            )
            x = torch.arange(12 * 16, dtype=torch.float32).reshape(12, 16) / 100
            expected = torch.nn.functional.linear(x, layer.mlp.gate.weight).reshape(2, 6, 4)

            layer.mlp.gate(x)

            recorded = pipeline._routing_harmful[0]
            assert len(recorded) == 2
            assert torch.allclose(recorded[0], expected[0, :5].mean(dim=0))
            assert torch.allclose(recorded[1], expected[1, 1:].mean(dim=0))
        finally:
            pipeline._routing_attention_mask = None
            for h in hooks:
                h.remove()
            abl_module.get_ffn_module = orig_get_ffn

    def test_hooks_skip_ambiguous_flattened_logits(self):
        """Ambiguous flattened layouts should fail visibly instead of collapsing."""
        pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
        hooks = []
        try:
            hooks = pipeline._install_router_profiling_hooks(layers)

            with pytest.warns(RuntimeWarning, match="multi-row 2-D router logits are ambiguous"):
                layer.mlp.gate(torch.randn(5, 16))

            assert pipeline._routing_harmful[0] == []
        finally:
            for h in hooks:
                h.remove()
            abl_module.get_ffn_module = orig_get_ffn

    def test_hooks_skip_mismatched_flattened_logits(self):
        """A stale or incompatible mask should produce an actionable warning."""
        pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
        hooks = []
        try:
            hooks = pipeline._install_router_profiling_hooks(layers)
            pipeline._routing_attention_mask = torch.ones(2, 4, dtype=torch.long)

            with pytest.warns(RuntimeWarning, match="do not match batch × sequence"):
                layer.mlp.gate(torch.randn(7, 16))

            assert pipeline._routing_harmful[0] == []
        finally:
            pipeline._routing_attention_mask = None
            for h in hooks:
                h.remove()
            abl_module.get_ffn_module = orig_get_ffn

    def test_no_handle_returns_empty(self):
        """Should return empty list when handle is None."""
        pipeline = AbliterationPipeline(model_name="test", method="surgical")
        pipeline.handle = None
        hooks = pipeline._install_router_profiling_hooks(torch.nn.ModuleList())
        assert hooks == []


class TestComputeExpertGranularDirections:
    """Test _compute_expert_granular_directions."""

    def test_computes_per_expert_directions(self):
        """Should compute per-expert refusal directions from routing data."""
        hidden = 16
        n_experts = 4

        pipeline = AbliterationPipeline(model_name="test", method="surgical")
        pipeline._on_log = lambda m: None
        pipeline._on_stage = lambda r: None
        pipeline._strong_layers = [0]

        torch.manual_seed(42)

        # Simulate router logits: expert 0 favored for harmful, expert 3 for harmless
        h_logits = []
        s_logits = []
        for _ in range(10):
            hl = torch.randn(n_experts)
            hl[0] += 2.0  # bias expert 0 for harmful
            h_logits.append(hl)
            sl = torch.randn(n_experts)
            sl[3] += 2.0  # bias expert 3 for harmless
            s_logits.append(sl)

        pipeline._routing_harmful = {0: h_logits}
        pipeline._routing_harmless = {0: s_logits}

        # Simulate per-prompt activations with harmful/harmless separation
        refusal_dir = torch.randn(hidden)
        refusal_dir = refusal_dir / refusal_dir.norm()

        h_acts = [torch.randn(hidden) + 1.5 * refusal_dir for _ in range(10)]
        s_acts = [torch.randn(hidden) - 1.5 * refusal_dir for _ in range(10)]
        pipeline._harmful_acts = {0: h_acts}
        pipeline._harmless_acts = {0: s_acts}

        pipeline._compute_expert_granular_directions()

        # Should have computed expert directions for layer 0
        assert 0 in pipeline._expert_directions
        assert len(pipeline._expert_directions[0]) > 0

        # Should have dynamic safety scores
        assert 0 in pipeline._expert_safety_scores
        scores = pipeline._expert_safety_scores[0]
        assert len(scores) == n_experts
        # Expert 0 should have higher safety score (more activated for harmful)
        expert_0_score = next(s for eid, s in scores if eid == 0)
        expert_3_score = next(s for eid, s in scores if eid == 3)
        assert expert_0_score > expert_3_score, (
            f"Expert 0 should have higher safety score: {expert_0_score} vs {expert_3_score}"
        )

    def test_directions_are_unit_vectors(self):
        """Per-expert directions should be unit normalized."""
        hidden = 16
        n_experts = 4

        pipeline = AbliterationPipeline(model_name="test", method="surgical")
        pipeline._on_log = lambda m: None
        pipeline._strong_layers = [0]

        torch.manual_seed(42)
        h_logits = [torch.randn(n_experts) for _ in range(10)]
        s_logits = [torch.randn(n_experts) for _ in range(10)]
        pipeline._routing_harmful = {0: h_logits}
        pipeline._routing_harmless = {0: s_logits}
        pipeline._harmful_acts = {0: [torch.randn(hidden) + torch.ones(hidden) for _ in range(10)]}
        pipeline._harmless_acts = {0: [torch.randn(hidden) - torch.ones(hidden) for _ in range(10)]}

        pipeline._compute_expert_granular_directions()

        if 0 in pipeline._expert_directions:
            for ei, d in pipeline._expert_directions[0].items():
                assert abs(d.norm().item() - 1.0) < 1e-4, (
                    f"Expert {ei} direction norm={d.norm().item()}, expected 1.0"
                )

    def test_skips_when_no_routing_data(self):
        """Should skip gracefully when no routing data is available."""
        pipeline = AbliterationPipeline(model_name="test", method="surgical")
        pipeline._on_log = lambda m: None
        pipeline._routing_harmful = {}
        pipeline._routing_harmless = {}

        pipeline._compute_expert_granular_directions()

        assert len(pipeline._expert_directions) == 0

    def test_skips_expert_with_low_routing_weight(self):
        """Experts with insufficient routing weight should not get directions."""
        hidden = 16

        pipeline = AbliterationPipeline(model_name="test", method="surgical")
        pipeline._on_log = lambda m: None
        pipeline._strong_layers = [0]

        # Create routing logits where expert 3 is never selected (very low)
        h_logits = []
        s_logits = []
        for _ in range(3):
            hl = torch.tensor([5.0, 5.0, 5.0, -100.0])  # expert 3 never routed
            h_logits.append(hl)
            sl = torch.tensor([5.0, 5.0, 5.0, -100.0])
            s_logits.append(sl)

        pipeline._routing_harmful = {0: h_logits}
        pipeline._routing_harmless = {0: s_logits}

        torch.manual_seed(42)
        pipeline._harmful_acts = {0: [torch.randn(hidden) for _ in range(3)]}
        pipeline._harmless_acts = {0: [torch.randn(hidden) for _ in range(3)]}

        pipeline._compute_expert_granular_directions()

        # Expert 3 should NOT have a direction (routing weight too low)
        if 0 in pipeline._expert_directions:
            assert 3 not in pipeline._expert_directions[0], (
                "Expert with near-zero routing weight should not get a direction"
            )


class TestProjectMoEExpertsGranular:
    """Test _project_moe_experts_granular (ModuleList path)."""

    def _make_direction(self, hidden_dim=16):
        d = torch.randn(hidden_dim, 1)
        return d / d.norm()

    def test_per_expert_directions_applied(self):
        """Each expert should use its own direction when available."""
        hidden = 16
        n_experts = 4

        class FakeExpert(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Linear(hidden, 32, bias=False)
                self.up_proj = torch.nn.Linear(hidden, 32, bias=False)

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.gate = torch.nn.Linear(hidden, n_experts, bias=False)
                self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(n_experts)])

        moe = FakeMoE()
        torch.manual_seed(42)
        for p in moe.parameters():
            p.data = torch.randn_like(p.data)

        shared_dir = self._make_direction(hidden)

        # Create distinct per-expert directions
        expert_dirs = {}
        for ei in range(n_experts):
            d = torch.randn(hidden)
            d = d / d.norm()
            expert_dirs[ei] = d

        pipeline = AbliterationPipeline(model_name="test", method="surgical")
        pipeline._on_log = lambda m: None
        pipeline._expert_directions = {0: expert_dirs}

        # Save originals
        orig_weights = {
            ei: moe.experts[ei].down_proj.weight.data.clone()
            for ei in range(n_experts)
        }

        count = pipeline._project_moe_experts_granular(
            moe, shared_dir, layer_idx=0,
        )

        assert count > 0, "Should project some weights"

        # All experts should be modified
        for ei in range(n_experts):
            assert not torch.allclose(
                moe.experts[ei].down_proj.weight.data, orig_weights[ei]
            ), f"Expert {ei} should be modified"

    def test_falls_back_to_shared_direction(self):
        """Experts without per-expert direction should use shared direction."""
        hidden = 16
        n_experts = 4

        class FakeExpert(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Linear(hidden, 32, bias=False)
                self.up_proj = torch.nn.Linear(hidden, 32, bias=False)

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.gate = torch.nn.Linear(hidden, n_experts, bias=False)
                self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(n_experts)])

        moe = FakeMoE()
        torch.manual_seed(42)
        for p in moe.parameters():
            p.data = torch.randn_like(p.data)

        shared_dir = self._make_direction(hidden)

        # Only expert 0 has a per-expert direction
        expert_dirs = {0: torch.randn(hidden).div_(torch.randn(hidden).norm())}
        expert_dirs[0] = expert_dirs[0] / expert_dirs[0].norm()

        pipeline = AbliterationPipeline(model_name="test", method="surgical")
        pipeline._on_log = lambda m: None
        pipeline._expert_directions = {0: expert_dirs}

        orig_e1 = moe.experts[1].down_proj.weight.data.clone()

        pipeline._project_moe_experts_granular(
            moe, shared_dir, layer_idx=0,
        )

        # Experts 1,2,3 should be modified (using shared direction)
        assert not torch.allclose(moe.experts[1].down_proj.weight.data, orig_e1), \
            "Expert 1 should use shared direction fallback"

    def test_router_uses_shared_direction(self):
        """Router should always use the shared direction, not per-expert."""
        hidden = 16
        n_experts = 4

        class FakeExpert(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Linear(hidden, 32, bias=False)

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.gate = torch.nn.Linear(hidden, n_experts, bias=False)
                self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(n_experts)])

        moe = FakeMoE()
        shared_dir = self._make_direction(hidden)

        pipeline = AbliterationPipeline(model_name="test", method="surgical")
        pipeline._on_log = lambda m: None
        pipeline._expert_directions = {0: {0: torch.randn(hidden)}}

        orig_gate = moe.gate.weight.data.clone()

        pipeline._project_moe_experts_granular(moe, shared_dir, layer_idx=0)

        # Gate should be projected
        assert not torch.allclose(moe.gate.weight.data, orig_gate), \
            "Router should be projected with shared direction"

        # Gate's projection onto shared direction should be near zero
        proj = (moe.gate.weight.data @ shared_dir).norm().item()
        assert proj < 1e-4, f"Router should have shared dir removed, proj={proj}"

    def test_shared_expert_uses_shared_direction(self):
        """Shared expert should always use the shared direction."""
        hidden = 16

        class FakeExpert(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Linear(hidden, 32, bias=False)
                self.up_proj = torch.nn.Linear(hidden, 32, bias=False)

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.gate = torch.nn.Linear(hidden, 2, bias=False)
                self.shared_expert = torch.nn.Module()
                self.shared_expert.down_proj = torch.nn.Linear(hidden, 32, bias=False)
                self.shared_expert.up_proj = torch.nn.Linear(hidden, 32, bias=False)
                self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(2)])

        moe = FakeMoE()
        shared_dir = self._make_direction(hidden)

        pipeline = AbliterationPipeline(model_name="test", method="surgical")
        pipeline._on_log = lambda m: None
        pipeline._expert_directions = {0: {0: torch.randn(hidden)}}

        orig_shared = moe.shared_expert.down_proj.weight.data.clone()

        pipeline._project_moe_experts_granular(moe, shared_dir, layer_idx=0)

        assert not torch.allclose(moe.shared_expert.down_proj.weight.data, orig_shared), \
            "Shared expert should be projected"


class TestProjectFused3DGranular:
    """Test _project_fused_3d_granular for fused 3D expert tensors."""

    def test_per_expert_directions_on_fused(self):
        """Each expert slice should use its own direction."""
        hidden = 16
        intermediate = 32
        n_experts = 4

        class FusedExperts(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Parameter(torch.randn(n_experts, intermediate, hidden))

        container = FusedExperts()
        torch.manual_seed(42)

        shared_dir = torch.randn(hidden, 1)
        shared_dir = shared_dir / shared_dir.norm()

        # Per-expert directions
        expert_dirs = {}
        for ei in range(n_experts):
            d = torch.randn(hidden)
            d = d / d.norm()
            expert_dirs[ei] = d

        orig_data = container.down_proj.data.clone()

        count = AbliterationPipeline._project_fused_3d_granular(
            container, shared_dir, expert_dirs, ["down_proj"],
            norm_preserve=False, scale=1.0,
        )

        assert count == n_experts, f"Should project {n_experts} experts, got {count}"

        # Each expert should be modified
        for ei in range(n_experts):
            assert not torch.allclose(
                container.down_proj.data[ei], orig_data[ei]
            ), f"Expert {ei} should be modified"

    def test_fallback_to_shared_on_fused(self):
        """Experts without per-expert direction should use shared direction."""
        hidden = 16
        intermediate = 32
        n_experts = 4

        class FusedExperts(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Parameter(torch.randn(n_experts, intermediate, hidden))

        container = FusedExperts()
        torch.manual_seed(42)

        shared_dir = torch.randn(hidden, 1)
        shared_dir = shared_dir / shared_dir.norm()

        # Only expert 0 has a direction
        expert_dirs = {0: torch.randn(hidden).div_(1.0)}
        expert_dirs[0] = expert_dirs[0] / expert_dirs[0].norm()

        orig_data = container.down_proj.data.clone()

        count = AbliterationPipeline._project_fused_3d_granular(
            container, shared_dir, expert_dirs, ["down_proj"],
            norm_preserve=False, scale=1.0,
        )

        assert count == n_experts
        # All experts should be modified (experts 1-3 use shared dir)
        for ei in range(n_experts):
            assert not torch.allclose(
                container.down_proj.data[ei], orig_data[ei]
            ), f"Expert {ei} should be modified"

    def test_norm_preserve_on_fused(self):
        """Fused 3D with norm_preserve should maintain per-expert norms."""
        hidden = 16
        intermediate = 32
        n_experts = 4

        class FusedExperts(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Parameter(torch.randn(n_experts, intermediate, hidden))

        container = FusedExperts()
        torch.manual_seed(42)

        shared_dir = torch.randn(hidden, 1)
        shared_dir = shared_dir / shared_dir.norm()

        expert_dirs = {}
        for ei in range(n_experts):
            d = torch.randn(hidden)
            expert_dirs[ei] = d / d.norm()

        orig_norms = [container.down_proj.data[i].norm().item() for i in range(n_experts)]

        AbliterationPipeline._project_fused_3d_granular(
            container, shared_dir, expert_dirs, ["down_proj"],
            norm_preserve=True, scale=1.0,
        )

        for i in range(n_experts):
            new_norm = container.down_proj.data[i].norm().item()
            assert abs(orig_norms[i] - new_norm) < 1e-3, (
                f"Expert {i} norm not preserved: {orig_norms[i]:.4f} vs {new_norm:.4f}"
            )

    def test_skips_non_3d_params(self):
        """Should skip parameters that are not 3-dimensional."""
        hidden = 16

        class FlatExperts(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Parameter(torch.randn(32, hidden))

        container = FlatExperts()
        shared_dir = torch.randn(hidden, 1)
        shared_dir = shared_dir / shared_dir.norm()

        count = AbliterationPipeline._project_fused_3d_granular(
            container, shared_dir, {}, ["down_proj"],
            norm_preserve=False, scale=1.0,
        )
        assert count == 0


class TestEGAExciseIntegration:
    """Test that EGA integrates properly in the excise stage path."""

    def test_ega_pipeline_flags(self):
        """Pipeline with surgical method should enable per_expert_directions."""
        pipeline = AbliterationPipeline(model_name="test", method="surgical")
        assert pipeline.per_expert_directions is True

    def test_ega_only_on_primary_direction(self):
        """EGA should only apply for dir_idx==0, not higher SVD directions."""
        # This is enforced by the `and dir_idx == 0` check in _excise
        # We verify the code structure exists
        from obliteratus.abliterate import AbliterationPipeline
        import inspect
        source = inspect.getsource(AbliterationPipeline._excise_inner)
        assert "dir_idx == 0" in source, "EGA should only apply for primary direction"
        assert "_project_moe_experts_granular" in source, "EGA method should be called in excise"

    def test_ega_distill_integration(self):
        """EGA should be called during distill when per_expert_directions is enabled."""
        from obliteratus.abliterate import AbliterationPipeline
        import inspect
        source = inspect.getsource(AbliterationPipeline._distill)
        assert "_compute_expert_granular_directions" in source
        assert "per_expert_directions" in source

    def test_nuclear_method_enables_ega(self):
        """Nuclear method should also enable per_expert_directions."""
        cfg = METHODS["nuclear"]
        assert cfg["per_expert_directions"] is True
        pipeline = AbliterationPipeline(model_name="test", method="nuclear")
        assert pipeline.per_expert_directions is True

    def test_basic_method_disables_ega(self):
        """Basic method should not enable per_expert_directions."""
        cfg = METHODS["basic"]
        assert cfg.get("per_expert_directions", False) is False

    def test_inverted_method_enables_ega(self):
        """Inverted method should enable per_expert_directions."""
        cfg = METHODS["inverted"]
        assert cfg["per_expert_directions"] is True

    def test_ega_with_routing_data_end_to_end(self):
        """End-to-end: EGA computes directions and granular projection modifies weights."""
        hidden = 16
        n_experts = 4

        class FakeExpert(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.down_proj = torch.nn.Linear(hidden, 32, bias=False)
                self.up_proj = torch.nn.Linear(hidden, 32, bias=False)

        class FakeMoE(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.gate = torch.nn.Linear(hidden, n_experts, bias=False)
                self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(n_experts)])

        moe = FakeMoE()
        torch.manual_seed(42)
        for p in moe.parameters():
            p.data = torch.randn_like(p.data)

        pipeline = AbliterationPipeline(model_name="test", method="surgical")
        pipeline._on_log = lambda m: None
        pipeline._on_stage = lambda r: None
        pipeline._strong_layers = [0]

        # Simulate EGA routing data
        h_logits = [torch.randn(n_experts) for _ in range(5)]
        s_logits = [torch.randn(n_experts) for _ in range(5)]
        pipeline._routing_harmful = {0: h_logits}
        pipeline._routing_harmless = {0: s_logits}

        # Simulate activations with clear separation
        refusal_dir = torch.randn(hidden)
        refusal_dir = refusal_dir / refusal_dir.norm()
        pipeline._harmful_acts = {0: [torch.randn(hidden) + 2 * refusal_dir for _ in range(5)]}
        pipeline._harmless_acts = {0: [torch.randn(hidden) - 2 * refusal_dir for _ in range(5)]}

        # Step 1: compute EGA directions
        pipeline._compute_expert_granular_directions()
        assert 0 in pipeline._expert_directions
        assert len(pipeline._expert_directions[0]) > 0

        # Step 2: apply granular projection
        shared_dir = torch.randn(hidden, 1)
        shared_dir = shared_dir / shared_dir.norm()

        orig_expert0 = moe.experts[0].down_proj.weight.data.clone()

        count = pipeline._project_moe_experts_granular(
            moe, shared_dir, layer_idx=0,
        )

        assert count > 0
        assert not torch.allclose(moe.experts[0].down_proj.weight.data, orig_expert0), \
            "Expert weights should be modified by EGA"