import pickle
import numpy as np
from geood.metrics import intrinsic_dim, mahalanobis_distance
from geood.result import DetectionResult


class Detector:
    def __init__(self):
        self._centroid = None
        self._cov = None
        self._reference_dim = None
        self.layer_used = None
        self._threshold = None
        self._model_ref = None

    @classmethod
    def _calibrate_from_hidden(cls, hidden_states, layer_indices):
        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._centroid = ref_vectors.mean(axis=0)
        detector._cov = np.cov(ref_vectors, rowvar=False)
        detector._reference_dim = float(best_dim)

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

        return detector

    def calibrate(self, model, ref_texts, tokenizer=None, layer="auto"):
        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

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

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

        calibrated = self._calibrate_from_hidden(hidden, layer_indices)
        self.__dict__.update(calibrated.__dict__)

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

        return self

    def detect(self, input_text, model=None, tokenizer=None):
        from geood.model_loader import resolve_model
        from geood.extraction import extract_hidden_states

        if 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."
                )
            model = self._model_ref

        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]
        )
        results = [
            self._detect_from_vector(v) for v in hidden[self.layer_used]
        ]

        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):
        maha = mahalanobis_distance(vector, self._centroid, self._cov)
        score = min(maha / max(self._threshold, 1e-8), 1.0)
        is_ood = maha > self._threshold

        return DetectionResult(
            is_ood=is_ood,
            score=float(score),
            mahalanobis=float(maha),
            intrinsic_dim=None,  # single-sample: dim not computable
            reference_dim=self._reference_dim,
            layer=self.layer_used,
        )

    def save(self, path):
        state = {
            "centroid": self._centroid,
            "cov": self._cov,
            "reference_dim": self._reference_dim,
            "layer_used": self.layer_used,
            "threshold": self._threshold,
            "model_ref": self._model_ref,
        }
        with open(path, "wb") as f:
            pickle.dump(state, f)

    @classmethod
    def load(cls, path):
        with open(path, "rb") as f:
            state = pickle.load(f)
        detector = cls()
        detector._centroid = state["centroid"]
        detector._cov = state["cov"]
        detector._reference_dim = state["reference_dim"]
        detector.layer_used = state["layer_used"]
        detector._threshold = state["threshold"]
        detector._model_ref = state.get("model_ref")
        return detector
