"""
Face crop classifier.

Loads a MobileNetV3 model via ONNX/TensorRT FP32, applies temperature-scaled
softmax, and returns per-crop probabilities for the target class.
"""

import logging
from pathlib import Path
from typing import Any, List, Optional

import numpy as np

logger = logging.getLogger(__name__)

# ImageNet normalisation constants
_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
_STD  = np.array([0.229, 0.224, 0.225], dtype=np.float32)


def calibrated_softmax(logits: np.ndarray, temperature: float = 1.0) -> np.ndarray:
    """Softmax with temperature scaling; returns class probabilities."""
    scaled = logits / max(temperature, 1e-8)
    shifted = scaled - scaled.max(axis=-1, keepdims=True)   # numerical stability
    exp = np.exp(shifted)
    return exp / exp.sum(axis=-1, keepdims=True)


class FaceClassifier:
    """Classify face crops and return calibrated target-class probabilities."""

    def __init__(
        self,
        engine_path: str,
        temperature: float = 1.0,
        input_size: int = 224,
        device: str = "cuda",
    ):
        self.engine_path = engine_path
        self.temperature = temperature
        self.input_size = input_size
        self.device = device
        self._session: Optional[Any] = None
        self._load_model()

    def _load_model(self):
        path = Path(self.engine_path)
        if not path.exists():
            logger.warning("Classifier engine not found at %s; using placeholder", self.engine_path)
            return

        suffix = path.suffix.lower()
        if suffix == ".onnx":
            try:
                import onnxruntime as ort
                providers = (["CUDAExecutionProvider"] if self.device.startswith("cuda")
                             else ["CPUExecutionProvider"])
                self._session = ort.InferenceSession(str(path), providers=providers)
                logger.info("Loaded ONNX classifier from %s", self.engine_path)
            except Exception as exc:
                logger.warning("Failed to load ONNX classifier: %s", exc)
        elif suffix in (".trt", ".engine", ".plan"):
            try:
                import tensorrt as trt
                with trt.Logger() as trt_logger, open(path, "rb") as f:
                    runtime = trt.Runtime(trt_logger)
                    self._session = runtime.deserialize_cuda_engine(f.read())
                logger.info("Loaded TensorRT classifier from %s", self.engine_path)
            except Exception as exc:
                logger.warning("Failed to load TensorRT classifier: %s", exc)
        else:
            logger.warning("Unsupported classifier format: %s", suffix)

    def classify(self, crop_paths: List[str], batch_size: int = 16) -> List[float]:
        """Return target-class probabilities (p ∈ [0,1]) for each crop path."""
        results: List[float] = []
        for i in range(0, len(crop_paths), batch_size):
            batch = crop_paths[i : i + batch_size]
            results.extend(self._classify_batch(batch))
        return results

    def _classify_batch(self, crop_paths: List[str]) -> List[float]:
        preprocessed = []
        for path in crop_paths:
            try:
                preprocessed.append(self._preprocess(path))
            except Exception as exc:
                logger.warning("Preprocessing failed for %s: %s", path, exc)
                preprocessed.append(np.zeros((3, self.input_size, self.input_size), dtype=np.float32))

        batch = np.stack(preprocessed, axis=0)  # (N, 3, H, W)
        logits = self._infer(batch)             # (N, 2)
        probs = calibrated_softmax(logits, self.temperature)
        return probs[:, 1].tolist()             # target-class column

    def _preprocess(self, crop_path: str) -> np.ndarray:
        from PIL import Image
        image = Image.open(crop_path).convert("RGB")
        if image.size != (self.input_size, self.input_size):
            image = image.resize((self.input_size, self.input_size), Image.Resampling.BILINEAR)
        arr = np.array(image, dtype=np.float32) / 255.0
        arr = (arr - _MEAN) / _STD
        image.close()
        return np.transpose(arr, (2, 0, 1))    # HWC → CHW

    def _infer(self, batch: np.ndarray) -> np.ndarray:
        if self._session is None:
            return self._placeholder_logits(batch.shape[0])

        try:
            if hasattr(self._session, "run"):
                input_name = self._session.get_inputs()[0].name
                output = self._session.run(None, {input_name: batch})
                return np.array(output[0])

            # TensorRT path
            import pycuda.autoinit  # noqa: F401
            import pycuda.driver as cuda

            context = self._session.create_execution_context()
            out_shape = (batch.shape[0], 2)
            d_in  = cuda.mem_alloc(batch.nbytes)
            d_out = cuda.mem_alloc(np.prod(out_shape) * np.dtype(np.float32).itemsize)
            stream = cuda.Stream()
            cuda.memcpy_htod_async(d_in, batch, stream)
            context.execute_async_v2(bindings=[int(d_in), int(d_out)], stream_handle=stream.handle)
            output = np.empty(out_shape, dtype=np.float32)
            cuda.memcpy_dtoh_async(output, d_out, stream)
            stream.synchronize()
            return output
        except Exception as exc:
            logger.warning("Classifier inference failed: %s", exc)
            return self._placeholder_logits(batch.shape[0])

    @staticmethod
    def _placeholder_logits(n: int) -> np.ndarray:
        """Return neutral logits when no model is loaded."""
        return np.zeros((n, 2), dtype=np.float32)