from __future__ import annotations

MODEL_NOT_FITTED_ERROR = "Model not fitted. Call fit() first."


"""PatchCore-lite-map: patch memory bank + kNN distance anomaly map.

This is a pragmatic, industrial-friendly baseline inspired by PatchCore:
- Extract conv patch embeddings (injectable embedder).
- Build a memory bank from normal patches (optional coreset subsampling).
- Score each patch by kNN distance to the bank.
- Upsample patch scores to an anomaly map and aggregate to an image score.

Unlike the full PatchCore implementation, this version is intentionally small:
- no multi-layer fusion
- no projection fitting
- no fancy reweighting

Default behavior is **offline-safe**:
- the default embedder uses torchvision backbones with `pretrained=False`
  unless explicitly enabled by the caller.
"""

import math
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Optional, Tuple, Union, cast

import numpy as np
from numpy.typing import NDArray

from pyimgano.features.torchvision_conv_patch_embedder import TorchvisionConvPatchEmbedder

from ._legacy_x import MISSING, resolve_legacy_x_keyword
from .knn_index import KNNIndex, build_knn_index
from .patchknn_core import AggregationMethod, aggregate_patch_scores, reshape_patch_scores
from .registry import register_model


@dataclass
class _EmbeddedImage:
    patch_embeddings: NDArray
    grid_shape: Tuple[int, int]
    original_size: Tuple[int, int]


def _embedder_to_payload(embedder: TorchvisionConvPatchEmbedder) -> dict[str, object]:
    payload: dict[str, object] = {
        "type": "torchvision_conv_patch_embedder",
        "config": {
            "backbone": str(embedder.backbone),
            "node": str(embedder.node),
            "pretrained": bool(embedder.pretrained),
            "device": str(embedder.device),
            "input_color": str(embedder.input_color),
            "image_size": int(embedder.image_size),
            "normalize": bool(embedder.normalize),
            "eps": float(embedder.eps),
        },
    }

    model = getattr(embedder, "_model", None)
    state_dict = getattr(model, "state_dict", None)
    if model is not None and callable(state_dict):
        raw_state = state_dict()
        normalized_state: dict[str, object] = {}
        for key, value in dict(raw_state).items():
            detach = getattr(value, "detach", None)
            cpu = getattr(value, "cpu", None)
            if callable(detach) and callable(cpu):
                try:
                    normalized_state[str(key)] = detach().cpu()
                    continue
                except Exception:
                    pass
            normalized_state[str(key)] = value
        payload["model_state_dict"] = normalized_state

    return payload


def _embedder_from_payload(payload: dict[str, object]) -> TorchvisionConvPatchEmbedder:
    if str(payload.get("type", "")) != "torchvision_conv_patch_embedder":
        raise ValueError("Unsupported patch embedder checkpoint payload.")

    config = dict(cast(dict[str, object], payload.get("config", {})))
    embedder = TorchvisionConvPatchEmbedder(
        backbone=str(config.get("backbone", "resnet18")),
        node=str(config.get("node", "layer3")),
        pretrained=bool(config.get("pretrained", False)),
        device=str(config.get("device", "cpu")),
        input_color=cast(str, config.get("input_color", "rgb")),
        image_size=int(config.get("image_size", 224)),
        normalize=bool(config.get("normalize", True)),
        eps=float(config.get("eps", 1e-12)),
    )

    model_state = payload.get("model_state_dict", None)
    if isinstance(model_state, dict):
        embedder._ensure_ready()
        model = getattr(embedder, "_model", None)
        load_state_dict = getattr(model, "load_state_dict", None)
        if callable(load_state_dict):
            load_state_dict(dict(model_state), strict=False)
    return embedder


@register_model(
    "vision_patchcore_lite_map",
    tags=("vision", "deep", "patchcore", "memory_bank", "lite", "patchknn", "numpy", "pixel_map"),
    metadata={
        "description": "PatchCore-lite anomaly map: conv patch embeddings + memory bank kNN distance",
        "paper": "PatchCore (Roth et al., CVPR 2022) - lite baseline",
        "year": 2022,
    },
)
class VisionPatchCoreLiteMap:
    """PatchCore-lite anomaly map (memory-bank kNN on patch embeddings).

    The embedder is injectable so unit tests (and offline industrial deployments)
    can run without any implicit weight downloads.
    """

    def __init__(
        self,
        *,
        embedder: Optional[TorchvisionConvPatchEmbedder] = None,
        contamination: float = 0.1,
        pretrained: bool = False,
        # kNN / memory bank
        knn_backend: str = "sklearn",
        metric: str = "euclidean",
        n_neighbors: int = 1,
        coreset_sampling_ratio: float = 1.0,
        random_seed: int = 0,
        # aggregation
        aggregation_method: AggregationMethod = "topk_mean",
        aggregation_topk: float = 0.01,
        # default embedder knobs
        backbone: str = "resnet18",
        node: str = "layer3",
        device: str = "cpu",
        image_size: int = 224,
        normalize_embeddings: bool = True,
        eps: float = 1e-12,
    ) -> None:
        if embedder is None:
            embedder = TorchvisionConvPatchEmbedder(
                backbone=str(backbone),
                node=str(node),
                pretrained=bool(pretrained),
                device=str(device),
                image_size=int(image_size),
                normalize=bool(normalize_embeddings),
                eps=float(eps),
            )

        self.embedder = embedder
        self.contamination = float(contamination)
        if not (0.0 < self.contamination < 0.5):
            raise ValueError(f"contamination must be in (0, 0.5). Got {self.contamination}.")

        self.pretrained = bool(pretrained)
        self.knn_backend = str(knn_backend)
        self.metric = str(metric)
        self.n_neighbors = int(n_neighbors)
        if self.n_neighbors < 1:
            raise ValueError(f"n_neighbors must be >= 1. Got {self.n_neighbors}.")

        self.coreset_sampling_ratio = float(coreset_sampling_ratio)
        if not (0.0 < self.coreset_sampling_ratio <= 1.0):
            raise ValueError(
                f"coreset_sampling_ratio must be in (0, 1]. Got {self.coreset_sampling_ratio}."
            )
        self.random_seed = int(random_seed)

        self.aggregation_method = aggregation_method
        self.aggregation_topk = float(aggregation_topk)

        self.decision_scores_: Optional[NDArray] = None
        self.threshold_: Optional[float] = None

        self._memory_bank: Optional[NDArray] = None
        self._knn_index: Optional[KNNIndex] = None
        self._n_neighbors_fit: Optional[int] = None

    @property
    def memory_bank_size_(self) -> int:
        if self._memory_bank is None:
            raise RuntimeError(MODEL_NOT_FITTED_ERROR)
        return int(self._memory_bank.shape[0])

    def _embed(self, image: Union[str, np.ndarray]) -> _EmbeddedImage:
        patch_embeddings, grid_shape, original_size = self.embedder.embed(image)
        patch_embeddings_np = np.asarray(patch_embeddings, dtype=np.float32)
        if patch_embeddings_np.ndim != 2:
            raise ValueError(f"Expected 2D patch embeddings, got shape {patch_embeddings_np.shape}")

        grid_h, grid_w = int(grid_shape[0]), int(grid_shape[1])
        if patch_embeddings_np.shape[0] != grid_h * grid_w:
            raise ValueError(
                "Patch embedding count does not match grid shape. "
                f"Got {patch_embeddings_np.shape[0]} patches for grid {grid_h}x{grid_w}."
            )

        original_h, original_w = int(original_size[0]), int(original_size[1])
        if original_h <= 0 or original_w <= 0:
            raise ValueError(f"Invalid original_size: {original_size}")

        return _EmbeddedImage(
            patch_embeddings=patch_embeddings_np,
            grid_shape=(grid_h, grid_w),
            original_size=(original_h, original_w),
        )

    def fit(self, x: object = MISSING, y=None, **kwargs: object):
        del y
        items = list(
            cast(
                Iterable[Union[str, np.ndarray]],
                resolve_legacy_x_keyword(x, kwargs, method_name="fit"),
            )
        )
        if not items:
            raise ValueError("X must contain at least one training image.")

        embedded_train = [self._embed(item) for item in items]
        memory_bank = np.concatenate([e.patch_embeddings for e in embedded_train], axis=0)

        if self.coreset_sampling_ratio < 1.0:
            rng = np.random.default_rng(self.random_seed)
            n_total = int(memory_bank.shape[0])
            n_keep = max(1, int(math.ceil(self.coreset_sampling_ratio * n_total)))
            n_keep = min(n_keep, n_total)
            keep_idx = rng.choice(n_total, size=n_keep, replace=False)
            memory_bank = memory_bank[keep_idx]

        self._memory_bank = memory_bank

        effective_k = min(max(1, self.n_neighbors), int(memory_bank.shape[0]))
        self._n_neighbors_fit = effective_k
        self._knn_index = build_knn_index(
            backend=self.knn_backend,
            n_neighbors=effective_k,
            metric=str(self.metric),
        )
        self._knn_index.fit(memory_bank)

        self.decision_scores_ = self.decision_function(items)
        self.threshold_ = float(np.quantile(self.decision_scores_, 1.0 - self.contamination))
        return self

    def _patch_scores(self, embedded: _EmbeddedImage) -> NDArray:
        if self._knn_index is None:
            raise RuntimeError(MODEL_NOT_FITTED_ERROR)
        if self._n_neighbors_fit is None:
            raise RuntimeError("Internal error: missing fitted neighbor count.")

        distances, _indices = self._knn_index.kneighbors(
            embedded.patch_embeddings,
            n_neighbors=self._n_neighbors_fit,
        )
        distances_np = np.asarray(distances, dtype=np.float32)
        if distances_np.ndim != 2:
            raise RuntimeError(f"Expected 2D kNN distances, got shape {distances_np.shape}")

        # PatchCore-style: distance to closest normal patch (or among k neighbors).
        return distances_np.min(axis=1)

    def decision_function(self, x: object = MISSING, **kwargs: object) -> NDArray:
        items = list(
            cast(
                Iterable[Union[str, np.ndarray]],
                resolve_legacy_x_keyword(x, kwargs, method_name="decision_function"),
            )
        )
        scores = np.zeros(len(items), dtype=np.float64)
        for i, item in enumerate(items):
            embedded = self._embed(item)
            patch_scores = self._patch_scores(embedded)
            scores[i] = aggregate_patch_scores(
                patch_scores,
                method=self.aggregation_method,
                topk=self.aggregation_topk,
            )
        return scores

    def predict(self, x: object = MISSING, **kwargs: object) -> NDArray:
        if self.threshold_ is None:
            raise RuntimeError(MODEL_NOT_FITTED_ERROR)
        scores = self.decision_function(
            cast(
                Iterable[Union[str, np.ndarray]],
                resolve_legacy_x_keyword(x, kwargs, method_name="predict"),
            )
        )
        return (scores > self.threshold_).astype(np.int64)

    def get_anomaly_map(self, image: Union[str, np.ndarray]) -> NDArray:
        embedded = self._embed(image)
        patch_scores = self._patch_scores(embedded)
        patch_grid = reshape_patch_scores(
            patch_scores,
            grid_h=embedded.grid_shape[0],
            grid_w=embedded.grid_shape[1],
        )

        try:
            import cv2  # type: ignore
        except Exception as exc:  # pragma: no cover
            raise ImportError(
                "opencv-python is required to upsample anomaly maps.\n"
                "Install it via:\n  pip install 'opencv-python'\n"
                f"Original error: {exc}"
            ) from exc

        original_h, original_w = embedded.original_size
        upsampled = cv2.resize(
            np.asarray(patch_grid, dtype=np.float32),
            (original_w, original_h),
            interpolation=cv2.INTER_LINEAR,
        )
        return np.asarray(upsampled, dtype=np.float32)

    def predict_anomaly_map(self, x: object = MISSING, **kwargs: object) -> NDArray:
        items = list(
            cast(
                Iterable[Union[str, np.ndarray]],
                resolve_legacy_x_keyword(x, kwargs, method_name="predict_anomaly_map"),
            )
        )
        maps = [self.get_anomaly_map(item) for item in items]
        return np.stack(maps)

    def save_checkpoint(self, path: str | Path) -> Path:
        if self._memory_bank is None or self._knn_index is None or self._n_neighbors_fit is None:
            raise RuntimeError(MODEL_NOT_FITTED_ERROR)

        if not isinstance(self.embedder, TorchvisionConvPatchEmbedder):
            raise TypeError(
                "PatchCore-lite-map checkpointing currently supports only "
                "TorchvisionConvPatchEmbedder."
            )

        from pyimgano.models.serialization import save_model

        payload = {
            "schema_version": 1,
            "detector": "vision_patchcore_lite_map",
            "config": {
                "contamination": float(self.contamination),
                "pretrained": bool(self.pretrained),
                "knn_backend": str(self.knn_backend),
                "metric": str(self.metric),
                "n_neighbors": int(self.n_neighbors),
                "coreset_sampling_ratio": float(self.coreset_sampling_ratio),
                "random_seed": int(self.random_seed),
                "aggregation_method": str(self.aggregation_method),
                "aggregation_topk": float(self.aggregation_topk),
            },
            "embedder": _embedder_to_payload(self.embedder),
            "state": {
                "memory_bank": np.asarray(self._memory_bank, dtype=np.float32),
                "decision_scores_": (
                    np.asarray(self.decision_scores_, dtype=np.float64)
                    if self.decision_scores_ is not None
                    else None
                ),
                "threshold_": (float(self.threshold_) if self.threshold_ is not None else None),
                "n_neighbors_fit": int(self._n_neighbors_fit),
            },
        }
        return save_model(payload, path)

    def load_checkpoint(self, path: str | Path) -> None:
        from pyimgano.models.serialization import load_model

        payload = load_model(path)
        if not isinstance(payload, dict):
            raise ValueError("Invalid PatchCore-lite-map checkpoint payload: expected a dict.")
        if str(payload.get("detector", "")) != "vision_patchcore_lite_map":
            raise ValueError(
                "Invalid PatchCore-lite-map checkpoint payload: detector marker mismatch."
            )

        config = dict(cast(dict[str, object], payload.get("config", {})))
        self.contamination = float(config.get("contamination", self.contamination))
        self.pretrained = bool(config.get("pretrained", self.pretrained))
        self.knn_backend = str(config.get("knn_backend", self.knn_backend))
        self.metric = str(config.get("metric", self.metric))
        self.n_neighbors = int(config.get("n_neighbors", self.n_neighbors))
        self.coreset_sampling_ratio = float(
            config.get("coreset_sampling_ratio", self.coreset_sampling_ratio)
        )
        self.random_seed = int(config.get("random_seed", self.random_seed))
        self.aggregation_method = cast(
            AggregationMethod, config.get("aggregation_method", self.aggregation_method)
        )
        self.aggregation_topk = float(config.get("aggregation_topk", self.aggregation_topk))

        embedder_payload = payload.get("embedder", None)
        if not isinstance(embedder_payload, dict):
            raise ValueError("Invalid PatchCore-lite-map checkpoint payload: missing embedder.")
        self.embedder = _embedder_from_payload(embedder_payload)

        state = payload.get("state", None)
        if not isinstance(state, dict):
            raise ValueError(
                "Invalid PatchCore-lite-map checkpoint payload: missing detector state."
            )

        self._memory_bank = np.asarray(state["memory_bank"], dtype=np.float32)
        self.decision_scores_ = (
            np.asarray(state["decision_scores_"], dtype=np.float64)
            if state.get("decision_scores_", None) is not None
            else None
        )
        self.threshold_ = (
            float(state["threshold_"]) if state.get("threshold_", None) is not None else None
        )
        self._n_neighbors_fit = int(state["n_neighbors_fit"])
        self._knn_index = build_knn_index(
            backend=self.knn_backend,
            n_neighbors=self._n_neighbors_fit,
            metric=str(self.metric),
        )
        self._knn_index.fit(self._memory_bank)


__all__ = ["VisionPatchCoreLiteMap"]
