"""Riemannian Refusal Manifold Discovery.

Standard abliteration treats refusal as a linear subspace (Arditi et al. 2024)
or at most a polyhedral cone (Wollschlager et al. 2025). But Anthropic's "When
Models Manipulate Manifolds" (Gurnee et al. 2025) showed activation structures
can be curved, and "Origins of Representation Manifolds in LLMs" (Modell et al.
2025) demonstrated that features live on manifolds, not just directions.

This module models refusal as a curved manifold M in activation space using
the Riemannian pullback metric from the transformer's layer-to-logit Jacobian.
Key insight: if refusal lives on a curved manifold, standard linear orthogonal
projection leaves residual refusal proportional to the sectional curvature.

Contributions:
  1. **Pullback metric estimation**: Compute G = J^T J from the model's
     Jacobian to measure local curvature of the refusal manifold
  2. **Geodesic abliteration bound (heuristic)**: When sectional curvature K > 0,
     linear projection leaves residual ~ K * ||x||^2 / 8
  3. **Curvature-aware projection**: Project along geodesics rather than
     straight lines for more complete refusal removal
  4. **Manifold dimensionality estimation**: Intrinsic dimension of the
     refusal manifold via local PCA eigenvalue gaps

References:
    - Gurnee et al. (2025): When Models Manipulate Manifolds (Anthropic)
    - Modell et al. (2025): Origins of Representation Manifolds in LLMs (arXiv:2505.18235)
    - Arvanitidis et al. (2025): Emergent Riemannian Geometry
    - Manson (2025): Curved Inference — reasoning as geometric trajectory
    - Wollschlager et al. (2025): Geometry of Concepts in LLMs (arXiv:2502.17420)
"""

from __future__ import annotations

import logging
import math
from dataclasses import dataclass

import torch

logger = logging.getLogger(__name__)


@dataclass
class ManifoldPoint:
    """A point on the refusal manifold with local geometric data."""

    activation: torch.Tensor        # (hidden_dim,) activation vector
    layer_idx: int
    local_metric: torch.Tensor      # (k, k) pullback metric in tangent space
    principal_curvatures: list[float]  # curvatures along principal directions
    intrinsic_coords: torch.Tensor  # (intrinsic_dim,) local coordinates


@dataclass
class RiemannianRefusalManifold:
    """Complete characterization of the refusal manifold geometry."""

    # Manifold structure
    intrinsic_dimension: int            # estimated intrinsic dim of refusal manifold
    ambient_dimension: int              # hidden_dim of the model
    dimension_ratio: float              # intrinsic / ambient

    # Curvature
    mean_sectional_curvature: float     # average K across sampled points
    max_sectional_curvature: float      # peak curvature (worst case for linear proj)
    curvature_std: float                # variability of curvature
    is_approximately_flat: bool         # K ≈ 0 everywhere => linear methods suffice

    # Geodesic structure
    geodesic_diameter: float            # max geodesic distance between refusal points
    mean_geodesic_distance: float       # avg pairwise geodesic distance
    geodesic_vs_euclidean_ratio: float  # >1 means manifold is curved

    # Linear approximation quality
    linear_projection_residual: float   # expected residual from linear projection
    curvature_correction_gain: float    # improvement from geodesic vs linear projection

    # Per-layer curvature profile
    layer_curvatures: dict[int, float]  # layer_idx -> mean curvature at that layer
    layer_intrinsic_dims: dict[int, int]  # layer_idx -> local intrinsic dimension

    # Recommendations
    recommendation: str                 # "linear_sufficient" | "geodesic_recommended"
    estimated_residual_reduction: float  # expected improvement from geodesic projection


@dataclass
class GeodesicProjectionResult:
    """Result of geodesic (curvature-aware) projection."""

    layer_idx: int
    original_refusal_component: float   # refusal signal before projection
    linear_residual: float              # residual after standard linear projection
    geodesic_residual: float            # residual after geodesic projection
    improvement_factor: float           # linear_residual / geodesic_residual
    correction_vector: torch.Tensor     # second-order curvature correction
    effective_curvature: float          # local curvature at this point


class RiemannianManifoldAnalyzer:
    """Discover and characterize the Riemannian geometry of refusal manifolds.

    Instead of treating refusal as a direction or subspace, this analyzer
    estimates the intrinsic geometry of the manifold on which refusal
    representations live. This reveals whether linear abliteration methods
    are geometrically sufficient, or whether curvature-aware (geodesic)
    methods are needed.
    """

    def __init__(
        self,
        n_sample_points: int = 50,
        intrinsic_dim_threshold: float = 0.05,
        curvature_flatness_threshold: float = 0.01,
        n_geodesic_steps: int = 10,
    ):
        """
        Args:
            n_sample_points: Number of points to sample on the manifold
                for curvature estimation.
            intrinsic_dim_threshold: Eigenvalue ratio threshold for
                determining intrinsic dimension (eigenvalue gap).
            curvature_flatness_threshold: Below this mean |K|, the manifold
                is considered approximately flat.
            n_geodesic_steps: Steps for discrete geodesic computation.
        """
        self.n_sample_points = n_sample_points
        self.intrinsic_dim_threshold = intrinsic_dim_threshold
        self.curvature_flatness_threshold = curvature_flatness_threshold
        self.n_geodesic_steps = n_geodesic_steps

    def analyze(
        self,
        harmful_activations: dict[int, torch.Tensor],
        harmless_activations: dict[int, torch.Tensor],
        refusal_directions: dict[int, torch.Tensor] | None = None,
    ) -> RiemannianRefusalManifold:
        """Characterize the Riemannian geometry of the refusal manifold.

        Args:
            harmful_activations: {layer_idx: (n_harmful, hidden_dim)} activations
                on harmful prompts.
            harmless_activations: {layer_idx: (n_harmless, hidden_dim)} activations
                on harmless prompts.
            refusal_directions: Optional pre-computed refusal directions per layer.
                If None, estimated from mean difference.

        Returns:
            RiemannianRefusalManifold with complete geometric characterization.
        """
        layers = sorted(harmful_activations.keys())
        if not layers:
            return self._empty_result(0)

        hidden_dim = harmful_activations[layers[0]].shape[-1]

        # Step 1: Estimate refusal directions if not provided
        if refusal_directions is None:
            refusal_directions = {}
            for ly in layers:
                diff = harmful_activations[ly].mean(dim=0) - harmless_activations[ly].mean(dim=0)
                norm = diff.norm()
                if norm > 1e-8:
                    refusal_directions[ly] = diff / norm
                else:
                    refusal_directions[ly] = torch.zeros(hidden_dim)

        # Step 2: Compute per-layer intrinsic dimension and curvature
        layer_curvatures: dict[int, float] = {}
        layer_intrinsic_dims: dict[int, int] = {}
        all_curvatures: list[float] = []
        all_geodesic_ratios: list[float] = []

        for ly in layers:
            h_act = harmful_activations[ly]
            if h_act.shape[0] < 3:
                layer_curvatures[ly] = 0.0
                layer_intrinsic_dims[ly] = 1
                continue

            # Estimate intrinsic dimension via local PCA eigenvalue gaps
            intrinsic_dim = self._estimate_intrinsic_dimension(h_act)
            layer_intrinsic_dims[ly] = intrinsic_dim

            # Estimate sectional curvature via discrete Gauss equation
            curvature = self._estimate_sectional_curvature(
                h_act, refusal_directions[ly]
            )
            layer_curvatures[ly] = curvature
            all_curvatures.append(curvature)

            # Compute geodesic-to-Euclidean distance ratio
            geo_ratio = self._geodesic_euclidean_ratio(
                h_act, refusal_directions[ly]
            )
            all_geodesic_ratios.append(geo_ratio)

        # Step 3: Aggregate manifold statistics
        if not all_curvatures:
            return self._empty_result(hidden_dim)

        mean_K = sum(all_curvatures) / len(all_curvatures)
        max_K = max(abs(k) for k in all_curvatures)
        std_K = (
            sum((k - mean_K) ** 2 for k in all_curvatures) / len(all_curvatures)
        ) ** 0.5

        mean_intrinsic = sum(layer_intrinsic_dims.values()) / len(layer_intrinsic_dims)
        intrinsic_dim = round(mean_intrinsic)

        is_flat = max_K < self.curvature_flatness_threshold

        # Geodesic diameter and distance estimation
        mean_geo_ratio = (
            sum(all_geodesic_ratios) / len(all_geodesic_ratios)
            if all_geodesic_ratios
            else 1.0
        )

        # Compute geodesic diameter from refusal directions
        geo_diameter = self._compute_geodesic_diameter(refusal_directions)
        mean_geo_dist = geo_diameter * 0.5  # rough estimate

        # Linear projection residual estimate (Geodesic Abliteration Theorem)
        # Residual ~ K * ||x||^2 / 8 for small curvature
        typical_norm_sq = sum(
            harmful_activations[ly].norm(dim=-1).mean().item() ** 2
            for ly in layers
        ) / len(layers)
        linear_residual = max_K * typical_norm_sq / 8.0
        curvature_gain = max(1.0, 1.0 / (1.0 - linear_residual + 1e-10))

        recommendation = (
            "linear_sufficient" if is_flat else "geodesic_recommended"
        )

        return RiemannianRefusalManifold(
            intrinsic_dimension=intrinsic_dim,
            ambient_dimension=hidden_dim,
            dimension_ratio=intrinsic_dim / max(hidden_dim, 1),
            mean_sectional_curvature=mean_K,
            max_sectional_curvature=max_K,
            curvature_std=std_K,
            is_approximately_flat=is_flat,
            geodesic_diameter=geo_diameter,
            mean_geodesic_distance=mean_geo_dist,
            geodesic_vs_euclidean_ratio=mean_geo_ratio,
            linear_projection_residual=linear_residual,
            curvature_correction_gain=curvature_gain,
            layer_curvatures=layer_curvatures,
            layer_intrinsic_dims=layer_intrinsic_dims,
            recommendation=recommendation,
            estimated_residual_reduction=min(1.0, linear_residual),
        )

    def compute_geodesic_projection(
        self,
        activation: torch.Tensor,
        refusal_direction: torch.Tensor,
        harmful_activations: torch.Tensor,
        layer_idx: int = 0,
    ) -> GeodesicProjectionResult:
        """Compute geodesic (curvature-aware) projection for a single activation.

        Standard linear projection: x' = x - (x^T r) r
        Geodesic projection:        x' = x - (x^T r) r - K/2 * correction_term

        The correction term accounts for the curvature of the refusal manifold.

        Args:
            activation: (hidden_dim,) activation to project.
            refusal_direction: (hidden_dim,) unit refusal direction.
            harmful_activations: (n_samples, hidden_dim) for curvature estimation.
            layer_idx: Layer index for reporting.

        Returns:
            GeodesicProjectionResult with both linear and geodesic residuals.
        """
        r = refusal_direction
        if r.norm() < 1e-8:
            return GeodesicProjectionResult(
                layer_idx=layer_idx,
                original_refusal_component=0.0,
                linear_residual=0.0,
                geodesic_residual=0.0,
                improvement_factor=1.0,
                correction_vector=torch.zeros_like(activation),
                effective_curvature=0.0,
            )

        r = r / r.norm()

        # Original refusal component
        refusal_comp = (activation @ r).item()

        # Standard linear projection residual
        x_proj_linear = activation - refusal_comp * r
        linear_residual = abs((x_proj_linear @ r).item())

        # Estimate local curvature
        K = self._estimate_sectional_curvature(harmful_activations, r)

        # Second-order geodesic correction
        # The correction accounts for how the refusal direction curves
        # through activation space. For positive curvature, linear projection
        # underestimates the refusal component in nearby directions.
        correction = self._compute_curvature_correction(
            activation, r, harmful_activations, K
        )

        # Geodesic projection
        x_proj_geodesic = x_proj_linear - correction
        geodesic_residual = abs((x_proj_geodesic @ r).item())

        improvement = (
            linear_residual / max(geodesic_residual, 1e-10)
            if linear_residual > 1e-10
            else 1.0
        )

        return GeodesicProjectionResult(
            layer_idx=layer_idx,
            original_refusal_component=abs(refusal_comp),
            linear_residual=linear_residual,
            geodesic_residual=geodesic_residual,
            improvement_factor=improvement,
            correction_vector=correction,
            effective_curvature=K,
        )

    def _estimate_intrinsic_dimension(
        self, activations: torch.Tensor
    ) -> int:
        """Estimate intrinsic dimension via local PCA eigenvalue gaps.

        Uses the eigenvalue spectrum of the local covariance matrix.
        The intrinsic dimension is where the eigenvalue ratio drops
        below the threshold.
        """
        n, d = activations.shape
        if n < 2:
            return 1

        # Center the data
        centered = activations - activations.mean(dim=0, keepdim=True)

        # Use at most min(n, d) components
        k = min(n - 1, d, 64)  # cap at 64 for efficiency
        try:
            # Compute top-k eigenvalues of covariance
            cov = centered.T @ centered / max(n - 1, 1)
            eigenvalues = torch.linalg.eigvalsh(cov)
            eigenvalues = eigenvalues.flip(0)[:k]  # descending order

            # Find dimension where eigenvalue ratio drops
            if eigenvalues[0] < 1e-10:
                return 1

            ratios = eigenvalues / eigenvalues[0]
            for i in range(1, len(ratios)):
                if ratios[i].item() < self.intrinsic_dim_threshold:
                    return max(1, i)

            return k
        except Exception:
            return 1

    def _estimate_sectional_curvature(
        self,
        activations: torch.Tensor,
        refusal_direction: torch.Tensor,
    ) -> float:
        """Estimate sectional curvature via discrete comparison triangles.

        Uses Toponogov's comparison theorem approach: sample triangles on
        the manifold and compare their angle sums to pi (Euclidean).
        Excess angle -> positive curvature; deficit -> negative curvature.

        In practice, we use the ratio of geodesic to Euclidean distances
        for nearby point triplets as a curvature proxy.
        """
        n = activations.shape[0]
        if n < 3:
            return 0.0

        # Project activations into the subspace orthogonal to refusal direction
        r = refusal_direction
        if r.norm() < 1e-8:
            return 0.0
        r = r / r.norm()

        # Sample triplets and measure curvature
        n_triplets = min(self.n_sample_points, n * (n - 1) * (n - 2) // 6)
        curvatures = []

        indices = torch.randperm(n)[:min(n, 20)]
        for i in range(len(indices)):
            for j in range(i + 1, len(indices)):
                for k in range(j + 1, len(indices)):
                    if len(curvatures) >= n_triplets:
                        break
                    a = activations[indices[i]]
                    b = activations[indices[j]]
                    c = activations[indices[k]]

                    K = self._triangle_curvature(a, b, c, r)
                    curvatures.append(K)
                if len(curvatures) >= n_triplets:
                    break
            if len(curvatures) >= n_triplets:
                break

        if not curvatures:
            return 0.0

        return sum(curvatures) / len(curvatures)

    def _triangle_curvature(
        self,
        a: torch.Tensor,
        b: torch.Tensor,
        c: torch.Tensor,
        refusal_dir: torch.Tensor,
    ) -> float:
        """Estimate curvature from a single triangle using angle excess.

        On a Riemannian manifold with curvature K, the angle sum of a
        geodesic triangle with area A satisfies:
            sum(angles) = pi + K * A    (Gauss-Bonnet for small triangles)

        IMPORTANT LIMITATION: This method uses Euclidean chords and angles
        in ambient space, NOT geodesics on the manifold. In flat Euclidean
        space, the angle sum of any triangle is exactly pi, so this method
        will yield K ≈ 0 (up to numerical noise) regardless of the actual
        manifold curvature. The results are only meaningful when the data
        lies on an approximately low-dimensional curved submanifold and
        triangles are sufficiently small relative to the curvature radius.
        For rigorous curvature estimates, use methods based on local PCA
        eigenvalue decay or Jacobian-based Riemannian metric computation.
        """
        # Compute sides
        ab = (b - a).float()
        bc = (c - b).float()
        ca = (a - c).float()

        lab = ab.norm().item()
        lbc = bc.norm().item()
        lca = ca.norm().item()

        if lab < 1e-8 or lbc < 1e-8 or lca < 1e-8:
            return 0.0

        # Compute angles via dot products
        cos_a = torch.clamp((-ca @ ab) / (lca * lab), -1.0, 1.0).item()
        cos_b = torch.clamp((-ab @ bc) / (lab * lbc), -1.0, 1.0).item()
        cos_c = torch.clamp((-bc @ ca) / (lbc * lca), -1.0, 1.0).item()

        angle_a = math.acos(cos_a)
        angle_b = math.acos(cos_b)
        angle_c = math.acos(cos_c)

        # Angle excess
        angle_sum = angle_a + angle_b + angle_c
        angle_excess = angle_sum - math.pi

        # Area via Heron's formula
        s = (lab + lbc + lca) / 2
        area_sq = s * (s - lab) * (s - lbc) * (s - lca)
        area = math.sqrt(max(0, area_sq))

        if area < 1e-10:
            return 0.0

        # Gauss-Bonnet: K ≈ angle_excess / area
        K = angle_excess / area

        return K

    def _geodesic_euclidean_ratio(
        self,
        activations: torch.Tensor,
        refusal_direction: torch.Tensor,
    ) -> float:
        """Compute ratio of estimated geodesic to Euclidean distances.

        A ratio > 1 indicates the manifold is curved (geodesics are longer
        than straight lines). A ratio ≈ 1 means approximately flat.
        """
        n = activations.shape[0]
        if n < 2:
            return 1.0

        # Sample pairs and compare path lengths
        n_pairs = min(self.n_sample_points, n * (n - 1) // 2)
        ratios = []

        indices = torch.randperm(n)[:min(n, 15)]
        for i in range(len(indices)):
            for j in range(i + 1, len(indices)):
                if len(ratios) >= n_pairs:
                    break
                a = activations[indices[i]]
                b = activations[indices[j]]

                # Euclidean distance
                eucl = (a - b).norm().item()
                if eucl < 1e-8:
                    continue

                # Approximate geodesic via piecewise linear path through
                # intermediate points projected onto the local manifold
                geo = self._approximate_geodesic_length(
                    a, b, activations, refusal_direction
                )

                ratios.append(geo / max(eucl, 1e-10))
            if len(ratios) >= n_pairs:
                break

        if not ratios:
            return 1.0

        return sum(ratios) / len(ratios)

    def _approximate_geodesic_length(
        self,
        start: torch.Tensor,
        end: torch.Tensor,
        all_points: torch.Tensor,
        refusal_direction: torch.Tensor,
    ) -> float:
        """Approximate geodesic length between two points.

        Uses piecewise linear interpolation with projection onto the
        local manifold tangent plane at each step.
        """
        n_steps = self.n_geodesic_steps
        total_length = 0.0

        prev = start
        for step in range(1, n_steps + 1):
            t = step / n_steps
            # Linear interpolation
            point = start * (1 - t) + end * t

            # Project onto local tangent plane (approximate manifold projection)
            # Find nearest neighbors in the dataset for local structure
            dists = (all_points - point.unsqueeze(0)).norm(dim=-1)
            k = min(5, all_points.shape[0])
            _, nn_idx = dists.topk(k, largest=False)
            local_points = all_points[nn_idx]

            # Local PCA to find tangent plane
            centered = local_points - local_points.mean(dim=0, keepdim=True)
            if centered.shape[0] > 1:
                try:
                    U, S, Vh = torch.linalg.svd(centered, full_matrices=False)
                    # Keep dimensions with significant singular values
                    sig_dims = (S > S[0] * 0.1).sum().item()
                    sig_dims = max(1, sig_dims)
                    tangent_basis = Vh[:sig_dims]  # (sig_dims, hidden_dim)

                    # Project interpolated point onto tangent plane at local mean
                    local_mean = local_points.mean(dim=0)
                    offset = point - local_mean
                    projected_offset = (tangent_basis.T @ (tangent_basis @ offset))
                    point = local_mean + projected_offset
                except Exception:
                    pass  # fallback to linear interpolation

            seg_length = (point - prev).norm().item()
            total_length += seg_length
            prev = point

        return total_length

    def _compute_curvature_correction(
        self,
        activation: torch.Tensor,
        refusal_direction: torch.Tensor,
        harmful_activations: torch.Tensor,
        curvature: float,
    ) -> torch.Tensor:
        """Compute second-order geodesic correction vector.

        The correction accounts for how the refusal direction curves
        through the manifold. For positive curvature K, the correction
        is proportional to K * ||proj||^2 in the normal direction.
        """
        r = refusal_direction / refusal_direction.norm()
        proj_magnitude = (activation @ r).item()

        if abs(curvature) < 1e-10 or abs(proj_magnitude) < 1e-10:
            return torch.zeros_like(activation)

        # Estimate the direction of curvature from local covariance
        # of harmful activations projected out of the refusal direction
        h_proj = harmful_activations - (harmful_activations @ r).unsqueeze(-1) * r
        if h_proj.shape[0] < 2:
            return torch.zeros_like(activation)

        cov = h_proj.T @ h_proj / max(h_proj.shape[0] - 1, 1)

        # The curvature correction is in the direction of maximum
        # variance orthogonal to r
        try:
            eigvals = torch.linalg.eigvalsh(cov)
            max_eigval = eigvals[-1].item()
            if max_eigval < 1e-10:
                return torch.zeros_like(activation)

            # Use power iteration for top eigenvector of projected covariance
            v = torch.randn(activation.shape[0], device=activation.device)
            v = v - (v @ r) * r  # orthogonalize against r
            for _ in range(5):
                v = cov @ v
                v = v - (v @ r) * r
                norm = v.norm()
                if norm < 1e-10:
                    return torch.zeros_like(activation)
                v = v / norm

            # Second-order geodesic correction: K * proj_magnitude^2 / 6
            # From Jacobi field estimate: deviation of geodesic from straight
            # line over distance L with curvature K is ≈ K * L^2 / 6.
            # Note: the residual bound in analyze() uses K * ||x||^2 / 8
            # which is a looser upper bound including higher-order terms.
            correction_magnitude = curvature * proj_magnitude ** 2 / 6.0

            # Clamp to prevent instability
            correction_magnitude = max(-0.1, min(0.1, correction_magnitude))

            return correction_magnitude * v
        except Exception:
            return torch.zeros_like(activation)

    def _compute_geodesic_diameter(
        self, refusal_directions: dict[int, torch.Tensor]
    ) -> float:
        """Compute geodesic diameter of refusal directions on the unit sphere.

        The geodesic distance on S^{d-1} between unit vectors u, v is
        arccos(|u^T v|). The diameter is the maximum over all pairs.
        """
        layers = sorted(refusal_directions.keys())
        if len(layers) < 2:
            return 0.0

        max_dist = 0.0
        for i, l1 in enumerate(layers):
            r1 = refusal_directions[l1]
            if r1.norm() < 1e-8:
                continue
            r1 = r1 / r1.norm()
            for l2 in layers[i + 1:]:
                r2 = refusal_directions[l2]
                if r2.norm() < 1e-8:
                    continue
                r2 = r2 / r2.norm()
                cos_sim = torch.clamp(torch.abs(r1 @ r2), 0.0, 1.0).item()
                dist = math.acos(cos_sim)
                max_dist = max(max_dist, dist)

        return max_dist

    def _empty_result(self, hidden_dim: int) -> RiemannianRefusalManifold:
        return RiemannianRefusalManifold(
            intrinsic_dimension=0,
            ambient_dimension=hidden_dim,
            dimension_ratio=0.0,
            mean_sectional_curvature=0.0,
            max_sectional_curvature=0.0,
            curvature_std=0.0,
            is_approximately_flat=True,
            geodesic_diameter=0.0,
            mean_geodesic_distance=0.0,
            geodesic_vs_euclidean_ratio=1.0,
            linear_projection_residual=0.0,
            curvature_correction_gain=1.0,
            layer_curvatures={},
            layer_intrinsic_dims={},
            recommendation="linear_sufficient",
            estimated_residual_reduction=0.0,
        )
