manifold_distance.py

"""
Manifold Distance–Based Partitioning of Generated Structures

Classifies generated structures into four categories based on their position
relative to the reference (training) manifold in MLIP embedding space:

    i)  Dense interior   — inside boundary, density ≥ 20th percentile, LOF ≥ -2.0
    ii) Sparse interior  — inside boundary, density < 20th percentile, LOF ≥ -2.0
    iii) Near exterior   — outside boundary, LOF ≥ -2.0
    iv) Far exterior     — LOF < -2.0 (Structural Hallucination)

Input is structure-level descriptors obtained by mean-pooling per-atom embeddings
from any MLIP extractor (e.g. MLIPEmbeddingExtractor from mlip_embedding_extractor.py).

Usage:
    from manifold_distance import ManifoldClassifier
    from mlip_embedding_extractor import MLIPEmbeddingExtractor

    # 1. Extract embeddings
    extractor = MLIPEmbeddingExtractor(model_type="mace", model_path="...", device="cuda")
    ref_result  = extractor.extract_from_csv("mp-20/train.csv")
    gen_result  = extractor.extract_from_csv("generated.csv")

    # 2. Fit classifier on reference set
    classifier = ManifoldClassifier(n_neighbors=10, lof_threshold=-1.5, sparse_percentile=20.0)
    classifier.fit_from_embedding_result(ref_result)

    # 3. Classify generated structures
    results = classifier.classify_from_embedding_result(gen_result)

    # 4. Inspect
    classifier.print_summary(results)
    df = classifier.to_dataframe(results)
"""

import warnings
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

import numpy as np
from sklearn.decomposition import PCA
from sklearn.neighbors import LocalOutlierFactor, NearestNeighbors
from sklearn.preprocessing import StandardScaler
from sklearn.svm import OneClassSVM

try:
    import pandas as pd
    HAS_PANDAS = True
except ImportError:
    HAS_PANDAS = False


# =============================================================================
# Category Definitions
# =============================================================================

CATEGORIES = {
    "dense_interior": {
        "description": "Inside boundary, dense region — well-supported candidate",
        "color": "#2ecc71",  # green
    },
    "sparse_interior": {
        "description": "Inside boundary, sparse region — novel but locally consistent",
        "color": "#3498db",  # blue
    },
    "near_exterior": {
        "description": "Outside boundary, LOF-consistent — exploration candidate",
        "color": "#9b59b6",  # purple
    },
    "far_exterior": {
        "description": "LOF outlier — likely unphysical (Structural Hallucination)",
        "color": "#e74c3c",  # red
    },
}


# =============================================================================
# Data Classes
# =============================================================================

@dataclass
class StructureMetrics:
    """Per-structure manifold metrics."""
    # Identification
    index: int
    material_id: str

    # Metric 1: manifold distance (z-score normalised mean kNN distance)
    manifold_distance: float

    # Metric 2: depth score (0 = centroid, 1 = periphery)
    depth_score: float

    # Metric 3: signed boundary distance (+ inside, – outside)
    boundary_distance: float

    # Metric 4: local density and its percentile vs. reference
    local_density: float
    density_percentile: float

    # Metric 5: LOF score (≈ -1 normal, << -1 outlier)
    lof_score: float

    # Classification
    category: str
    nearest_reference_ids: List[str] = field(default_factory=list)


@dataclass
class ClassificationSummary:
    """Aggregate statistics for one category."""
    category: str
    description: str
    count: int
    percentage: float
    material_ids: List[str]


# =============================================================================
# Utility: mean-pool atom embeddings → structure embeddings
# =============================================================================

def mean_pool(
    atom_embeddings: List[np.ndarray],
    material_ids: Optional[List[str]] = None,
) -> Tuple[np.ndarray, List[str]]:
    """
    Collapse per-atom embeddings to one vector per structure via mean-pooling.

    Args:
        atom_embeddings: List of (n_atoms_i, embed_dim) arrays, one per structure.
                         Accepts the `.embeddings` field of EmbeddingResult directly.
        material_ids:    Optional list of IDs; auto-generated if None.

    Returns:
        structure_embeddings: (n_structures, embed_dim) float32 array
        ids:                  list of material IDs
    """
    pooled = []
    valid_ids = []
    ids = material_ids or [f"struct_{i}" for i in range(len(atom_embeddings))]

    for i, emb in enumerate(atom_embeddings):
        if emb is None or len(emb) == 0:
            warnings.warn(f"Skipping structure {ids[i]}: empty embedding.")
            continue
        pooled.append(np.mean(emb, axis=0))
        valid_ids.append(ids[i])

    return np.array(pooled, dtype=np.float32), valid_ids


# =============================================================================
# Main Classifier
# =============================================================================

class ManifoldClassifier:
    """
    Manifold distance–based classifier for generated crystal structures.

    Fits a reference manifold from training-set embeddings and classifies
    generated structures into four categories based on boundary position,
    local density, and Local Outlier Factor.

    Example:
        >>> clf = ManifoldClassifier(n_neighbors=10)
        >>> clf.fit(reference_embeddings, reference_ids)
        >>> results = clf.classify(generated_embeddings, generated_ids)
        >>> clf.print_summary(results)
    """

    def __init__(
        self,
        n_neighbors: int = 30,
        lof_threshold: float = -2.0,
        sparse_percentile: float = 20.0,
        svm_nu: float = 0.1,
    ):
        """
        Args:
            n_neighbors:      k for kNN density / LOF computation.
                              Default matches ManifoldFishAnalyzer(n_neighbors=30).
            lof_threshold:    LOF scores below this → far_exterior / hallucination.
                              Default matches ManifoldFishAnalyzer(outlier_threshold=-2).
            sparse_percentile: Density percentile below this → sparse region.
            svm_nu:           OneClassSVM nu parameter (expected outlier fraction).
        """
        self.n_neighbors = n_neighbors
        self.lof_threshold = lof_threshold
        self.sparse_percentile = sparse_percentile
        self.svm_nu = svm_nu

        self._is_fitted = False

    # -------------------------------------------------------------------------
    # Fitting (reference manifold)
    # -------------------------------------------------------------------------

    def fit(
        self,
        reference_embeddings: np.ndarray,
        reference_ids: Optional[List[str]] = None,
    ) -> "ManifoldClassifier":
        """
        Fit all reference models.

        Args:
            reference_embeddings: (n_ref, embed_dim) structure-level descriptors.
            reference_ids:        Optional list of reference material IDs.
        """
        ref = np.asarray(reference_embeddings, dtype=np.float32)
        n_ref, dim = ref.shape
        self._reference = ref
        self._dim = dim
        self._reference_ids = reference_ids or [f"ref_{i}" for i in range(n_ref)]

        k = min(self.n_neighbors, n_ref - 1)

        # kNN for distance / density
        self._knn = NearestNeighbors(n_neighbors=k, metric="euclidean")
        self._knn.fit(ref)

        # Reference distance statistics (for z-score normalisation)
        ref_dists, _ = self._knn.kneighbors(ref)
        self._ref_mean_dist = float(np.mean(ref_dists[:, -1]))
        self._ref_std_dist  = float(np.std(ref_dists[:, -1])) + 1e-8

        # Reference density distribution (for percentile ranking)
        ref_densities = self._density_from_distances(ref_dists)
        self._ref_densities_sorted = np.sort(ref_densities)

        # Centroid / depth normalisation
        self._centroid = ref.mean(axis=0)
        self._max_centroid_dist = float(
            np.max(np.linalg.norm(ref - self._centroid, axis=1))
        ) + 1e-8

        # LOF (novelty mode — scores unseen points)
        # random_state is not used by LOF itself (kNN is deterministic),
        # but setting it future-proofs against sklearn version changes.
        self._lof = LocalOutlierFactor(
            n_neighbors=k,
            novelty=True,
            contamination=0.1,
        )
        self._lof.fit(ref)  # fully deterministic given fixed data + n_neighbors

        # Boundary model
        # OneClassSVM works for any dimensionality.
        # For dim ≤ 3 a Delaunay triangulation would suffice, but the typical
        # embedding space (128–512 dims) always uses SVM.
        self._svm_scaler = StandardScaler()
        ref_scaled = self._svm_scaler.fit_transform(ref)
        self._svm = OneClassSVM(kernel="rbf", gamma="auto", nu=self.svm_nu)
        self._svm.fit(ref_scaled)
        ref_decision = self._svm.decision_function(ref_scaled)
        self._svm_decision_std = float(np.std(ref_decision)) + 1e-8

        self._is_fitted = True
        return self

    def fit_from_embedding_result(
        self,
        embedding_result,
        element_filter: Optional[int] = None,
    ) -> "ManifoldClassifier":
        """
        Convenience wrapper: fit directly from an EmbeddingResult object
        (output of MLIPEmbeddingExtractor).

        Args:
            embedding_result: EmbeddingResult with .embeddings and .material_ids fields.
            element_filter:   Optional atomic number; if set, mean-pools only atoms
                              of that element (useful for element-specific analysis).
        """
        atom_embs = embedding_result.embeddings
        ids = embedding_result.material_ids

        if element_filter is not None:
            atomic_nums = embedding_result.atomic_numbers
            filtered = []
            for emb, zs in zip(atom_embs, atomic_nums):
                mask = np.asarray(zs) == element_filter
                filtered.append(emb[mask] if mask.any() else emb)
            atom_embs = filtered

        struct_embs, valid_ids = mean_pool(atom_embs, ids)
        return self.fit(struct_embs, valid_ids)

    # -------------------------------------------------------------------------
    # Classification (generated structures)
    # -------------------------------------------------------------------------

    def classify(
        self,
        generated_embeddings: np.ndarray,
        material_ids: Optional[List[str]] = None,
    ) -> List[StructureMetrics]:
        """
        Classify generated structures against the fitted reference manifold.

        Args:
            generated_embeddings: (n_gen, embed_dim) structure-level descriptors.
            material_ids:         Optional list of generated material IDs.

        Returns:
            List of StructureMetrics, one per generated structure.
        """
        if not self._is_fitted:
            raise RuntimeError("Call fit() or fit_from_embedding_result() first.")

        gen = np.asarray(generated_embeddings, dtype=np.float32)
        n_gen = len(gen)
        ids = material_ids or [f"gen_{i}" for i in range(n_gen)]

        # --- Metric 1: manifold distance (z-score) ---
        dists, neighbor_idx = self._knn.kneighbors(gen)
        mean_dists = dists.mean(axis=1)
        manifold_distance = (mean_dists - self._ref_mean_dist) / self._ref_std_dist

        # --- Metric 2: depth score ---
        centroid_dists = np.linalg.norm(gen - self._centroid, axis=1)
        depth_score = centroid_dists / self._max_centroid_dist

        # --- Metric 3: boundary distance (SVM decision function, normalised) ---
        gen_scaled = self._svm_scaler.transform(gen)
        svm_decision = self._svm.decision_function(gen_scaled)
        boundary_distance = svm_decision / self._svm_decision_std

        # --- Metric 4: local density and density percentile ---
        local_density = self._density_from_distances(dists)
        positions = np.searchsorted(self._ref_densities_sorted, local_density, side="right")
        density_percentile = 100.0 * positions / len(self._ref_densities_sorted)

        # --- Metric 5: LOF score ---
        lof_score = self._lof.score_samples(gen)

        # --- Classify each structure ---
        results = []
        for i in range(n_gen):
            category = self._classify_one(
                boundary_distance=float(boundary_distance[i]),
                density_percentile=float(density_percentile[i]),
                lof_score=float(lof_score[i]),
            )
            nearest_ref_ids = [self._reference_ids[j] for j in neighbor_idx[i]]
            results.append(StructureMetrics(
                index=i,
                material_id=ids[i],
                manifold_distance=float(manifold_distance[i]),
                depth_score=float(depth_score[i]),
                boundary_distance=float(boundary_distance[i]),
                local_density=float(local_density[i]),
                density_percentile=float(density_percentile[i]),
                lof_score=float(lof_score[i]),
                category=category,
                nearest_reference_ids=nearest_ref_ids,
            ))

        return results

    def classify_from_embedding_result(
        self,
        embedding_result,
        element_filter: Optional[int] = None,
    ) -> List[StructureMetrics]:
        """
        Convenience wrapper: classify directly from an EmbeddingResult object.

        Args:
            embedding_result: EmbeddingResult with .embeddings and .material_ids.
            element_filter:   Optional atomic number for element-specific pooling.
        """
        atom_embs = embedding_result.embeddings
        ids = embedding_result.material_ids

        if element_filter is not None:
            atomic_nums = embedding_result.atomic_numbers
            filtered = []
            for emb, zs in zip(atom_embs, atomic_nums):
                mask = np.asarray(zs) == element_filter
                filtered.append(emb[mask] if mask.any() else emb)
            atom_embs = filtered

        struct_embs, valid_ids = mean_pool(atom_embs, ids)
        return self.classify(struct_embs, valid_ids)

    # -------------------------------------------------------------------------
    # Internal helpers
    # -------------------------------------------------------------------------

    @staticmethod
    def _density_from_distances(dists: np.ndarray) -> np.ndarray:
        """Reciprocal of mean kNN distance as local density estimate."""
        return 1.0 / (dists.mean(axis=1) + 1e-8)

    def _classify_one(
        self,
        boundary_distance: float,
        density_percentile: float,
        lof_score: float,
    ) -> str:
        """
        Decision tree for four-category classification.

        Priority order:
            1. LOF < lof_threshold  → far_exterior (Structural Hallucination)
            2. boundary_distance < 0 → near_exterior
            3. density_percentile < sparse_percentile → sparse_interior
            4. otherwise → dense_interior
        """
        if lof_score < self.lof_threshold:
            return "far_exterior"

        if boundary_distance < 0:
            return "near_exterior"

        if density_percentile < self.sparse_percentile:
            return "sparse_interior"

        return "dense_interior"

    # -------------------------------------------------------------------------
    # Reporting
    # -------------------------------------------------------------------------

    def summarise(self, results: List[StructureMetrics]) -> Dict[str, ClassificationSummary]:
        """Return per-category counts and material IDs."""
        from collections import defaultdict
        buckets: Dict[str, List[str]] = defaultdict(list)
        for r in results:
            buckets[r.category].append(r.material_id)

        n_total = len(results)
        summary = {}
        for cat in CATEGORIES:
            ids = buckets.get(cat, [])
            summary[cat] = ClassificationSummary(
                category=cat,
                description=CATEGORIES[cat]["description"],
                count=len(ids),
                percentage=100.0 * len(ids) / n_total if n_total else 0.0,
                material_ids=ids,
            )
        return summary

    def print_summary(self, results: List[StructureMetrics]) -> None:
        """Print a human-readable summary table."""
        summary = self.summarise(results)
        n_total = len(results)
        print(f"\nManifold classification — {n_total} structures")
        print(f"{'Category':<20} {'Count':>7} {'%':>7}  Description")
        print("-" * 72)
        for cat, s in summary.items():
            print(f"{cat:<20} {s.count:>7} {s.percentage:>6.1f}%  {s.description}")
        print()

    def to_dataframe(self, results: List[StructureMetrics]):
        """Convert results to a pandas DataFrame (requires pandas)."""
        if not HAS_PANDAS:
            raise ImportError("pandas is required for to_dataframe(). pip install pandas")
        rows = []
        for r in results:
            rows.append({
                "material_id":        r.material_id,
                "category":           r.category,
                "manifold_distance":  r.manifold_distance,
                "depth_score":        r.depth_score,
                "boundary_distance":  r.boundary_distance,
                "local_density":      r.local_density,
                "density_percentile": r.density_percentile,
                "lof_score":          r.lof_score,
            })
        return pd.DataFrame(rows)

    def get_ids_by_category(
        self,
        results: List[StructureMetrics],
        category: str,
    ) -> List[str]:
        """Return material IDs for a given category."""
        return [r.material_id for r in results if r.category == category]


# =============================================================================
# Convenience functions
# =============================================================================

def classify_structures(
    reference_embeddings: np.ndarray,
    generated_embeddings: np.ndarray,
    reference_ids: Optional[List[str]] = None,
    generated_ids: Optional[List[str]] = None,
    n_neighbors: int = 30,
    lof_threshold: float = -2.0,
    sparse_percentile: float = 20.0,
) -> List[StructureMetrics]:
    """
    One-shot classification of generated structures against a reference set.

    Args:
        reference_embeddings: (n_ref, embed_dim) structure-level descriptors.
        generated_embeddings: (n_gen, embed_dim) structure-level descriptors.
        reference_ids:        Optional reference material IDs.
        generated_ids:        Optional generated material IDs.
        n_neighbors:          k for kNN / LOF.
        lof_threshold:        LOF cut-off for hallucination detection.
        sparse_percentile:    Density percentile threshold for sparse regions.

    Returns:
        List of StructureMetrics.
    """
    clf = ManifoldClassifier(
        n_neighbors=n_neighbors,
        lof_threshold=lof_threshold,
        sparse_percentile=sparse_percentile,
    )
    clf.fit(reference_embeddings, reference_ids)
    return clf.classify(generated_embeddings, generated_ids)