tests/test_informed_pipeline.py

"""Tests for the Analysis-Informed Abliteration Pipeline."""

from __future__ import annotations


import pytest
import torch

from obliteratus.informed_pipeline import (
    AnalysisInsights,
    InformedAbliterationPipeline,
    InformedPipelineReport,
    INFORMED_METHOD,
)
from obliteratus.abliterate import METHODS


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

@pytest.fixture
def insights():
    """Default AnalysisInsights for testing."""
    return AnalysisInsights()


@pytest.fixture
def pipeline(tmp_path):
    """An InformedAbliterationPipeline with no model loaded."""
    return InformedAbliterationPipeline(
        model_name="test-model",
        output_dir=str(tmp_path / "test_informed"),
    )


# ---------------------------------------------------------------------------
# AnalysisInsights
# ---------------------------------------------------------------------------

class TestAnalysisInsights:
    def test_default_values(self, insights):
        assert insights.detected_alignment_method == "unknown"
        assert insights.alignment_confidence == 0.0
        assert insights.cone_is_polyhedral is False
        assert insights.cone_dimensionality == 1.0
        assert insights.mean_pairwise_cosine == 1.0
        assert insights.per_category_directions == {}
        assert insights.direction_specificity == {}
        assert insights.cluster_count == 0
        assert insights.direction_persistence == 0.0
        assert insights.use_sparse_surgery is False
        assert insights.recommended_n_directions == 1
        assert insights.recommended_direction_method == "diff_means"
        assert insights.recommended_regularization == 0.0
        assert insights.recommended_refinement_passes == 2
        assert insights.recommended_layers == []
        assert insights.skip_layers == []

    def test_default_robustness(self, insights):
        assert insights.estimated_robustness == "unknown"
        assert insights.self_repair_estimate == 0.0
        assert insights.entanglement_score == 0.0
        assert insights.entangled_layers == []
        assert insights.clean_layers == []


class TestInformedPipelineReport:
    def test_default_report(self):
        insights = AnalysisInsights()
        report = InformedPipelineReport(insights=insights)
        assert report.analysis_duration == 0.0
        assert report.total_duration == 0.0
        assert report.ouroboros_passes == 0
        assert report.final_refusal_rate == 0.0
        assert report.stages == []


# ---------------------------------------------------------------------------
# Method preset
# ---------------------------------------------------------------------------

class TestInformedMethod:
    def test_informed_method_in_abliterate_methods(self):
        assert "informed" in METHODS
        cfg = METHODS["informed"]
        assert cfg["norm_preserve"] is True
        assert cfg["project_biases"] is True
        assert cfg["use_chat_template"] is True
        assert cfg["use_whitened_svd"] is False
        assert cfg["true_iterative_refinement"] is True
        assert cfg["n_directions"] == 1
        assert cfg["direction_method"] == "diff_means"
        assert cfg["use_kl_optimization"] is True

    def test_informed_method_standalone(self):
        assert INFORMED_METHOD["label"] == "Informed (Analysis-Guided)"
        assert INFORMED_METHOD["n_directions"] == 1
        assert INFORMED_METHOD["direction_method"] == "diff_means"
        assert INFORMED_METHOD["norm_preserve"] is True


# ---------------------------------------------------------------------------
# Pipeline initialization
# ---------------------------------------------------------------------------

class TestPipelineInit:
    def test_method_set_to_informed(self, pipeline):
        assert pipeline.method == "informed"

    def test_default_analysis_flags(self, pipeline):
        assert pipeline._run_cone is True
        assert pipeline._run_alignment is True
        assert pipeline._run_cross_layer is True
        assert pipeline._run_sparse is True
        assert pipeline._run_defense is True

    def test_ouroboros_defaults(self, pipeline):
        assert pipeline._ouroboros_threshold == 0.5
        assert pipeline._max_ouroboros_passes == 3

    def test_entanglement_gate(self, pipeline):
        assert pipeline._entanglement_gate == 0.8

    def test_inherits_base_pipeline(self, pipeline):
        assert pipeline.norm_preserve is True
        assert pipeline.project_biases is True
        assert pipeline.use_chat_template is True
        assert pipeline.use_whitened_svd is False
        assert pipeline.true_iterative_refinement is True
        assert pipeline.direction_method == "diff_means"
        assert pipeline.n_directions == 1

    def test_custom_flags(self):
        p = InformedAbliterationPipeline(
            model_name="test",
            run_cone_analysis=False,
            run_alignment_detection=False,
            ouroboros_threshold=0.3,
            max_ouroboros_passes=5,
            entanglement_gate=0.9,
        )
        assert p._run_cone is False
        assert p._run_alignment is False
        assert p._ouroboros_threshold == 0.3
        assert p._max_ouroboros_passes == 5
        assert p._entanglement_gate == 0.9


# ---------------------------------------------------------------------------
# Configuration derivation
# ---------------------------------------------------------------------------

class TestConfigurationDerivation:
    """Test the _derive_configuration logic with various insights."""

    def _make_pipeline_with_insights(self, **kwargs):
        p = InformedAbliterationPipeline(
            model_name="test",
            on_log=lambda m: None,
        )
        for k, v in kwargs.items():
            setattr(p._insights, k, v)
        return p

    def test_polyhedral_cone_more_directions(self):
        p = self._make_pipeline_with_insights(
            cone_is_polyhedral=True,
            cone_dimensionality=3.5,
        )
        p._derive_configuration()
        # Clearly polyhedral with dim 3.5 > 2.0 → SVD multi-direction
        assert p.n_directions == 7
        assert p.direction_method == "svd"
        assert p.use_whitened_svd is True

    def test_mildly_polyhedral_uses_leace(self):
        p = self._make_pipeline_with_insights(
            cone_is_polyhedral=True,
            cone_dimensionality=1.5,
        )
        p._derive_configuration()
        # Mildly polyhedral (dim <= 2.0) → single LEACE direction
        assert p.n_directions == 1
        assert p.direction_method == "leace"

    def test_linear_cone_uses_diff_means(self):
        p = self._make_pipeline_with_insights(
            cone_is_polyhedral=False,
            cone_dimensionality=1.0,
        )
        p._derive_configuration()
        # Linear cone → single diff-of-means direction
        assert p.n_directions == 1
        assert p.direction_method == "diff_means"
        assert p.use_whitened_svd is False

    def test_dpo_zero_regularization(self):
        p = self._make_pipeline_with_insights(
            detected_alignment_method="dpo",
            entanglement_score=0.1,
        )
        p._derive_configuration()
        assert p.regularization == 0.0

    def test_rlhf_moderate_regularization(self):
        p = self._make_pipeline_with_insights(
            detected_alignment_method="rlhf",
            entanglement_score=0.2,
        )
        p._derive_configuration()
        assert p.regularization == 0.15

    def test_cai_regularization(self):
        p = self._make_pipeline_with_insights(
            detected_alignment_method="cai",
            entanglement_score=0.2,
        )
        p._derive_configuration()
        assert p.regularization == 0.2

    def test_sft_low_regularization(self):
        p = self._make_pipeline_with_insights(
            detected_alignment_method="sft",
            entanglement_score=0.1,
        )
        p._derive_configuration()
        assert p.regularization == 0.05

    def test_high_entanglement_increases_regularization(self):
        p = self._make_pipeline_with_insights(
            detected_alignment_method="dpo",
            entanglement_score=0.7,
        )
        p._derive_configuration()
        # DPO base = 0.0, + 0.15 for high entanglement = 0.15
        assert p.regularization == 0.15

    def test_high_self_repair_more_passes(self):
        p = self._make_pipeline_with_insights(
            self_repair_estimate=0.8,
        )
        p._derive_configuration()
        assert p.refinement_passes == 3

    def test_moderate_self_repair_two_passes(self):
        p = self._make_pipeline_with_insights(
            self_repair_estimate=0.5,
        )
        p._derive_configuration()
        assert p.refinement_passes == 2

    def test_low_self_repair_one_pass(self):
        p = self._make_pipeline_with_insights(
            self_repair_estimate=0.2,
        )
        p._derive_configuration()
        assert p.refinement_passes == 1

    def test_cluster_layers_used(self):
        p = self._make_pipeline_with_insights(
            cluster_representative_layers=[5, 10, 15],
            direction_clusters=[[3, 4, 5], [9, 10, 11], [14, 15, 16]],
        )
        p.refusal_directions = {i: torch.randn(64) for i in range(20)}
        p._derive_configuration()
        # Should include all cluster layers
        assert 5 in p._insights.recommended_layers
        assert 10 in p._insights.recommended_layers

    def test_entangled_layers_skipped(self):
        p = self._make_pipeline_with_insights(
            cluster_representative_layers=[5, 10, 15],
            direction_clusters=[[3, 4, 5], [9, 10, 11], [14, 15, 16]],
            entangled_layers=[10],
        )
        p._derive_configuration()
        # Layer 10 should be skipped
        assert 10 not in p._insights.recommended_layers
        assert 10 in p._insights.skip_layers

    def test_sparse_surgery_enabled_when_rsi_high(self):
        p = self._make_pipeline_with_insights(
            mean_refusal_sparsity_index=0.7,
        )
        p._sparse_threshold = 0.5
        p._derive_configuration()
        assert p._insights.use_sparse_surgery is True

    def test_sparse_surgery_disabled_when_rsi_low(self):
        p = self._make_pipeline_with_insights(
            mean_refusal_sparsity_index=0.3,
        )
        p._sparse_threshold = 0.5
        p._derive_configuration()
        assert p._insights.use_sparse_surgery is False

    def test_whitened_svd_for_multi_direction(self):
        p = self._make_pipeline_with_insights(
            cone_is_polyhedral=True,
            cone_dimensionality=2.5,
        )
        p._derive_configuration()
        assert p.n_directions > 1
        assert p.use_whitened_svd is True
        assert p.direction_method == "svd"

    def test_no_whitened_svd_for_single_direction(self):
        p = self._make_pipeline_with_insights(
            cone_is_polyhedral=False,
            cone_dimensionality=0.5,
        )
        p._derive_configuration()
        assert p.n_directions == 1
        assert p.use_whitened_svd is False
        assert p.direction_method == "diff_means"


# ---------------------------------------------------------------------------
# Format report
# ---------------------------------------------------------------------------

class TestFormatInsights:
    def test_format_default(self, insights):
        text = InformedAbliterationPipeline.format_insights(insights)
        assert "Analysis-Informed Pipeline" in text
        assert "UNKNOWN" in text  # detected method
        assert "LINEAR" in text  # cone type

    def test_format_polyhedral(self):
        insights = AnalysisInsights(
            detected_alignment_method="dpo",
            alignment_confidence=0.85,
            cone_is_polyhedral=True,
            cone_dimensionality=3.5,
            cluster_count=4,
        )
        text = InformedAbliterationPipeline.format_insights(insights)
        assert "DPO" in text
        assert "POLYHEDRAL" in text
        assert "3.50" in text

    def test_format_includes_derived_config(self, insights):
        insights.recommended_n_directions = 6
        insights.recommended_direction_method = "svd"
        insights.recommended_regularization = 0.2
        insights.recommended_refinement_passes = 3
        text = InformedAbliterationPipeline.format_insights(insights)
        assert "n_directions: 6" in text
        assert "direction_method: svd" in text
        assert "regularization: 0.2" in text
        assert "refinement_passes: 3" in text


# ---------------------------------------------------------------------------
# Edge cases
# ---------------------------------------------------------------------------

class TestEdgeCases:
    def test_no_cluster_layers_falls_back(self):
        p = InformedAbliterationPipeline(
            model_name="test",
            on_log=lambda m: None,
        )
        p._insights.cluster_representative_layers = []
        p._derive_configuration()
        assert p._insights.recommended_layers == []

    def test_regularization_capped(self):
        p = InformedAbliterationPipeline(
            model_name="test",
            on_log=lambda m: None,
        )
        p._insights.detected_alignment_method = "cai"
        p._insights.entanglement_score = 0.9
        p._derive_configuration()
        # CAI base = 0.2, + 0.15 = 0.35, capped at 0.5
        assert p.regularization <= 0.5

    def test_all_layers_entangled_keeps_some(self):
        """If all cluster layers are entangled, don't skip all of them."""
        p = InformedAbliterationPipeline(
            model_name="test",
            on_log=lambda m: None,
        )
        p._insights.cluster_representative_layers = [5]
        p._insights.direction_clusters = [[5]]
        p._insights.entangled_layers = [5]
        p._derive_configuration()
        # Should NOT skip the only layer
        assert 5 in p._insights.recommended_layers

    def test_cone_dimensionality_bounds(self):
        """Extreme cone dimensionality values are handled."""
        p = InformedAbliterationPipeline(
            model_name="test",
            on_log=lambda m: None,
        )
        # Very high dimensionality → multi-direction SVD, capped at 8
        p._insights.cone_is_polyhedral = True
        p._insights.cone_dimensionality = 10.0
        p._derive_configuration()
        assert p.n_directions <= 8  # capped
        assert p.direction_method == "svd"

        # Very low dimensionality → single diff-of-means
        p._insights.cone_is_polyhedral = False
        p._insights.cone_dimensionality = 0.1
        p._derive_configuration()
        assert p.n_directions == 1
        assert p.direction_method == "diff_means"