from __future__ import annotations

"""
DRAEM: SSIM-based Discriminatively Trained Reconstruction Embedding for Anomaly Detection.

DRAEM uses synthetic anomalies during training and a discriminative approach
for robust anomaly detection and localization.

Reference:
    Zavrtanik, V., Kristan, M., & Skočaj, D. (2021).
    DRAEM-A discriminatively trained reconstruction embedding for surface anomaly detection.
    In Proceedings of the IEEE/CVF International Conference on Computer Vision (pp. 8330-8339).
"""

import logging
import os
from pathlib import Path
from typing import Dict, Iterable, Optional, Union, cast

import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from numpy.typing import NDArray
from torch.optim import Adam
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms

from .baseCv import BaseVisionDeepDetector
from .deep_io import safe_torch_load
from .registry import register_model

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


logger = logging.getLogger(__name__)

ImageInput = Union[str, np.ndarray]


def _resolve_legacy_x_keyword(
    x: object, kwargs: Dict[str, object], *, method_name: str
) -> Iterable[ImageInput]:
    """Accept legacy `X=` inputs while keeping lowercase local naming."""

    legacy_x = kwargs.pop("X", _MISSING)
    if kwargs:
        unexpected = next(iter(kwargs))
        raise TypeError(f"{method_name}() got an unexpected keyword argument {unexpected!r}")
    if x is _MISSING:
        if legacy_x is _MISSING:
            raise TypeError(f"{method_name}() missing 1 required positional argument: 'x'")
        return cast(Iterable[ImageInput], legacy_x)
    if legacy_x is not _MISSING:
        raise TypeError(f"{method_name}() got multiple values for argument 'x'")
    return cast(Iterable[ImageInput], x)


class SimpleUNet(nn.Module):
    """Simple U-Net architecture for reconstruction."""

    def __init__(self, in_channels=3, out_channels=3):
        super().__init__()

        # Encoder
        self.enc1 = self._conv_block(in_channels, 64)
        self.enc2 = self._conv_block(64, 128)
        self.enc3 = self._conv_block(128, 256)
        self.enc4 = self._conv_block(256, 512)

        # Decoder
        self.dec4 = self._upconv_block(512, 256)
        self.dec3 = self._upconv_block(512, 128)  # 512 because of skip connection
        self.dec2 = self._upconv_block(256, 64)
        # The encoder downsamples 3 times (max-pool), so the decoder should
        # upsample 3 times to match the input resolution. Use a conv block for
        # the final stage to keep output spatial size aligned with the input.
        self.dec1 = self._conv_block(128, out_channels)

        self.final = nn.Conv2d(out_channels, out_channels, 1)

    def _conv_block(self, in_ch, out_ch):
        return nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        )

    def _upconv_block(self, in_ch, out_ch):
        return nn.Sequential(
            nn.ConvTranspose2d(in_ch, out_ch, 2, stride=2),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        # Encoder
        e1 = self.enc1(x)
        e2 = self.enc2(F.max_pool2d(e1, 2))
        e3 = self.enc3(F.max_pool2d(e2, 2))
        e4 = self.enc4(F.max_pool2d(e3, 2))

        # Decoder with skip connections
        d4 = self.dec4(e4)
        d4 = torch.cat([d4, e3], dim=1)

        d3 = self.dec3(d4)
        d3 = torch.cat([d3, e2], dim=1)

        d2 = self.dec2(d3)
        d2 = torch.cat([d2, e1], dim=1)

        d1 = self.dec1(d2)

        return torch.sigmoid(self.final(d1))


class ImagePathDataset(Dataset):
    """Dataset for loading images from paths or in-memory numpy arrays."""

    def __init__(self, image_paths, transform=None, rng: Optional[np.random.Generator] = None):
        self.image_paths = image_paths
        self.transform = transform
        self.rng = (
            rng
            if rng is not None
            else np.random.default_rng(int.from_bytes(os.urandom(8), "little"))
        )

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        item = self.image_paths[idx]
        if isinstance(item, np.ndarray):
            if item.dtype != np.uint8:
                raise ValueError(f"Expected uint8 RGB image, got dtype={item.dtype}")
            if item.ndim != 3 or item.shape[2] != 3:
                raise ValueError(f"Expected shape (H,W,3), got {item.shape}")
            img = np.ascontiguousarray(item)
        else:
            img_path = str(item)
            img = cv2.imread(img_path)
            if img is None:
                raise ValueError(f"Failed to load image: {img_path}")
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        if self.transform:
            img = self.transform(img)

        # Create synthetic anomaly for training
        augmented = self._add_synthetic_anomaly(img)

        return augmented, img

    def _add_synthetic_anomaly(self, img):
        """Add synthetic anomaly to image."""
        # Simple implementation: add random noise patches
        augmented = img.clone()
        rng = self.rng

        # Randomly add 1-3 anomalous regions
        n_anomalies = int(rng.integers(1, 4))

        for _ in range(n_anomalies):
            # Random patch size
            h = int(rng.integers(10, 30))
            w = int(rng.integers(10, 30))
            # Random position
            y = int(rng.integers(0, img.shape[1] - h))
            x = int(rng.integers(0, img.shape[2] - w))

            # Add random noise or texture
            if rng.random() > 0.5:
                # Random noise
                augmented[:, y : y + h, x : x + w] = torch.rand(3, h, w)
            else:
                # Invert region
                augmented[:, y : y + h, x : x + w] = 1 - augmented[:, y : y + h, x : x + w]

        return augmented


@register_model(
    "vision_draem",
    tags=("vision", "deep", "draem", "reconstruction", "synthetic", "numpy", "pixel_map"),
    metadata={
        "description": "DRAEM - Discriminatively trained reconstruction (ICCV 2021)",
        "paper": "DRAEM: Discriminatively Trained Reconstruction Embedding",
        "year": 2021,
    },
)
class VisionDRAEM(BaseVisionDeepDetector):
    """
    DRAEM anomaly detector using synthetic anomalies.

    Parameters
    ----------
    image_size : int, default=256
        Input image size
    epochs : int, default=100
        Number of training epochs
    batch_size : int, default=8
        Training batch size
    lr : float, default=0.0001
        Learning rate
    num_workers : int, default=0
        Number of workers for the training DataLoader.
    device : str, default='cpu'
        Device to run model on

    Examples
    --------
    >>> detector = VisionDRAEM(epochs=50, device='cuda')
    >>> detector.fit(train_images)
    >>> scores = detector.decision_function(test_images)
    >>> labels = detector.predict(test_images)  # 0=normal, 1=anomaly
    """

    def __init__(
        self,
        image_size: int = 256,
        epochs: int = 100,
        batch_size: int = 8,
        lr: float = 0.0001,
        num_workers: int = 0,
        device: str = "cpu",
        **kwargs,
    ):
        """Initialize DRAEM detector."""
        super().__init__(**kwargs)

        self.image_size = image_size
        self.epochs = epochs
        self.batch_size = batch_size
        self.lr = lr
        self.num_workers = num_workers
        self.device = device
        self._is_fitted = False

        # Build model
        self.model = SimpleUNet(in_channels=3, out_channels=3)
        self.model.to(self.device)

        # Image preprocessing
        self.transform = transforms.Compose(
            [
                transforms.ToPILImage(),
                transforms.Resize((image_size, image_size)),
                transforms.ToTensor(),
            ]
        )

        logger.info(
            "Initialized DRAEM with image_size=%d, epochs=%d, batch_size=%d, device=%s",
            image_size,
            epochs,
            batch_size,
            device,
        )

    def save_checkpoint(self, path: str | Path) -> Path:
        if not self._is_fitted or not hasattr(self, "threshold_"):
            raise RuntimeError(MODEL_NOT_FITTED_ERROR)

        out_path = Path(path)
        out_path.parent.mkdir(parents=True, exist_ok=True)

        raw_state = self.model.state_dict()
        model_state_dict: 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):
                model_state_dict[str(key)] = detach().cpu()
            else:
                model_state_dict[str(key)] = value

        torch.save(
            {
                "model_state_dict": model_state_dict,
                "decision_scores_": torch.as_tensor(
                    np.asarray(self.decision_scores_, dtype=np.float64),
                    dtype=torch.float64,
                ),
                "threshold_": float(self.threshold_),
                "is_fitted": True,
            },
            out_path,
        )
        return out_path

    def load_checkpoint(self, path: str | Path) -> None:
        state = safe_torch_load(path, map_location="cpu")
        if not isinstance(state, dict):
            raise ValueError("Invalid VisionDRAEM checkpoint payload.")

        model_state_dict = state.get("model_state_dict", None)
        if not isinstance(model_state_dict, dict):
            raise ValueError("VisionDRAEM checkpoint is missing model_state_dict.")

        self.model.load_state_dict(dict(model_state_dict), strict=False)
        self.model.to(self.device)
        self.model.eval()
        self.decision_scores_ = np.asarray(state["decision_scores_"], dtype=np.float64)
        self.threshold_ = float(state["threshold_"])
        self._is_fitted = bool(state.get("is_fitted", True))

    def fit(
        self,
        x: object = _MISSING,
        y: Optional[NDArray] = None,
        **kwargs: object,
    ) -> "VisionDRAEM":
        """
        Train DRAEM on normal images.

        Parameters
        ----------
        X : iterable of str
            Paths to normal training images
        y : array-like, optional
            Ignored

        Returns
        -------
        self : VisionDRAEM
        """
        del y
        logger.info("Training DRAEM detector")

        x_iter = _resolve_legacy_x_keyword(x, kwargs, method_name="fit")
        x_list = list(x_iter)
        if not x_list:
            raise ValueError("Training set cannot be empty")

        # Create dataset
        dataset = ImagePathDataset(x_list, transform=self.transform)
        dataloader = DataLoader(
            dataset,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=self.num_workers,
            pin_memory=str(self.device).startswith("cuda"),
        )

        # Setup optimizer
        optimizer = Adam(self.model.parameters(), lr=self.lr, weight_decay=0.0)

        # Training loop
        self.model.train()

        for epoch in range(self.epochs):
            epoch_loss = 0.0

            for augmented, original in dataloader:
                augmented = augmented.to(self.device)
                original = original.to(self.device)

                # Forward pass
                reconstructed = self.model(augmented)

                # L2 reconstruction loss
                loss = F.mse_loss(reconstructed, original)

                # Backward pass
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                epoch_loss += loss.detach().item()

            if (epoch + 1) % 10 == 0:
                avg_loss = epoch_loss / len(dataloader)
                logger.info("Epoch %d/%d, Loss: %.6f", epoch + 1, self.epochs, avg_loss)

        logger.info("DRAEM training completed")

        # Mark as fitted and compute training scores to establish a threshold.
        self._is_fitted = True
        self.decision_scores_ = self.decision_function(x_list)
        self._process_decision_scores()

        return self

    def predict(
        self,
        x: object = _MISSING,
        return_confidence: bool = False,
        **kwargs: object,
    ) -> NDArray:
        """
        Predict binary anomaly labels for test images.

        Parameters
        ----------
        X : iterable of str
            Paths to test images

        Returns
        -------
        labels : ndarray of shape (n_samples,)
            Binary labels (0 = normal, 1 = anomaly)
        """
        if return_confidence:
            raise NotImplementedError(
                f"return_confidence is not implemented for {self.__class__.__name__}"
            )

        if not self._is_fitted or not hasattr(self, "threshold_"):
            raise RuntimeError(MODEL_NOT_FITTED_ERROR)

        x_iter = _resolve_legacy_x_keyword(x, kwargs, method_name="predict")
        scores = self.decision_function(x_iter)
        return (scores >= self.threshold_).astype(int)

    def _load_item_rgb(self, item: ImageInput) -> np.ndarray:  # pragma: no cover
        """Load an input item as a contiguous RGB uint8 numpy array."""

        if isinstance(item, np.ndarray):
            if item.dtype != np.uint8:
                raise ValueError(f"Expected uint8 RGB image, got dtype={item.dtype}")
            if item.ndim != 3 or item.shape[2] != 3:
                raise ValueError(f"Expected shape (H,W,3), got {item.shape}")
            return np.ascontiguousarray(item)

        img = cv2.imread(str(item))
        if img is None:
            raise ValueError(f"Failed to load image: {item}")
        return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    def _score_item(self, item: ImageInput) -> float:  # pragma: no cover
        img = self._load_item_rgb(item)
        img_tensor = self.transform(img).unsqueeze(0).to(self.device)

        reconstructed = self.model(img_tensor)
        error = F.mse_loss(reconstructed, img_tensor, reduction="none")
        return float(error.mean().item())

    def decision_function(
        self,
        x: object = _MISSING,
        batch_size: Optional[int] = None,
        **kwargs: object,
    ) -> NDArray:
        """Compute anomaly scores."""
        # This detector scores one image at a time. Keep `batch_size` for
        # interface compatibility with BaseDeepLearningDetector.
        if batch_size is not None:
            batch_size_int = int(batch_size)
            if batch_size_int <= 0:
                raise ValueError(f"batch_size must be positive integer, got: {batch_size!r}")

        if not self._is_fitted:
            raise RuntimeError(MODEL_NOT_FITTED_ERROR)

        self.model.eval()

        x_iter = _resolve_legacy_x_keyword(x, kwargs, method_name="decision_function")
        x_list = list(x_iter)
        scores = np.zeros(len(x_list), dtype=np.float64)

        logger.info("Computing anomaly scores for %d images", len(x_list))

        with torch.no_grad():
            for idx, item in enumerate(x_list):
                try:
                    scores[idx] = self._score_item(item)
                except Exception as exc:
                    logger.warning("Failed to score item %d: %s", idx, exc)
                    scores[idx] = 0.0

        return scores

    def get_anomaly_map(self, image_path: ImageInput) -> NDArray:
        """Generate pixel-level anomaly heatmap."""
        if not self._is_fitted:
            raise RuntimeError(MODEL_NOT_FITTED_ERROR)

        if isinstance(image_path, np.ndarray):
            if image_path.dtype != np.uint8:
                raise ValueError(f"Expected uint8 RGB image, got dtype={image_path.dtype}")
            if image_path.ndim != 3 or image_path.shape[2] != 3:
                raise ValueError(f"Expected shape (H,W,3), got {image_path.shape}")
            img = np.ascontiguousarray(image_path)
            original_size = (int(img.shape[1]), int(img.shape[0]))  # (W, H)
        else:
            img = cv2.imread(str(image_path))
            if img is None:
                raise ValueError(f"Failed to load image: {image_path}")
            original_size = (int(img.shape[1]), int(img.shape[0]))  # (W, H)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        img_tensor = self.transform(img).unsqueeze(0).to(self.device)

        self.model.eval()
        with torch.no_grad():
            reconstructed = self.model(img_tensor)
            error = F.mse_loss(reconstructed, img_tensor, reduction="none")  # (1, C, H, W)
            anomaly_map = error.mean(dim=1).squeeze(0).cpu().numpy()

        anomaly_map = anomaly_map.astype(np.float32, copy=False)
        anomaly_map = cv2.resize(anomaly_map, original_size, interpolation=cv2.INTER_CUBIC)
        return anomaly_map

    def predict_anomaly_map(self, x: object = _MISSING, **kwargs: object) -> NDArray:
        """Generate pixel-level anomaly maps for a batch of images."""
        x_iter = _resolve_legacy_x_keyword(x, kwargs, method_name="predict_anomaly_map")
        maps = [self.get_anomaly_map(path) for path in x_iter]
        return np.stack(maps)
