from __future__ import annotations

"""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 typing import Iterable, Optional, Tuple, Union

import numpy as np
from numpy.typing import NDArray

from pyimgano.features.torchvision_conv_patch_embedder import TorchvisionConvPatchEmbedder

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]


@register_model(
    "vision_patchcore_lite_map",
    tags=("vision", "deep", "patchcore", "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(
                "coreset_sampling_ratio must be in (0, 1]. " f"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. Call fit() first.")
        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: Iterable[Union[str, np.ndarray]], y=None):
        items = list(X)
        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. Call fit() first.")
        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: Iterable[Union[str, np.ndarray]]) -> NDArray:
        items = list(X)
        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: Iterable[Union[str, np.ndarray]]) -> NDArray:
        if self.threshold_ is None:
            raise RuntimeError("Model not fitted. Call fit() first.")
        scores = self.decision_function(X)
        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: Iterable[Union[str, np.ndarray]]) -> NDArray:
        items = list(X)
        maps = [self.get_anomaly_map(item) for item in items]
        return np.stack(maps)


__all__ = ["VisionPatchCoreLiteMap"]
