"""Activation space geometry types derived from empirical clustering.

K-means clustering on 208 benchmarks x 16 zwiad metrics identified
5 natural geometry types (best silhouette at k=5). Each type has a
distinct metric signature that determines which steering methods are
most effective.

Classifier maps zwiad metrics -> GeometryType using the discriminating
metrics identified by PCA (PC1: linearity/coherence/alignment,
PC2: variance concentration/effective dims, PC3: multi-dir/concepts).
"""
from __future__ import annotations

import json
from enum import Enum
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple


class GeometryType(Enum):
    """Five empirical activation space geometry types.

    Derived from k=5 k-means clustering on 208 benchmarks
    using 16 zwiad geometry metrics.
    """
    LINEAR_FLAT = "linear_flat"
    LINEAR_MULTI = "linear_multi"
    CURVED_DIFFUSE = "curved_diffuse"
    UNSTABLE_CONCENTRATED = "unstable_concentrated"
    CONCENTRATED_FEW = "concentrated_few"


# Centroid signatures from k=5 clustering (key metrics only)
_CENTROIDS = {
    GeometryType.LINEAR_FLAT: {
        "linear": 0.980, "stability": 0.975,
        "n_concepts": 5, "var_pc1": 0.155,
        "curvature": 0.044, "coherence": 0.85,
    },
    GeometryType.LINEAR_MULTI: {
        "linear": 0.828, "stability": 0.968,
        "n_concepts": 8, "var_pc1": 0.148,
        "curvature": 0.116, "coherence": 0.72,
    },
    GeometryType.CURVED_DIFFUSE: {
        "linear": 0.565, "stability": 0.888,
        "n_concepts": 13, "var_pc1": 0.128,
        "curvature": 0.206, "coherence": 0.55,
    },
    GeometryType.UNSTABLE_CONCENTRATED: {
        "linear": 0.578, "stability": 0.844,
        "n_concepts": 8, "var_pc1": 0.344,
        "curvature": 0.204, "coherence": 0.58,
    },
    GeometryType.CONCENTRATED_FEW: {
        "linear": 0.751, "stability": 0.967,
        "n_concepts": 3, "var_pc1": 0.369,
        "curvature": 0.037, "coherence": 0.78,
    },
}

# Metric extraction keys (same as cluster_zwiad.py)
_METRIC_ALIASES = {
    "linear": ["linear_probe_accuracy", "linear_probe"],
    "stability": ["direction_stability_score",
                   "direction_stability"],
    "n_concepts": ["n_concepts"],
    "var_pc1": ["manifold_variance_pc1"],
    "curvature": ["manifold_curvature",
                   "manifold_curvature_proxy"],
    "coherence": ["concept_coherence"],
}


def _extract(metrics: Dict[str, Any], key: str) -> Optional[float]:
    for alias in _METRIC_ALIASES[key]:
        if alias in metrics and metrics[alias] is not None:
            return float(metrics[alias])
    return None


def classify_geometry(
    metrics: Dict[str, Any],
) -> Tuple[GeometryType, float]:
    """Classify activation space geometry from zwiad metrics.

    Uses weighted Euclidean distance to cluster centroids.
    Returns (geometry_type, confidence) where confidence is
    the margin between best and second-best match.
    """
    vals = {}
    for key in _METRIC_ALIASES:
        v = _extract(metrics, key)
        if v is None:
            return GeometryType.LINEAR_FLAT, 0.0
        vals[key] = v

    # Normalization ranges (from 208-benchmark dataset)
    _ranges = {
        "linear": (0.3, 1.0), "stability": (0.5, 1.0),
        "n_concepts": (1, 25), "var_pc1": (0.05, 0.6),
        "curvature": (0.0, 0.5), "coherence": (0.2, 1.0),
    }
    # Importance weights from PCA loadings
    _weights = {
        "linear": 3.0, "stability": 1.5,
        "n_concepts": 1.0, "var_pc1": 2.0,
        "curvature": 2.0, "coherence": 1.5,
    }

    distances = {}
    for gtype, centroid in _CENTROIDS.items():
        dist = 0.0
        for key in centroid:
            lo, hi = _ranges[key]
            span = hi - lo if hi > lo else 1.0
            norm_val = (vals[key] - lo) / span
            norm_cen = (centroid[key] - lo) / span
            dist += _weights[key] * (norm_val - norm_cen) ** 2
        distances[gtype] = dist ** 0.5

    ranked = sorted(distances.items(), key=lambda x: x[1])
    best_type, best_dist = ranked[0]
    second_dist = ranked[1][1]
    if best_dist > 0:
        confidence = min(
            (second_dist - best_dist) / best_dist, 1.0)
    else:
        confidence = 1.0
    return best_type, round(confidence, 3)


def select_representative_benchmarks(
    zwiad_dir: str, model_slug: str,
    per_type: int = 2,
) -> Dict[GeometryType, List[str]]:
    """Select representative benchmarks covering all 5 types.

    For each geometry type, picks the benchmarks closest to
    the centroid (most representative of that type).
    """
    zwiad_path = Path(zwiad_dir)
    prefix = model_slug.replace("/", "_")

    # Classify all benchmarks
    typed: Dict[GeometryType, List[Tuple[str, float]]] = {
        t: [] for t in GeometryType}
    for f in sorted(zwiad_path.glob(f"{prefix}__*.json")):
        bench = f.stem.replace(f"{prefix}__", "")
        data = json.loads(f.read_text())
        m = data.get("metrics", data)
        gtype, conf = classify_geometry(m)
        typed[gtype].append((bench, conf))

    # Pick top per_type per type (highest confidence = most
    # representative)
    selected = {}
    for gtype in GeometryType:
        candidates = typed[gtype]
        candidates.sort(key=lambda x: x[1], reverse=True)
        selected[gtype] = [
            b for b, _ in candidates[:per_type]]
    return selected


@dataclass
class GeometryProfile:
    """Full geometry characterization of a benchmark."""
    geometry_type: GeometryType
    confidence: float
    metrics: Dict[str, float]

    def to_dict(self) -> dict:
        return {
            "geometry_type": self.geometry_type.value,
            "confidence": self.confidence,
            "metrics": self.metrics,
        }

    @classmethod
    def from_dict(cls, d: dict) -> GeometryProfile:
        return cls(
            geometry_type=GeometryType(d["geometry_type"]),
            confidence=d["confidence"],
            metrics=d["metrics"],
        )


def profile_benchmark(
    metrics: Dict[str, Any],
) -> GeometryProfile:
    """Build full geometry profile from zwiad metrics."""
    gtype, conf = classify_geometry(metrics)
    key_metrics = {}
    for key in _METRIC_ALIASES:
        v = _extract(metrics, key)
        if v is not None:
            key_metrics[key] = v
    return GeometryProfile(
        geometry_type=gtype,
        confidence=conf,
        metrics=key_metrics,
    )
