from __future__ import annotations

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

import numpy as np
from numpy.typing import NDArray

from ._legacy_x import MISSING, resolve_legacy_x_keyword
from .anomalydino import (
    PatchEmbedder,
    TorchHubDinoV2Embedder,
    _embedder_from_checkpoint_payload,
    _embedder_to_checkpoint_payload,
)
from .knn_index import KNNIndex, build_knn_index
from .patchknn_core import AggregationMethod, aggregate_patch_scores, reshape_patch_scores
from .registry import register_model

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



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


@register_model(
    "vision_softpatch",
    tags=(
        "vision",
        "deep",
        "softpatch",
        "patchknn",
        "robust",
        "numpy",
        "pixel_map",
        "memory_bank",
        "neighbors",
    ),
    metadata={
        "description": "SoftPatch-inspired robust patch-memory detector (few-shot friendly)",
        "paper": "SoftPatch: Unsupervised Anomaly Detection with Noisy Data",
        "year": 2022,
    },
)
class VisionSoftPatch:
    """SoftPatch-inspired robust patch-memory detector (inference-first).

    Notes
    -----
    This detector is intentionally implemented in a Patch-kNN style:

    - Extract patch embeddings for each image.
    - Build a patch memory bank from normal/reference images.
    - Score patches by kNN distance to the memory bank.
    - Upsample patch scores to an anomaly map and aggregate to an image score.

    The embedder is injectable so unit tests can run without torch and without
    downloading weights.
    """

    def __init__(
        self,
        *,
        embedder: Optional[PatchEmbedder] = None,
        contamination: float = 0.1,
        pretrained: bool = False,
        knn_backend: str = "sklearn",
        n_neighbors: int = 1,
        coreset_sampling_ratio: float = 1.0,
        random_seed: int = 0,
        train_patch_outlier_quantile: float = 0.0,
        aggregation_method: AggregationMethod = "topk_mean",
        aggregation_topk: float = 0.01,
        device: str = "cpu",
        image_size: int = 518,
        dino_model_name: str = "dinov2_vits14",
    ) -> None:
        if embedder is None:
            if bool(pretrained):
                embedder = TorchHubDinoV2Embedder(
                    model_name=dino_model_name,
                    device=device,
                    image_size=image_size,
                )
            else:
                raise ValueError(
                    "vision_softpatch requires a patch embedder.\n"
                    "Pass embedder=... (recommended, offline) or set pretrained=True to allow "
                    "torch.hub to load DINOv2 weights (may download from the internet)."
                )

        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.knn_backend = str(knn_backend)
        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.train_patch_outlier_quantile = float(train_patch_outlier_quantile)
        if not (0.0 <= self.train_patch_outlier_quantile < 1.0):
            raise ValueError(
                "train_patch_outlier_quantile must be in [0, 1). "
                f"Got {self.train_patch_outlier_quantile}."
            )

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

        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
        self.filtered_patches_: int = 0

    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 self.threshold_ is None:
            raise RuntimeError(MODEL_NOT_FITTED_ERROR)

        from pyimgano.utils.optional_deps import require

        torch = require("torch", extra="torch", purpose="VisionSoftPatch checkpoint saving")

        out_path = Path(path)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        torch.save(
            {
                "embedder": _embedder_to_checkpoint_payload(self.embedder),
                "memory_bank": np.asarray(self._memory_bank, dtype=np.float32),
                "n_neighbors_fit": int(self._n_neighbors_fit),
                "decision_scores_": np.asarray(self.decision_scores_, dtype=np.float64),
                "threshold_": float(self.threshold_),
                "filtered_patches_": int(self.filtered_patches_),
            },
            out_path,
        )
        return out_path

    def load_checkpoint(self, path: str | Path) -> None:
        from pyimgano.utils.optional_deps import require

        torch = require("torch", extra="torch", purpose="VisionSoftPatch checkpoint loading")

        state = torch.load(Path(path), map_location="cpu", weights_only=False)
        if not isinstance(state, dict):
            raise ValueError("Invalid VisionSoftPatch checkpoint payload.")

        embedder_payload = state.get("embedder", None)
        if not isinstance(embedder_payload, dict):
            raise ValueError("VisionSoftPatch checkpoint is missing embedder payload.")
        self.embedder = _embedder_from_checkpoint_payload(dict(embedder_payload))
        self._memory_bank = np.asarray(state["memory_bank"], dtype=np.float32)
        self._n_neighbors_fit = min(
            int(state.get("n_neighbors_fit", self.n_neighbors)),
            int(self._memory_bank.shape[0]),
        )
        self._knn_index = build_knn_index(
            backend=self.knn_backend,
            n_neighbors=self._n_neighbors_fit,
        )
        self._knn_index.fit(self._memory_bank)
        self.decision_scores_ = np.asarray(state["decision_scores_"], dtype=np.float64)
        self.threshold_ = float(state["threshold_"])
        self.filtered_patches_ = int(state.get("filtered_patches_", 0))

    @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.train_patch_outlier_quantile > 0.0:
            mean = np.mean(memory_bank, axis=0, keepdims=True)
            distances = np.linalg.norm(memory_bank - mean, axis=1)
            keep_q = 1.0 - float(self.train_patch_outlier_quantile)
            threshold = float(np.quantile(distances, keep_q))
            keep_mask = distances <= threshold
            if int(np.sum(keep_mask)) < 1:
                keep_mask[np.argmin(distances)] = True
            self.filtered_patches_ = int(memory_bank.shape[0] - int(np.sum(keep_mask)))
            memory_bank = memory_bank[keep_mask]
        else:
            self.filtered_patches_ = 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)
        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}")

        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)
