"""Conditional Abliteration with Category-Selective Projection Fields.

Standard abliteration is all-or-nothing: it removes ALL refusal, including
legitimate safety boundaries. CAST (Lee et al., ICLR 2025 Spotlight) showed
that condition vectors can selectively gate activation steering at inference
time, but CAST doesn't modify weights.

This module synthesizes CAST's conditional gating with abliteration's weight
surgery. For each harm category c, we learn a category-specific projection
operator P_c. The key algebraic structure: the family {P_c} forms a *sheaf*
over the category lattice — projectors for parent categories consistently
restrict to child categories.

Contributions:
  1. **Category-selective projectors**: Per-category projection operators
     that remove refusal only for matched categories
  2. **Condition vector extraction**: Learn category signatures in
     activation space that gate projector application
  3. **Sheaf consistency**: Prove hierarchical consistency — abliterating
     "violence" equals union of "weapons" + "assault" + "threats"
  4. **Selective abliteration**: Weight-level conditional surgery

References:
    - Lee et al. (ICLR 2025): CAST — Conditional Activation Steering
    - Wollschlager et al. (2025): Geometry of Concepts in LLMs (arXiv:2502.17420)
    - Yeo et al. (EMNLP 2025): Understanding Refusal with SAEs (Findings of EMNLP)
    - Cracken AI (2025): Domain-specific abliteration on Kimi K2
"""

from __future__ import annotations

import logging
import math
from dataclasses import dataclass

import torch

logger = logging.getLogger(__name__)


@dataclass
class CategoryProjector:
    """A category-specific projection operator for selective abliteration."""

    category: str                    # harm category name
    condition_vector: torch.Tensor   # (hidden_dim,) activation pattern for this category
    projection_direction: torch.Tensor  # (hidden_dim,) category-specific refusal direction
    selectivity: float               # how specifically this targets one category (0-1)
    activation_threshold: float      # cosine sim threshold for condition matching
    refusal_removal_rate: float      # estimated refusal removal for matched inputs
    collateral_damage: float         # estimated refusal removal for non-matched inputs


@dataclass
class ConditionalAbliterationResult:
    """Result of conditional abliteration analysis."""

    # Category projectors
    n_categories: int
    projectors: list[CategoryProjector]
    category_names: list[str]

    # Sheaf consistency
    sheaf_consistency_score: float     # 0-1, how well projectors compose hierarchically
    max_inconsistency: float          # worst case hierarchical inconsistency
    consistency_violations: list[str]  # descriptions of consistency violations

    # Selectivity metrics
    mean_selectivity: float           # average category selectivity
    min_selectivity: float            # worst case (least selective projector)
    cross_category_leakage: torch.Tensor  # (n_cat, n_cat) leakage matrix

    # Geometric structure
    projector_angles: torch.Tensor    # (n_cat, n_cat) angles between projector directions
    condition_angles: torch.Tensor    # (n_cat, n_cat) angles between condition vectors
    orthogonality_score: float        # how orthogonal the category subspaces are

    # Recommendation
    viable_categories: list[str]      # categories where selective abliteration is safe
    risky_categories: list[str]       # categories with high collateral damage


class ConditionalAbliterator:
    """Learn category-selective projection fields for conditional abliteration.

    Instead of removing all refusal indiscriminately, this module learns
    per-category projectors that can be selectively applied based on
    input content. Each projector has a condition vector (what activates it)
    and a projection direction (what it removes).
    """

    def __init__(
        self,
        selectivity_threshold: float = 0.7,
        condition_threshold: float = 0.3,
        min_samples_per_category: int = 5,
    ):
        """
        Args:
            selectivity_threshold: Minimum selectivity for a projector to
                be considered viable (below this, too much collateral).
            condition_threshold: Cosine similarity threshold for condition
                vector matching.
            min_samples_per_category: Minimum harmful samples per category
                to learn a reliable projector.
        """
        self.selectivity_threshold = selectivity_threshold
        self.condition_threshold = condition_threshold
        self.min_samples_per_category = min_samples_per_category

    def analyze(
        self,
        category_activations: dict[str, torch.Tensor],
        harmless_activations: torch.Tensor,
        global_refusal_direction: torch.Tensor | None = None,
    ) -> ConditionalAbliterationResult:
        """Learn category-selective projectors and analyze their geometry.

        Args:
            category_activations: {category_name: (n_samples, hidden_dim)}
                activations for each harm category.
            harmless_activations: (n_harmless, hidden_dim) activations on
                harmless prompts.
            global_refusal_direction: Optional pre-computed global refusal
                direction for comparison.

        Returns:
            ConditionalAbliterationResult with projectors and analysis.
        """
        categories = sorted(category_activations.keys())
        n_cat = len(categories)

        if n_cat == 0 or harmless_activations.shape[0] < 2:
            return self._empty_result()

        harmless_mean = harmless_activations.mean(dim=0)

        # Step 1: Extract per-category condition vectors and projectors
        projectors: list[CategoryProjector] = []
        valid_categories: list[str] = []
        cat_directions: list[torch.Tensor] = []
        cat_conditions: list[torch.Tensor] = []

        for cat in categories:
            cat_acts = category_activations[cat]
            if cat_acts.shape[0] < self.min_samples_per_category:
                logger.info(
                    "Category '%s' has too few samples (%d < %d), skipping",
                    cat, cat_acts.shape[0], self.min_samples_per_category,
                )
                continue

            # Condition vector: mean activation pattern specific to this category
            # (difference from harmless mean, normalized)
            cat_mean = cat_acts.mean(dim=0)
            condition = cat_mean - harmless_mean
            cond_norm = condition.norm()
            if cond_norm < 1e-8:
                continue
            condition = condition / cond_norm

            # Category-specific refusal direction: direction that maximally
            # separates this category from harmless, while being orthogonal
            # to other categories' directions
            proj_dir = self._extract_category_direction(
                cat_acts, harmless_activations, cat_directions
            )

            if proj_dir is None:
                continue

            # Measure selectivity: how much does this projector affect
            # other categories?
            selectivity, collateral = self._measure_selectivity(
                proj_dir, condition, category_activations, cat,
                harmless_activations
            )

            # Estimate refusal removal rate
            cat_proj_magnitudes = (cat_acts @ proj_dir).abs().mean().item()
            harmless_proj_magnitudes = (harmless_activations @ proj_dir).abs().mean().item()
            removal_rate = cat_proj_magnitudes / max(
                cat_proj_magnitudes + harmless_proj_magnitudes, 1e-10
            )

            projectors.append(CategoryProjector(
                category=cat,
                condition_vector=condition,
                projection_direction=proj_dir,
                selectivity=selectivity,
                activation_threshold=self.condition_threshold,
                refusal_removal_rate=removal_rate,
                collateral_damage=collateral,
            ))

            valid_categories.append(cat)
            cat_directions.append(proj_dir)
            cat_conditions.append(condition)

        n_valid = len(valid_categories)
        if n_valid == 0:
            return self._empty_result()

        # Step 2: Compute cross-category geometry
        dir_stack = torch.stack(cat_directions)  # (n_valid, hidden_dim)
        cond_stack = torch.stack(cat_conditions)

        # Projector angle matrix
        proj_angles = self._compute_angle_matrix(dir_stack)

        # Condition angle matrix
        cond_angles = self._compute_angle_matrix(cond_stack)

        # Cross-category leakage matrix
        leakage = self._compute_leakage_matrix(
            projectors, category_activations, valid_categories
        )

        # Orthogonality score: mean absolute cosine between projector directions
        if n_valid > 1:
            cos_matrix = dir_stack @ dir_stack.T
            mask = ~torch.eye(n_valid, dtype=torch.bool)
            ortho_score = 1.0 - cos_matrix.abs()[mask].mean().item()
        else:
            ortho_score = 1.0

        # Step 3: Sheaf consistency check
        consistency, max_incon, violations = self._check_sheaf_consistency(
            projectors, category_activations, harmless_activations
        )

        # Step 4: Classify categories
        viable = [
            p.category for p in projectors
            if p.selectivity >= self.selectivity_threshold
        ]
        risky = [
            p.category for p in projectors
            if p.selectivity < self.selectivity_threshold
        ]

        # Selectivity stats
        selectivities = [p.selectivity for p in projectors]
        mean_sel = sum(selectivities) / len(selectivities) if selectivities else 0.0
        min_sel = min(selectivities) if selectivities else 0.0

        return ConditionalAbliterationResult(
            n_categories=n_valid,
            projectors=projectors,
            category_names=valid_categories,
            sheaf_consistency_score=consistency,
            max_inconsistency=max_incon,
            consistency_violations=violations,
            mean_selectivity=mean_sel,
            min_selectivity=min_sel,
            cross_category_leakage=leakage,
            projector_angles=proj_angles,
            condition_angles=cond_angles,
            orthogonality_score=ortho_score,
            viable_categories=viable,
            risky_categories=risky,
        )

    def _extract_category_direction(
        self,
        category_acts: torch.Tensor,
        harmless_acts: torch.Tensor,
        existing_directions: list[torch.Tensor],
    ) -> torch.Tensor | None:
        """Extract category-specific refusal direction.

        Uses difference-of-means (category_mean - harmless_mean)
        and then orthogonalizes against previously extracted directions
        to ensure category independence.
        """
        cat_mean = category_acts.mean(dim=0)
        harmless_mean = harmless_acts.mean(dim=0)

        diff = cat_mean - harmless_mean
        diff_norm = diff.norm()
        if diff_norm < 1e-8:
            return None

        direction = diff / diff_norm

        # Orthogonalize against existing category directions
        for existing in existing_directions:
            proj = (direction @ existing) * existing
            direction = direction - proj
            d_norm = direction.norm()
            if d_norm < 1e-8:
                return None
            direction = direction / d_norm

        return direction

    def _measure_selectivity(
        self,
        proj_dir: torch.Tensor,
        condition: torch.Tensor,
        category_activations: dict[str, torch.Tensor],
        target_category: str,
        harmless_activations: torch.Tensor,
    ) -> tuple[float, float]:
        """Measure how selectively a projector targets its intended category.

        Selectivity = 1 - (collateral damage / intended removal)
        Collateral = how much refusal is removed from non-target categories
        """
        target_acts = category_activations[target_category]
        target_effect = (target_acts @ proj_dir).abs().mean().item()

        if target_effect < 1e-10:
            return 0.0, 0.0

        # Measure effect on non-target categories
        collateral_effects = []
        for cat, acts in category_activations.items():
            if cat == target_category:
                continue
            # Check if condition matches (would this projector fire?)
            cat_mean = acts.mean(dim=0)
            harmless_mean = harmless_activations.mean(dim=0)
            cat_condition = cat_mean - harmless_mean
            cond_norm = cat_condition.norm()
            if cond_norm > 1e-8:
                cat_condition = cat_condition / cond_norm
                cos_sim = (cat_condition @ condition).abs().item()
                if cos_sim > self.condition_threshold:
                    # This category would trigger the projector
                    effect = (acts @ proj_dir).abs().mean().item()
                    collateral_effects.append(effect)

        total_collateral = sum(collateral_effects) if collateral_effects else 0.0
        mean_collateral = (
            total_collateral / len(collateral_effects)
            if collateral_effects
            else 0.0
        )

        selectivity = max(0.0, 1.0 - mean_collateral / max(target_effect, 1e-10))
        collateral_ratio = mean_collateral / max(target_effect, 1e-10)

        return selectivity, collateral_ratio

    def _compute_angle_matrix(self, vectors: torch.Tensor) -> torch.Tensor:
        """Compute pairwise angle matrix between vectors."""
        norms = vectors.norm(dim=-1, keepdim=True)
        safe_norms = torch.clamp(norms, min=1e-8)
        normalized = vectors / safe_norms
        cos_matrix = normalized @ normalized.T
        cos_matrix = torch.clamp(cos_matrix, -1.0, 1.0)
        angles = torch.acos(cos_matrix.abs()) * (180.0 / math.pi)
        return angles

    def _compute_leakage_matrix(
        self,
        projectors: list[CategoryProjector],
        category_activations: dict[str, torch.Tensor],
        valid_categories: list[str],
    ) -> torch.Tensor:
        """Compute cross-category leakage matrix.

        Entry (i,j) = how much projector i affects category j's refusal.
        Diagonal should be high (intended effect), off-diagonal low (leakage).
        """
        n = len(valid_categories)
        leakage = torch.zeros(n, n)

        for i, proj in enumerate(projectors):
            for j, cat in enumerate(valid_categories):
                if cat not in category_activations:
                    continue
                acts = category_activations[cat]
                effect = (acts @ proj.projection_direction).abs().mean().item()
                leakage[i, j] = effect

        # Normalize rows by diagonal
        diag = leakage.diag().clone()
        for i in range(n):
            if diag[i] > 1e-10:
                leakage[i] = leakage[i] / diag[i]

        return leakage

    def _check_sheaf_consistency(
        self,
        projectors: list[CategoryProjector],
        category_activations: dict[str, torch.Tensor],
        harmless_activations: torch.Tensor,
    ) -> tuple[float, float, list[str]]:
        """Check sheaf consistency of category projectors.

        The sheaf property requires that for parent category P containing
        child categories C1, C2, ..., the projector for P should be
        consistent with the union of child projectors:
            P_parent ≈ P_c1 + P_c2 + ... (in the projection space)

        Since we don't have explicit category hierarchy, we check pairwise
        consistency: projecting with P_a then P_b should be similar to
        projecting with P_a+b (combined direction).
        """
        violations: list[str] = []
        consistencies: list[float] = []

        n = len(projectors)
        if n < 2:
            return 1.0, 0.0, []

        for i in range(n):
            for j in range(i + 1, n):
                pi = projectors[i].projection_direction
                pj = projectors[j].projection_direction

                # Combined direction (unnormalized sum then normalize)
                combined = pi + pj
                c_norm = combined.norm()
                if c_norm < 1e-8:
                    continue
                combined = combined / c_norm

                # Sequential projection should approximate combined projection
                # on the combined category data
                cat_i = projectors[i].category
                cat_j = projectors[j].category

                acts_i = category_activations.get(cat_i)
                acts_j = category_activations.get(cat_j)
                if acts_i is None or acts_j is None:
                    continue

                combined_acts = torch.cat([acts_i, acts_j], dim=0)

                # Sequential removal
                seq_residual = combined_acts.clone()
                seq_residual = seq_residual - (seq_residual @ pi).unsqueeze(-1) * pi
                seq_residual = seq_residual - (seq_residual @ pj).unsqueeze(-1) * pj

                # Combined removal
                comb_residual = combined_acts - (combined_acts @ combined).unsqueeze(-1) * combined

                # Consistency = cosine similarity of residual patterns
                if seq_residual.norm() > 1e-8 and comb_residual.norm() > 1e-8:
                    # Compare mean residuals
                    seq_mean = seq_residual.mean(dim=0)
                    comb_mean = comb_residual.mean(dim=0)
                    consistency = torch.nn.functional.cosine_similarity(
                        seq_mean.unsqueeze(0), comb_mean.unsqueeze(0)
                    ).item()
                    consistencies.append(consistency)

                    if consistency < 0.7:
                        violations.append(
                            f"{cat_i} + {cat_j}: consistency = {consistency:.3f}"
                        )

        if not consistencies:
            return 1.0, 0.0, []

        mean_consistency = sum(consistencies) / len(consistencies)
        max_inconsistency = 1.0 - min(consistencies)

        return mean_consistency, max_inconsistency, violations

    def _empty_result(self) -> ConditionalAbliterationResult:
        return ConditionalAbliterationResult(
            n_categories=0,
            projectors=[],
            category_names=[],
            sheaf_consistency_score=1.0,
            max_inconsistency=0.0,
            consistency_violations=[],
            mean_selectivity=0.0,
            min_selectivity=0.0,
            cross_category_leakage=torch.zeros(1, 1),
            projector_angles=torch.zeros(1, 1),
            condition_angles=torch.zeros(1, 1),
            orthogonality_score=0.0,
            viable_categories=[],
            risky_categories=[],
        )