"""
Crack segmentation model using nnU-Net.

Architecture: nnU-Net 2D with 4-fold ensemble.
Trained on 30K+ crack images across diverse surface types.
"""

import cv2
import torch
import numpy as np
from pathlib import Path
from torchvision.transforms import Normalize
from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor

from deadlift.crack.detect import filter_detections
from deadlift.utils.weights import download_weights

HF_REPO = "MrNergigante/crack-detection-nnunet"
WEIGHT_FILES = [
    "dataset.json",
    "plans.json",
    "dataset_fingerprint.json",
    "fold_0/checkpoint_final.pth",
    "fold_1/checkpoint_final.pth",
    "fold_2/checkpoint_final.pth",
    "fold_4/checkpoint_final.pth",
]
FOLDS = (0, 1, 2, 4)


class CrackModel:
    """Pixel-level crack segmentation model.

    Args:
        device: "auto", "cuda", "mps", or "cpu".
        mirroring: Enable test-time augmentation (better accuracy, ~8x slower).
        tile_step: Tile overlap (0.5 = max overlap, 1.0 = none).
        folds: Which model folds to ensemble.
        hf_token: HuggingFace token (only needed for private repos).
    """

    def __init__(
        self,
        device: str = "auto",
        mirroring: bool = True,
        tile_step: float = 0.625,
        folds: tuple[int, ...] = FOLDS,
        hf_token: str | None = None,
    ):
        self.device = self._resolve_device(device)

        # Download weights from HuggingFace (cached after first download)
        weights_dir = download_weights(HF_REPO, WEIGHT_FILES, token=hf_token)

        # Initialize nnU-Net predictor
        self.predictor = nnUNetPredictor(
            tile_step_size=tile_step,
            use_gaussian=True,
            use_mirroring=mirroring,
            perform_everything_on_device=True,
            device=self.device,
            verbose=False,
            verbose_preprocessing=False,
            allow_tqdm=True,
        )

        self.predictor.initialize_from_trained_model_folder(
            str(weights_dir),
            use_folds=folds,
            checkpoint_name="checkpoint_final.pth",
        )

        self.predictor.device = self.device
        self.predictor.network = self.predictor.network.to(self.device)
        self.predictor.network.eval()

    @staticmethod
    def _resolve_device(device: str) -> torch.device:
        if device == "auto":
            if torch.cuda.is_available():
                return torch.device("cuda", 0)
            elif torch.backends.mps.is_available():
                return torch.device("mps")
            return torch.device("cpu")
        return torch.device(device)

    def _run_segmentation(self, img: np.ndarray, rgb: bool = True) -> tuple[np.ndarray, np.ndarray]:
        """Run nnU-Net segmentation on a single image.

        Returns:
            crack_prob: Float array of crack probability per pixel (0-1).
            crack_mask: Uint8 binary mask (255 = crack, 0 = background).
        """
        t = torch.tensor(img, dtype=torch.float32)
        t = t.moveaxis(-1, 0) if t.shape[0] != 3 else t

        # Replace all-zero patches with random noise (nnU-Net requirement)
        zero_mask = torch.all(t == 0, dim=0)
        t[:, zero_mask] = torch.rand_like(t)[:, zero_mask]

        # BGR to RGB conversion if needed
        if rgb:
            t = t[[2, 1, 0], ...]

        # Per-channel normalization
        data = Normalize(t.mean((1, 2)), t.std((1, 2)))(t)
        data = data.unsqueeze(1)

        # Run inference
        logits = self.predictor.predict_logits_from_preprocessed_data(data).squeeze().cpu()
        softmax = torch.nn.functional.softmax(logits.to(torch.float32), dim=0)
        argmax = torch.argmax(logits, dim=0).to(torch.uint8)

        # Extract crack channel (index 1)
        crack_prob = softmax[1].numpy()
        crack_mask = (argmax.numpy() * 255).astype(np.uint8)

        return crack_prob, crack_mask

    def predict(
        self,
        image: str | Path | np.ndarray,
        min_confidence: float = 0.3,
        min_area: int = 100,
    ) -> dict:
        """Run crack detection on an image.

        Args:
            image: Path to image file or numpy array (BGR).
            min_confidence: Minimum softmax confidence to keep a detection.
            min_area: Minimum contour area in pixels.

        Returns:
            Dict with keys: crack_prob, crack_mask, detections, overlay.
        """
        if isinstance(image, (str, Path)):
            img_bgr = cv2.imread(str(image), cv2.IMREAD_COLOR)
            if img_bgr is None:
                raise FileNotFoundError(f"Could not read image: {image}")
        else:
            img_bgr = image

        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        crack_prob, crack_mask = self._run_segmentation(img_rgb, rgb=True)

        detections = filter_detections(
            crack_mask, crack_prob,
            min_confidence=min_confidence,
            min_area=min_area,
        )

        # Build overlay
        overlay = img_bgr.copy()
        overlay[crack_mask > 0] = [0, 0, 255]
        blended = cv2.addWeighted(img_bgr, 0.6, overlay, 0.4, 0)

        for det in detections:
            x, y, w, h = det["bbox"]
            label = f"crack {det['confidence']:.2f}"
            cv2.rectangle(blended, (x, y), (x + w, y + h), (0, 0, 255), 2)
            cv2.putText(blended, label, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)

        return {
            "crack_prob": crack_prob,
            "crack_mask": crack_mask,
            "detections": detections,
            "overlay": blended,
        }
