"""
Shared helpers for geometry tools.

Enums, typed specs, math utilities, token resolution, and direction
vector extraction used across all geometry tools.
"""

import logging
from enum import Enum
from typing import Any

import numpy as np
from pydantic import BaseModel

from ...steering_store import SteeringVectorRegistry

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Enums and typed specs
# ---------------------------------------------------------------------------


class DirectionType(str, Enum):
    """Valid direction source types for geometry tools."""

    TOKEN = "token"
    NEURON = "neuron"
    RESIDUAL = "residual"
    FFN_OUTPUT = "ffn_output"
    ATTENTION_OUTPUT = "attention_output"
    HEAD_OUTPUT = "head_output"
    STEERING_VECTOR = "steering_vector"


class DirectionSpec(BaseModel):
    """Typed specification for a direction vector."""

    type: DirectionType
    value: str | int | None = None


# ---------------------------------------------------------------------------
# Math helpers
# ---------------------------------------------------------------------------


def _cosine_sim(a: np.ndarray, b: np.ndarray) -> float:
    """Cosine similarity clamped to [-1, 1]."""
    na = float(np.linalg.norm(a))
    nb = float(np.linalg.norm(b))
    if na < 1e-8 or nb < 1e-8:
        return 0.0
    return float(np.clip(np.dot(a, b) / (na * nb), -1.0, 1.0))


def _angle_deg(cos_sim: float) -> float:
    """Cosine similarity -> angle in degrees."""
    return float(np.degrees(np.arccos(np.clip(cos_sim, -1.0, 1.0))))


def _angle_between(a: np.ndarray, b: np.ndarray) -> float:
    """Angle in degrees between two vectors in the full space."""
    return _angle_deg(_cosine_sim(a, b))


def _unit(v: np.ndarray) -> np.ndarray:
    """Unit-normalise a vector. Returns zero vector if norm is near-zero."""
    n = float(np.linalg.norm(v))
    if n < 1e-8:
        return np.zeros_like(v)
    return v / n


def _gram_schmidt(vectors: list[np.ndarray]) -> list[np.ndarray]:
    """Gram-Schmidt orthogonalisation. Returns orthonormal basis."""
    basis: list[np.ndarray] = []
    for v in vectors:
        w = v.copy().astype(np.float64)
        for b in basis:
            w -= float(np.dot(w, b)) * b
        n = float(np.linalg.norm(w))
        if n > 1e-10:
            basis.append((w / n).astype(np.float32))
    return basis


def _auto_layers(num_layers: int, max_points: int = 20) -> list[int]:
    """Auto-select evenly spaced layers, always including 0 and last."""
    if num_layers <= max_points:
        return list(range(num_layers))
    step = max(1, num_layers // max_points)
    layers = list(range(0, num_layers, step))
    if (num_layers - 1) not in layers:
        layers.append(num_layers - 1)
    return layers


# ---------------------------------------------------------------------------
# Token helpers
# ---------------------------------------------------------------------------


def _resolve_token_to_id(tokenizer: Any, token: str) -> int | None:
    """Encode a token string to a single token ID (tries bare + space-prefixed)."""
    for variant in (token, " " + token):
        ids = tokenizer.encode(variant, add_special_tokens=False)
        if ids:
            return ids[0]
    return None


def _get_unembed_vec_np(model: Any, token_id: int) -> np.ndarray | None:
    """Get unembedding vector for a token as numpy float32."""
    from ..residual_tools import _get_unembed_vector

    vec: Any = _get_unembed_vector(model, token_id)
    if vec is None:
        return None
    raw: list[float] = vec.tolist() if hasattr(vec, "tolist") else list(vec)
    return np.array(raw, dtype=np.float32)


def _token_text_at(tokenizer: Any, prompt: str, position: int) -> str:
    """Decode a single token from a prompt at the given position."""
    ids = tokenizer.encode(prompt, add_special_tokens=True)
    idx = position if position >= 0 else len(ids) + position
    idx = max(0, min(idx, len(ids) - 1))
    return tokenizer.decode([ids[idx]])


# ---------------------------------------------------------------------------
# Direction vector extraction
# ---------------------------------------------------------------------------


def _parse_direction_spec(raw: dict) -> tuple[DirectionSpec | None, str | None]:
    """Parse a raw dict into a typed DirectionSpec. Returns (spec, error)."""
    if not isinstance(raw, dict) or "type" not in raw:
        return None, "Each direction must be a dict with 'type'."
    try:
        return DirectionSpec(**raw), None
    except Exception as exc:
        return None, f"Invalid direction spec: {exc}"


def _extract_direction_vector(
    model: Any,
    config: Any,
    tokenizer: Any,
    meta: Any,
    prompt: str,
    layer: int,
    spec: DirectionSpec,
    token_position: int,
    decomp: dict | None,
) -> tuple[str, np.ndarray | None, str | None]:
    """Extract a direction vector from a typed spec.

    Returns (label, vector_np, error_message).
    decomp is the result of _run_decomposition_forward if already computed.
    """
    from ..residual_tools import _extract_position

    if spec.type == DirectionType.TOKEN:
        if spec.value is None:
            return ("", None, "Token direction requires 'value'.")
        tid = _resolve_token_to_id(tokenizer, str(spec.value))
        if tid is None:
            return ("", None, f"Cannot encode token '{spec.value}'.")
        tvec = _get_unembed_vec_np(model, tid)
        if tvec is None:
            return ("", None, f"No unembed vector for token '{spec.value}'.")
        return (f"token:{spec.value}", tvec, None)

    if spec.type == DirectionType.STEERING_VECTOR:
        if spec.value is None:
            return ("", None, "Steering vector requires 'value' (name).")
        reg = SteeringVectorRegistry.get()
        sv_result = reg.fetch(str(spec.value))
        if sv_result is None:
            return ("", None, f"Steering vector '{spec.value}' not found.")
        sv_vec, _ = sv_result
        return (f"sv:{spec.value}", sv_vec.astype(np.float32), None)

    if spec.type == DirectionType.NEURON:
        if spec.value is None:
            return ("", None, "Neuron direction requires 'value' (neuron index).")
        import mlx.core as mx
        from chuk_lazarus.introspection.hooks import ModelHooks

        helper = ModelHooks(model, model_config=config)
        model_layers = helper._get_layers()
        if layer >= len(model_layers):
            return ("", None, f"Layer {layer} out of range.")
        target = model_layers[layer]
        down_weight = target.mlp.down_proj.weight
        if isinstance(down_weight, mx.array):
            col: Any = down_weight[:, int(spec.value)]
            nraw: list[float] = col.tolist() if hasattr(col, "tolist") else list(col)
            return (f"neuron:{spec.value}", np.array(nraw, dtype=np.float32), None)
        return ("", None, f"Cannot extract neuron {spec.value} direction.")

    # Remaining types need decomposition data
    if decomp is None:
        return ("", None, f"Decomposition data not available for type '{spec.type.value}'.")

    if spec.type == DirectionType.RESIDUAL:
        h = decomp["hidden_states"].get(layer)
        if h is None:
            return ("", None, f"No hidden state at layer {layer}.")
        rvec: Any = _extract_position(h, token_position)
        rraw: list[float] = rvec.tolist() if hasattr(rvec, "tolist") else list(rvec)
        return ("residual", np.array(rraw, dtype=np.float32), None)

    if spec.type == DirectionType.FFN_OUTPUT:
        ffn = decomp["ffn_outputs"].get(layer)
        if ffn is None:
            return ("", None, f"No FFN output at layer {layer}.")
        fvec: Any = _extract_position(ffn, token_position)
        fraw: list[float] = fvec.tolist() if hasattr(fvec, "tolist") else list(fvec)
        return ("ffn_output", np.array(fraw, dtype=np.float32), None)

    if spec.type == DirectionType.ATTENTION_OUTPUT:
        attn = decomp["attn_outputs"].get(layer)
        if attn is None:
            return ("", None, f"No attention output at layer {layer}.")
        avec: Any = _extract_position(attn, token_position)
        araw: list[float] = avec.tolist() if hasattr(avec, "tolist") else list(avec)
        return ("attn_output", np.array(araw, dtype=np.float32), None)

    if spec.type == DirectionType.HEAD_OUTPUT:
        if spec.value is None:
            return ("", None, "Head output requires 'value' (head index).")
        import mlx.core as mx
        from chuk_lazarus.introspection.hooks import ModelHooks

        attn_out = decomp["attn_outputs"].get(layer)
        if attn_out is None:
            return ("", None, f"No attention output at layer {layer}.")

        helper = ModelHooks(model, model_config=config)
        model_layers = helper._get_layers()
        target = model_layers[layer]
        hidden_dim = meta.hidden_dim
        num_heads = meta.num_attention_heads
        head_dim = hidden_dim // num_heads
        head_idx = int(spec.value)
        if head_idx < 0 or head_idx >= num_heads:
            return ("", None, f"Head {head_idx} out of [0, {num_heads - 1}].")

        # Use o_proj weight column slice as the head's output direction
        o_proj = target.self_attn.o_proj
        w_slice = o_proj.weight[:, head_idx * head_dim : (head_idx + 1) * head_dim]
        head_dir: Any = mx.mean(w_slice, axis=1)
        hraw: list[float] = head_dir.tolist() if hasattr(head_dir, "tolist") else list(head_dir)
        return (f"head:{spec.value}", np.array(hraw, dtype=np.float32), None)

    return ("", None, f"Unhandled direction type '{spec.type.value}'.")
