"""Core detector: calibrate once, detect many times."""

from __future__ import annotations

import json
import os
from typing import Union

import numpy as np

from geood.metrics import intrinsic_dim, mahalanobis_distance
from geood.result import DetectionResult

_NPZ_EXT = ".npz"


class Detector:
    """Geometric OOD detector calibrated on in-distribution hidden states."""

    def __init__(self) -> None:
        self._centroid: np.ndarray | None = None  # in PCA space
        self._cov: np.ndarray | None = None       # in PCA space
        self._pca_mean: np.ndarray | None = None   # original space mean
        self._pca_components: np.ndarray | None = None  # projection matrix
        self._reference_dim: float | None = None
        self.layer_used: int | None = None
        self._threshold: float | None = None
        self._model_ref: str | None = None
        self._cached_model: object | None = None
        self._cached_tokenizer: object | None = None

    def __repr__(self) -> str:
        if self.layer_used is None:
            return "Detector(uncalibrated)"
        return (
            f"Detector(layer={self.layer_used}, "
            f"ref_dim={self._reference_dim:.1f}, "
            f"threshold={self._threshold:.1f})"
        )

    @classmethod
    def calibrate_from_vectors(
        cls,
        hidden_states: dict[int, list[np.ndarray]],
        layer_indices: list[int],
    ) -> Detector:
        """Calibrate from pre-extracted hidden state vectors.

        Useful for testing or when you have already extracted hidden
        states yourself.

        Args:
            hidden_states: ``{layer_index: [vector_per_sample, ...]}``.
            layer_indices: Which layers to consider for auto-selection.

        Returns:
            Calibrated detector.
        """
        detector = cls()
        best_layer = None
        best_dim = -1

        for idx in layer_indices:
            vectors = np.array(hidden_states[idx])
            dim = intrinsic_dim(vectors)
            if dim > best_dim:
                best_dim = dim
                best_layer = idx

        detector.layer_used = best_layer
        ref_vectors = np.array(hidden_states[best_layer])
        detector._reference_dim = float(best_dim)

        # PCA projection to avoid singular covariance (n < d)
        mean = ref_vectors.mean(axis=0)
        centered = ref_vectors - mean
        U, S, Vt = np.linalg.svd(centered, full_matrices=False)
        # Keep components explaining 95% variance
        var_explained = np.cumsum(S ** 2) / np.sum(S ** 2)
        n_components = int(np.searchsorted(var_explained, 0.95) + 1)
        detector._pca_mean = mean
        detector._pca_components = Vt[:n_components]  # (k, d)

        # Project into PCA subspace
        projected = centered @ detector._pca_components.T  # (n, k)
        detector._centroid = projected.mean(axis=0)
        detector._cov = np.cov(projected, rowvar=False)

        maha_dists = np.array([
            mahalanobis_distance(v, detector._centroid, detector._cov)
            for v in projected
        ])
        detector._threshold = float(maha_dists.mean() + 2 * maha_dists.std())

        return detector

    # Keep old name for backwards compat
    _calibrate_from_hidden = calibrate_from_vectors

    def calibrate(
        self,
        model: Union[str, object],
        ref_texts: list[str],
        tokenizer: object | None = None,
        layer: Union[str, int] = "auto",
    ) -> Detector:
        """Calibrate the detector from reference texts.

        Runs a forward pass on *ref_texts*, extracts hidden states at
        candidate layers, and selects the layer with highest intrinsic
        dimensionality.  Computes centroid, covariance, and Mahalanobis
        threshold for later detection.

        The model is cached internally for subsequent ``detect()`` calls.

        Args:
            model: HuggingFace model name or loaded model object.
            ref_texts: In-distribution reference texts (recommended: 50+).
            tokenizer: Required when *model* is an object.
            layer: ``"auto"`` or an integer layer index.

        Returns:
            self (for chaining).
        """
        from geood.model_loader import resolve_model
        from geood.extraction import (
            extract_hidden_states, get_layer_count, get_candidate_layers,
        )

        resolved_model, resolved_tokenizer, should_cleanup = resolve_model(
            model, tokenizer,
        )
        self._model_ref = model if isinstance(model, str) else None
        # Cache model so detect() doesn't reload every time
        self._cached_model = resolved_model
        self._cached_tokenizer = resolved_tokenizer

        n_layers = get_layer_count(resolved_model)
        layer_indices = (
            get_candidate_layers(n_layers) if layer == "auto" else [layer]
        )

        hidden = extract_hidden_states(
            resolved_model, resolved_tokenizer, ref_texts, layer_indices,
        )

        calibrated = self.calibrate_from_vectors(hidden, layer_indices)
        # Preserve cached model refs
        cached_model = self._cached_model
        cached_tokenizer = self._cached_tokenizer
        model_ref = self._model_ref
        self.__dict__.update(calibrated.__dict__)
        self._cached_model = cached_model
        self._cached_tokenizer = cached_tokenizer
        self._model_ref = model_ref

        return self

    def detect(
        self,
        input_text: Union[str, list[str]],
        model: Union[str, object, None] = None,
        tokenizer: object | None = None,
    ) -> Union[DetectionResult, list[DetectionResult]]:
        """Score one or more texts against the calibrated reference.

        If the detector was calibrated with a model, it reuses the cached
        model.  Pass *model* explicitly to override.

        Args:
            input_text: A single string or a list of strings.
            model: Override model (str or object).  If *None*, reuses
                the cached model from calibration.
            tokenizer: Required when *model* is an object.

        Returns:
            A single :class:`DetectionResult` when *input_text* is a
            string, or a list of results when it is a list.
        """
        from geood.model_loader import resolve_model
        from geood.extraction import extract_hidden_states

        # Reuse cached model when possible
        if model is None and self._cached_model is not None:
            resolved_model = self._cached_model
            resolved_tokenizer = self._cached_tokenizer
            should_cleanup = False
        elif model is None:
            if self._model_ref is None:
                raise ValueError(
                    "No model available. Either pass model= to detect() "
                    "or calibrate with a model name string."
                )
            resolved_model, resolved_tokenizer, should_cleanup = resolve_model(
                self._model_ref, tokenizer,
            )
            # Cache for next call
            self._cached_model = resolved_model
            self._cached_tokenizer = resolved_tokenizer
        else:
            resolved_model, resolved_tokenizer, should_cleanup = resolve_model(
                model, tokenizer,
            )

        texts = [input_text] if isinstance(input_text, str) else input_text
        hidden = extract_hidden_states(
            resolved_model, resolved_tokenizer, texts, [self.layer_used],
        )

        vectors = hidden[self.layer_used]

        # Compute intrinsic dim for batch (>= 2 samples)
        batch_dim = None
        if len(vectors) >= 2:
            batch_dim = intrinsic_dim(np.array(vectors))

        results = [
            self._detect_from_vector(v, batch_dim) for v in vectors
        ]

        if should_cleanup:
            del resolved_model
            import torch
            torch.cuda.empty_cache()

        return results[0] if isinstance(input_text, str) else results

    def _detect_from_vector(
        self, vector: np.ndarray, batch_dim: int | None = None,
    ) -> DetectionResult:
        # Project into the same PCA subspace used during calibration
        projected = (vector - self._pca_mean) @ self._pca_components.T
        maha = mahalanobis_distance(projected, self._centroid, self._cov)
        raw_score = maha / max(self._threshold, 1e-8)
        score = max(0.0, min(raw_score, 1.0))
        is_ood = maha > self._threshold

        return DetectionResult(
            is_ood=is_ood,
            score=float(score),
            mahalanobis=float(maha),
            intrinsic_dim=batch_dim,
            reference_dim=self._reference_dim,
            layer=self.layer_used,
        )

    def save(self, path: str) -> None:
        """Serialize the detector to disk.

        Uses numpy ``.npz`` for arrays and JSON-encoded metadata.
        No pickle — safe to load from untrusted sources.

        The file is always saved with a ``.npz`` extension.  If *path*
        does not end in ``.npz``, the extension is appended.
        """
        if not path.endswith(_NPZ_EXT):
            path = path + _NPZ_EXT
        meta_str = json.dumps({
            "reference_dim": self._reference_dim,
            "layer_used": self.layer_used,
            "threshold": self._threshold,
            "model_ref": self._model_ref,
        })
        meta_arr = np.frombuffer(meta_str.encode(), dtype=np.uint8)
        np.savez_compressed(
            path,
            centroid=self._centroid,
            cov=self._cov,
            pca_mean=self._pca_mean,
            pca_components=self._pca_components,
            meta=meta_arr,
        )

    @classmethod
    def load(cls, path: str) -> Detector:
        """Load a detector from a saved ``.npz`` file.

        Safe to use with untrusted files — no pickle deserialization.
        """
        if not path.endswith(_NPZ_EXT) and not os.path.exists(path):
            path = path + _NPZ_EXT
        data = np.load(path, allow_pickle=False)
        meta = json.loads(data["meta"].tobytes().decode())
        detector = cls()
        detector._centroid = data["centroid"]
        detector._cov = data["cov"]
        detector._pca_mean = data["pca_mean"]
        detector._pca_components = data["pca_components"]
        detector._reference_dim = meta["reference_dim"]
        detector.layer_used = meta["layer_used"]
        detector._threshold = meta["threshold"]
        detector._model_ref = meta.get("model_ref")
        return detector
