"""
room_classifier.py — Classify room/environment type from detected objects.

Two modes:
  1. rule_based  — fast heuristic from object co-occurrence (always available).
  2. mobilenet   — MobileNetV2 scene classifier (TensorRT FP16) from a
                   Places365-pretrained model exported to ONNX.
                   Falls back to rule_based if engine file is missing.

Rule-based logic:
  Each room type has a set of indicator objects (COCO class IDs) with weights.
  Score = weighted sum of detected indicators, normalised to softmax.
  Highest score wins.

MobileNet mode:
  Input: 224×224 RGB float32 CHW
  Output: [1, 8] logits (one per ROOM_* constant)
  Engine: /config/room_classifier.engine  (built via scripts/build_scene_trt.py)
"""

from __future__ import annotations
import math
from typing import List

import numpy as np


# Room type constants (must match RoomClassification.msg ROOM_* values)
ROOM_UNKNOWN     = 0
ROOM_LIVING_ROOM = 1
ROOM_KITCHEN     = 2
ROOM_HALLWAY     = 3
ROOM_BEDROOM     = 4
ROOM_GARAGE      = 5
ROOM_OUTDOOR     = 6
ROOM_OFFICE      = 7

ROOM_NAMES = {
    ROOM_UNKNOWN:     'unknown',
    ROOM_LIVING_ROOM: 'living_room',
    ROOM_KITCHEN:     'kitchen',
    ROOM_HALLWAY:     'hallway',
    ROOM_BEDROOM:     'bedroom',
    ROOM_GARAGE:      'garage',
    ROOM_OUTDOOR:     'outdoor',
    ROOM_OFFICE:      'office',
}

N_ROOMS = len(ROOM_NAMES)

# ── Rule-based indicator weights ──────────────────────────────────────────────
# {room_type: {class_id: weight}}
_INDICATORS = {
    ROOM_LIVING_ROOM: {57: 3.0, 56: 2.0, 62: 2.5, 65: 1.5},     # couch, chair, tv, bed
    ROOM_KITCHEN:     {72: 4.0, 69: 3.5, 68: 3.0, 71: 2.0, 46: 1.5},  # fridge, oven, microwave, sink, banana
    ROOM_HALLWAY:     {80: 3.0, 28: 1.5, 25: 1.5},               # door, backpack, umbrella
    ROOM_BEDROOM:     {65: 4.0, 56: 1.5, 62: 1.0},               # bed, chair, tv
    ROOM_GARAGE:      {2: 3.0, 3: 3.0, 7: 2.0, 1: 2.0},          # car, motorcycle, truck, bicycle
    ROOM_OUTDOOR:     {14: 2.0, 17: 2.0, 0: 0.5, 2: 1.5},        # bird, horse, person(weak), car
    ROOM_OFFICE:      {63: 3.0, 66: 2.5, 73: 2.0, 56: 1.5},      # laptop, keyboard, book, chair
}


def _softmax(x: np.ndarray) -> np.ndarray:
    e = np.exp(x - x.max())
    return e / e.sum()


def classify_rule_based(
    class_ids: List[int],
    confidences: List[float],
) -> tuple[int, str, float, list[float]]:
    """
    Classify room type from detected object class IDs.

    Returns:
        (room_type, room_name, confidence, class_scores[N_ROOMS])
    """
    scores = np.zeros(N_ROOMS, dtype=np.float32)
    for cid, conf in zip(class_ids, confidences):
        for room_type, indicators in _INDICATORS.items():
            if cid in indicators:
                scores[room_type] += indicators[cid] * conf

    if scores.max() < 0.1:
        # No indicator objects found — unknown
        uniform = np.ones(N_ROOMS, dtype=np.float32) / N_ROOMS
        return ROOM_UNKNOWN, ROOM_NAMES[ROOM_UNKNOWN], 0.0, uniform.tolist()

    probs       = _softmax(scores)
    room_type   = int(probs.argmax())
    confidence  = float(probs[room_type])
    return room_type, ROOM_NAMES[room_type], confidence, probs.tolist()


# ── MobileNetV2 TRT/ONNX backend ─────────────────────────────────────────────

class _MobileNetTRT:
    """TensorRT MobileNetV2 scene classifier (Jetson Orin)."""

    _INPUT_SIZE = 224

    def __init__(self, engine_path: str):
        import tensorrt as trt
        import pycuda.driver as cuda
        import pycuda.autoinit  # noqa: F401

        self._cuda = cuda
        logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, 'rb') as f, trt.Runtime(logger) as rt:
            self._engine = rt.deserialize_cuda_engine(f.read())
        self._context = self._engine.create_execution_context()
        self._inputs, self._outputs, self._bindings = [], [], []
        for i in range(self._engine.num_io_tensors):
            name  = self._engine.get_tensor_name(i)
            dtype = trt.nptype(self._engine.get_tensor_dtype(name))
            shape = tuple(self._engine.get_tensor_shape(name))
            nbytes = int(np.prod(shape)) * np.dtype(dtype).itemsize
            host   = cuda.pagelocked_empty(shape, dtype)
            device = cuda.mem_alloc(nbytes)
            self._bindings.append(int(device))
            if self._engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                self._inputs.append({'host': host, 'device': device})
            else:
                self._outputs.append({'host': host, 'device': device})
        self._stream = cuda.Stream()

    def predict(self, bgr: np.ndarray) -> np.ndarray:
        import cv2
        rgb = cv2.cvtColor(
            cv2.resize(bgr, (self._INPUT_SIZE, self._INPUT_SIZE)), cv2.COLOR_BGR2RGB
        )
        blob = (rgb.astype(np.float32) / 255.0).transpose(2, 0, 1)[None]
        np.copyto(self._inputs[0]['host'], blob.ravel())
        self._cuda.memcpy_htod_async(
            self._inputs[0]['device'], self._inputs[0]['host'], self._stream
        )
        self._context.execute_async_v2(self._bindings, self._stream.handle)
        for out in self._outputs:
            self._cuda.memcpy_dtoh_async(out['host'], out['device'], self._stream)
        self._stream.synchronize()
        return self._outputs[0]['host'].copy().ravel()   # logits [N_ROOMS]


class _MobileNetONNX:
    """ONNX Runtime fallback for non-Jetson environments."""

    _INPUT_SIZE = 224

    def __init__(self, onnx_path: str):
        import onnxruntime as ort
        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        self._session = ort.InferenceSession(onnx_path, providers=providers)
        self._input_name = self._session.get_inputs()[0].name

    def predict(self, bgr: np.ndarray) -> np.ndarray:
        import cv2
        rgb = cv2.cvtColor(
            cv2.resize(bgr, (self._INPUT_SIZE, self._INPUT_SIZE)), cv2.COLOR_BGR2RGB
        )
        blob = (rgb.astype(np.float32) / 255.0).transpose(2, 0, 1)[None]
        result = self._session.run(None, {self._input_name: blob})
        return result[0].ravel()


class RoomClassifier:
    """
    High-level room classifier.  Tries TRT engine, then ONNX, then rule-based.
    """

    def __init__(self, engine_path: str = '', onnx_path: str = ''):
        self._backend = None
        self._method  = 'rule_based'

        if engine_path:
            try:
                self._backend = _MobileNetTRT(engine_path)
                self._method  = 'mobilenet_trt'
                return
            except Exception:
                pass
        if onnx_path:
            try:
                self._backend = _MobileNetONNX(onnx_path)
                self._method  = 'mobilenet_onnx'
                return
            except Exception:
                pass

    def classify(
        self,
        bgr: np.ndarray | None,
        class_ids: List[int],
        confidences: List[float],
    ) -> tuple[int, str, float, list[float], str]:
        """
        Classify the current frame.

        Returns:
            (room_type, room_name, confidence, class_scores, method)
        """
        if self._backend is not None and bgr is not None:
            logits = self._backend.predict(bgr)
            probs  = _softmax(logits.astype(np.float32))
            room   = int(probs.argmax())
            return room, ROOM_NAMES[room], float(probs[room]), probs.tolist(), self._method

        room, name, conf, scores = classify_rule_based(class_ids, confidences)
        return room, name, conf, scores, 'rule_based'
