# SPDX-License-Identifier: Apache-2.0
"""Host pre-processing of METEOR's runtimes: eight camera frames + calibration + ego speed -> the graph's feed.

The contract (research/meteor/SPEC.md sections 2-3; ``meteor_v157.param.yaml``; tier4/METEOR ``hf/onnx_smoke_test.py:
34-56``, the stated reference for feeding the graph, and ``deploy/cpp/realtime_main.cpp:123-140``):

- ``imgs`` uint8 [1, 8, 3, 432, 768], RGB, CHW, cameras in the training order (``config.CAMERAS``). A frame of another
  size is resized with OpenCV ``INTER_AREA`` (``ttaw.image_area.meteor_resize``, bit-exact), anisotropically if the
  aspect differs; frames are raw and unrectified (training never undistorted).
- ``K`` float32 [1, 8, 3, 3] at 768x432: the intrinsics of the image as captured, row 0 x 768 / w0 and row 1 x 432 / h0
  (``extract_gt.py:128-130``); only fx, fy, cx, cy reach the graph.
- ``T_cam_ego`` float32 [1, 8, 4, 4] = inv(T_ego_cam): ``np.linalg.inv`` on float32 like ``onnx_smoke_test.py:46`` (the
  research goldens' feed, byte-identical; ``inverse="float64"`` gives the C++ runtime's double inverse cast to float).
- ``v0`` float32 [1]: the ego speed in m/s.
- An absent camera (the 7-camera rig, or a dropped slot) is an all-zero image plus the DONOR's K and pose
  (``bevlane/dataset.py:17-21``: CAM_BACK_NARROW <- CAM_BACK_WIDE, CAM_FRONT_NARROW <- CAM_FRONT_WIDE). With the
  released ``/255`` input a uint8 zero image is exactly the trained "absent" input; with ``input_norm="imagenet"`` the
  camera must be zeroed after normalising (``present`` mask, :func:`normalize_images`).

:func:`load_scene_frame` reads METEOR's demo-scene layout (``manifest.json``, ``img/NNNN_<CAM>.jpg``,
``ego_motion.npz``), which the public-dataset samples reuse; it reproduces ``research/meteor/scripts/meteor_io.py``'s
feed byte for byte (``tests/test_host_preprocess.py`` checks the sha256 of the shipped sample against its golden).
:func:`assemble_frame` is the API path (images of any size + calibration + speed).
"""
from __future__ import annotations

import hashlib
import json
import math
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Mapping, Optional, Sequence, Tuple

import numpy as np

from ..reference import config as C
from ..ttaw.image_area import meteor_resize, scale_intrinsics
from ..ttaw.io import load_image

__all__ = ["MeteorFrame", "assemble_frame", "load_scene_frame", "load_manifest", "scene_dir_of", "invert_extrinsics",
           "prepare_camera", "normalize_images", "pose_from_transform"]

F32 = np.float32


@dataclass
class MeteorFrame:
    """One network input (the ONNX feed) plus what the host post-processing needs."""

    imgs: np.ndarray                      # uint8 [1, 8, 3, 432, 768] RGB
    K: np.ndarray                         # float32 [1, 8, 3, 3] at 768x432
    T_cam_ego: np.ndarray                 # float32 [1, 8, 4, 4] ego -> camera
    v0: np.ndarray                        # float32 [1] m/s
    present: np.ndarray                   # bool [8]
    pose: Optional[Tuple[float, float, float]] = None   # ego (x, y, yaw) in a world frame (seg fusion, yaw tracks)
    meta: Dict[str, Any] = field(default_factory=dict)

    def feed(self) -> Dict[str, np.ndarray]:
        """The ONNX input dict (``imgs``, ``K``, ``T_cam_ego``, ``v0``)."""
        return {"imgs": self.imgs, "K": self.K, "T_cam_ego": self.T_cam_ego, "v0": self.v0}

    def sha256(self) -> Dict[str, str]:
        return {k: hashlib.sha256(np.ascontiguousarray(v).tobytes()).hexdigest() for k, v in self.feed().items()}


def invert_extrinsics(T_ego_cam: Any, inverse: str = "float32") -> np.ndarray:
    """camera -> ego (4x4) to the graph's ego -> camera, float32. ``float32``: ``np.linalg.inv`` of the float32 matrix
    (``onnx_smoke_test.py:46``, the research feed); ``float64``: double inverse cast to float (``realtime_main.cpp:
    123-128``)."""
    if inverse == "float32":
        return np.linalg.inv(np.asarray(T_ego_cam, F32).reshape(4, 4)).astype(F32)
    if inverse == "float64":
        return np.linalg.inv(np.asarray(T_ego_cam, np.float64).reshape(4, 4)).astype(F32)
    raise ValueError(f"inverse must be 'float32' or 'float64', not {inverse!r}")


def prepare_camera(image: Any, K: Optional[Any] = None) -> Tuple[np.ndarray, Optional[np.ndarray], Tuple[int, int]]:
    """One frame (anything ``ttaw.io.load_image`` takes; RGB) -> (uint8 [3, 432, 768] RGB CHW, K at 768x432 or None,
    the source (h, w)). INTER_AREA when the size differs (``ttaw.image_area``); K is scaled per axis."""
    rgb = load_image(image)
    src_hw = (int(rgb.shape[0]), int(rgb.shape[1]))
    if src_hw != (C.IMG_H, C.IMG_W):
        if src_hw[0] < C.IMG_H or src_hw[1] < C.IMG_W:
            raise ValueError(f"camera frame {src_hw[1]}x{src_hw[0]} is smaller than METEOR's 768x432 input "
                             "(INTER_AREA up-scaling is not supported)")
        rgb = meteor_resize(rgb)
    k = None if K is None else scale_intrinsics(K, src_hw)
    return np.ascontiguousarray(rgb.transpose(2, 0, 1)), k, src_hw


def pose_from_transform(T_world_from_ego: Any) -> Tuple[float, float, float]:
    """(x, y, yaw) of a 4x4 world <- ego transform (METEOR's ``ego_motion.npz`` ``pose`` convention)."""
    T = np.asarray(T_world_from_ego, np.float64).reshape(4, 4)
    return float(T[0, 3]), float(T[1, 3]), float(math.atan2(T[1, 0], T[0, 0]))


def assemble_frame(images: Mapping[str, Any], intrinsics: Mapping[str, Any], T_ego_cam: Mapping[str, Any],
                   v0: float, *, present: Optional[Mapping[str, bool]] = None, inverse: str = "float32",
                   pose: Optional[Tuple[float, float, float]] = None) -> MeteorFrame:
    """The API path: per camera name an image (any size; None = absent), the intrinsics of that image and its
    camera -> ego transform; ``v0`` in m/s. Absent cameras get a zero image and their donor's K and pose; a camera
    without a donor (only the two narrow cameras have one) must be present. A narrow camera sent as an all-zero image
    is absent too (METEOR's data convention, ``absent_zero.jpg``): its calibration is replaced by the donor's, as in
    training (``bevlane/dataset.py:131-145``)."""
    imgs = np.zeros((1, C.N_CAMS, 3, C.IMG_H, C.IMG_W), np.uint8)
    K = np.zeros((1, C.N_CAMS, 3, 3), F32)
    T = np.zeros((1, C.N_CAMS, 4, 4), F32)
    flags = np.zeros(C.N_CAMS, bool)
    src_sizes: Dict[str, Any] = {}
    cams: Dict[str, Tuple[np.ndarray, np.ndarray]] = {}
    for name in C.CAMERAS:
        img = images.get(name)
        ok = img is not None and (present is None or bool(present.get(name, True)))
        if not ok:
            continue
        chw, k, src_hw = prepare_camera(img, intrinsics.get(name))
        if name in C.DONOR and not chw.any():
            continue                  # an all-zero narrow camera is METEOR's "absent" convention: zero + donor pose
        if k is None or T_ego_cam.get(name) is None:
            raise ValueError(f"{name}: intrinsics and T_ego_cam are required for a present camera")
        i = C.CAMERAS.index(name)
        imgs[0, i] = chw
        flags[i] = True
        cams[name] = (k, np.asarray(T_ego_cam[name], np.float64).reshape(4, 4))
        src_sizes[name] = list(src_hw)
    donors: Dict[str, str] = {}
    for i, name in enumerate(C.CAMERAS):
        src = name if flags[i] else C.DONOR.get(name)
        if src is None or src not in cams:
            raise ValueError(f"{name} is absent and has no present donor camera (donors: {C.DONOR})")
        if src != name:
            donors[name] = src
        k, t = cams[src]
        K[0, i] = np.asarray(k, F32)
        T[0, i] = invert_extrinsics(t, inverse)
    return MeteorFrame(imgs, K, T, np.array([v0], F32), flags, pose,
                       {"source_hw": src_sizes, "donors": donors, "inverse": inverse})


def scene_dir_of(root: os.PathLike) -> Path:
    """``<root>/scenes.txt`` -> ``<root>/<first scene>`` (METEOR demo-root layout)."""
    root = Path(root)
    return root / (root / "scenes.txt").read_text().split()[0]


def load_manifest(scene_dir: os.PathLike) -> Dict[str, Any]:
    return json.loads((Path(scene_dir) / "manifest.json").read_text())


def load_scene_frame(scene_dir: os.PathLike, frame: int, manifest: Optional[Mapping[str, Any]] = None, *,
                     inverse: str = "float32") -> MeteorFrame:
    """Frame ``frame`` of a METEOR demo-scene directory (``manifest.json`` cams[c].K at 768x432 and T_ego_cam,
    frames[i].imgs[c]; ``ego_motion.npz`` v0 / pose / wp / valid), as ``meteor_io.load_frame`` feeds the graph."""
    scene_dir = Path(scene_dir)
    m = manifest or load_manifest(scene_dir)
    fr = m["frames"][frame]
    imgs = np.zeros((1, C.N_CAMS, 3, C.IMG_H, C.IMG_W), np.uint8)
    K = np.zeros((1, C.N_CAMS, 3, 3), F32)
    T = np.zeros((1, C.N_CAMS, 4, 4), F32)
    for i, cam in enumerate(C.CAMERAS):
        chw, _, _ = prepare_camera(scene_dir / fr["imgs"][cam])
        imgs[0, i] = chw
        K[0, i] = np.asarray(m["cams"][cam]["K"], F32)
        T[0, i] = invert_extrinsics(m["cams"][cam]["T_ego_cam"], inverse)
    present = np.asarray(m.get("present_mask", [1] * C.N_CAMS), bool).reshape(C.N_CAMS)
    ego = np.load(scene_dir / "ego_motion.npz")
    v0 = float(ego["v0"][frame])
    pose = tuple(float(v) for v in ego["pose"][frame]) if "pose" in ego.files else None
    meta: Dict[str, Any] = {"scene": m.get("scene", scene_dir.name), "frame": int(frame)}
    for key in ("wp", "valid"):
        if key in ego.files:
            meta[key] = np.asarray(ego[key][frame]).tolist()
    return MeteorFrame(imgs, K, T, np.array([v0], F32), present, pose, meta)


def normalize_images(imgs_u8: np.ndarray, mode: str = "onnx", present: Optional[Sequence[bool]] = None) -> np.ndarray:
    """Host oracle of the device input normalisation: uint8 [1, 8, 3, H, W] -> float32 [8, 3, H, W]; ``onnx`` = x / 255
    (as released), ``imagenet`` = (x / 255 - mean) / std with absent cameras zeroed afterwards (D12)."""
    x = np.asarray(imgs_u8).astype(F32) / F32(255.0)
    if mode == "imagenet":
        mean = np.asarray(C.IMAGENET_MEAN, F32).reshape(1, 1, 3, 1, 1)
        std = np.asarray(C.IMAGENET_STD, F32).reshape(1, 1, 3, 1, 1)
        x = (x - mean) / std
        if present is not None:
            x = x * np.asarray(present, F32).reshape(1, -1, 1, 1, 1)
    elif mode != "onnx":
        raise ValueError(f"input_norm must be one of {C.INPUT_NORMS}, not {mode!r}")
    return x.reshape(C.N_CAMS, 3, C.IMG_H, C.IMG_W)