# SPDX-License-Identifier: Apache-2.0
"""C08 outputs: the result classes of the 13 bundles and their ``POST /predict`` JSON (BUNDLE_CONVENTIONS.md 7.4).

Every class has ``to_dicts()`` (one plain dict per detection / pose / class) and ``to_dict(output_format)`` (the
whole response body: ``model``, ``frame_id``, the payload, ``meta``, ``timing_ms``). The Python API returns these
objects and the server returns ``to_dict(...)`` of the same object, so both agree bit for bit.

==================  =================================  ==================================================
class               families                           payload
==================  =================================  ==================================================
``Detections3D``    CenterPoint, TransFusion, BEVFusion, ``detections`` [label, label_id, score, center, size,
                    PointPainting, StreamPETR, BEVDet,  yaw, velocity], ``num_detections``; rows sorted by
                    BEVFormer, PTv3-det                 descending score at construction
``Detections2D``    YOLOX                              ``detections`` [label, label_id, score, box_xyxy] +
                                                       ``extras`` (e.g. ``{"semseg": Mask2D}``)
``Segmentation3D``  FRNet, PTv3-seg                    ``labels`` (npz uint8/uint16 [N]), ``scores``,
                                                       ``class_names``, ``class_counts``
``Mask2D``          SceneSeg (YOLOX semseg extra)      ``mask`` (png or npz), ``class_names``, ``class_counts``
``Trajectory``      Diffusion Planner                  ``trajectory`` [[x, y, yaw, ...] x T], ``columns``,
                                                       ``turn_indicator``, ``predicted_agents``
==================  =================================  ==================================================
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Sequence

import numpy as np

from .io import encode_array, encode_png, to_jsonable

__all__ = ["Detections3D", "Detections2D", "Segmentation3D", "Mask2D", "Trajectory", "label_name"]


def label_name(labels: Sequence[str], label_id: int) -> str:
    """``labels[label_id]``, or the id as text when it is out of range."""
    return labels[label_id] if 0 <= label_id < len(labels) else str(label_id)


def _envelope(model: str, frame_id: str, meta: dict, timing_ms: dict) -> Dict[str, Any]:
    return {"model": model, "frame_id": frame_id, "meta": to_jsonable(meta),
            "timing_ms": {k: round(float(v), 3) for k, v in timing_ms.items()}}


def _sorted_by_score(scores: np.ndarray) -> np.ndarray:
    return np.argsort(-scores, kind="stable")


@dataclass
class Detections3D:
    """3D boxes in ``frame_id`` (Autoware ``base_link``: x forward, y left, z up; metres, radians).

    ``boxes`` float32 ``[N, 7]`` = x, y, z (box centre), length, width, height, yaw; ``scores`` ``[N]``;
    ``label_ids`` int32 ``[N]`` (index into ``labels``); ``velocities`` ``[N, 2]`` (vx, vy) m/s or None.
    Rows are sorted by descending score (stable) at construction."""

    boxes: np.ndarray
    scores: np.ndarray
    label_ids: np.ndarray
    velocities: Optional[np.ndarray] = None
    labels: Sequence[str] = ()
    model: str = ""
    frame_id: str = "base_link"
    timing_ms: dict = field(default_factory=dict)
    meta: dict = field(default_factory=dict)

    def __post_init__(self) -> None:
        self.boxes = np.asarray(self.boxes, np.float32).reshape(-1, 7)
        self.scores = np.asarray(self.scores, np.float32).reshape(-1)
        self.label_ids = np.asarray(self.label_ids, np.int32).reshape(-1)
        n = len(self.scores)
        if len(self.boxes) != n or len(self.label_ids) != n:
            raise ValueError(f"boxes {self.boxes.shape}, scores {self.scores.shape}, labels {self.label_ids.shape}")
        if self.velocities is not None:
            self.velocities = np.asarray(self.velocities, np.float32).reshape(n, 2)
        order = _sorted_by_score(self.scores)
        self.boxes, self.scores, self.label_ids = self.boxes[order], self.scores[order], self.label_ids[order]
        if self.velocities is not None:
            self.velocities = self.velocities[order]

    def __len__(self) -> int:
        return int(self.scores.shape[0])

    @property
    def label_names(self) -> List[str]:
        return [label_name(self.labels, i) for i in self.label_ids.tolist()]

    def to_dicts(self) -> List[Dict[str, Any]]:
        out = []
        for i, name in enumerate(self.label_names):
            x, y, z, length, width, height, yaw = (float(v) for v in self.boxes[i])
            d = {"label": name, "label_id": int(self.label_ids[i]), "score": round(float(self.scores[i]), 4),
                 "center": [round(x, 3), round(y, 3), round(z, 3)],
                 "size": [round(length, 3), round(width, 3), round(height, 3)], "yaw": round(yaw, 4)}
            if self.velocities is not None:
                d["velocity"] = [round(float(v), 3) for v in self.velocities[i]]
            out.append(d)
        return out

    def to_dict(self, output_format: str = "json") -> Dict[str, Any]:
        body = _envelope(self.model, self.frame_id, self.meta, self.timing_ms)
        body.update(num_detections=len(self), detections=self.to_dicts())
        if output_format == "npz":  # lossless arrays for programmatic clients
            body["arrays"] = {"boxes": encode_array(self.boxes, key="boxes"),
                              "scores": encode_array(self.scores, key="scores"),
                              "label_ids": encode_array(self.label_ids, key="label_ids")}
            if self.velocities is not None:
                body["arrays"]["velocities"] = encode_array(self.velocities, key="velocities")
        return body


@dataclass
class Mask2D:
    """A per-pixel class map (H, W) (uint8 for <= 256 classes) in the source image's pixel grid."""

    mask: np.ndarray
    class_names: Sequence[str] = ()
    model: str = ""
    frame_id: str = "camera"
    encoding: str = "png"
    timing_ms: dict = field(default_factory=dict)
    meta: dict = field(default_factory=dict)

    def __post_init__(self) -> None:
        self.mask = np.asarray(self.mask)
        if self.mask.ndim != 2:
            raise ValueError(f"mask must be (H, W), got {self.mask.shape}")
        if self.encoding not in ("png", "npz"):
            raise ValueError("encoding must be 'png' or 'npz'")

    def class_counts(self) -> Dict[str, int]:
        ids, counts = np.unique(self.mask, return_counts=True)
        return {label_name(self.class_names, int(i)): int(c) for i, c in zip(ids, counts)}

    def to_dicts(self) -> List[Dict[str, Any]]:
        ids, counts = np.unique(self.mask, return_counts=True)
        return [{"label": label_name(self.class_names, int(i)), "label_id": int(i), "pixels": int(c)}
                for i, c in zip(ids, counts)]

    def payload(self, output_format: str = "json") -> Dict[str, Any]:
        """The encoded mask alone (also used when the mask is an extra of another output)."""
        if self.encoding == "png" and output_format != "npz" and self.mask.dtype == np.uint8:
            return encode_png(self.mask, key="mask")
        return encode_array(self.mask, key="mask")

    def to_dict(self, output_format: str = "json") -> Dict[str, Any]:
        body = _envelope(self.model, self.frame_id, self.meta, self.timing_ms)
        body.update(mask=self.payload(output_format), class_names=list(self.class_names),
                    class_counts=self.class_counts())
        return body


@dataclass
class Detections2D:
    """2-D boxes in original image pixels: ``boxes_xyxy`` float32 ``[N, 4]``, ``scores``, ``label_ids``.
    ``extras`` holds companion outputs encoded into the body by name (``{"semseg": Mask2D(...)}``).
    Rows are sorted by descending score (stable) at construction."""

    boxes_xyxy: np.ndarray
    scores: np.ndarray
    label_ids: np.ndarray
    labels: Sequence[str] = ()
    extras: Dict[str, Any] = field(default_factory=dict)
    model: str = ""
    frame_id: str = "camera"
    timing_ms: dict = field(default_factory=dict)
    meta: dict = field(default_factory=dict)

    def __post_init__(self) -> None:
        self.boxes_xyxy = np.asarray(self.boxes_xyxy, np.float32).reshape(-1, 4)
        self.scores = np.asarray(self.scores, np.float32).reshape(-1)
        self.label_ids = np.asarray(self.label_ids, np.int32).reshape(-1)
        if not (len(self.boxes_xyxy) == len(self.scores) == len(self.label_ids)):
            raise ValueError("boxes_xyxy, scores and label_ids lengths differ")
        order = _sorted_by_score(self.scores)
        self.boxes_xyxy, self.scores, self.label_ids = self.boxes_xyxy[order], self.scores[order], self.label_ids[order]

    def __len__(self) -> int:
        return int(self.scores.shape[0])

    def to_dicts(self) -> List[Dict[str, Any]]:
        return [{"label": label_name(self.labels, int(self.label_ids[i])), "label_id": int(self.label_ids[i]),
                 "score": round(float(self.scores[i]), 4),
                 "box_xyxy": [round(float(v), 2) for v in self.boxes_xyxy[i]]} for i in range(len(self))]

    def to_dict(self, output_format: str = "json") -> Dict[str, Any]:
        body = _envelope(self.model, self.frame_id, self.meta, self.timing_ms)
        body.update(num_detections=len(self), detections=self.to_dicts())
        for name, extra in self.extras.items():
            if hasattr(extra, "payload"):
                body[name] = extra.payload(output_format)
            elif isinstance(extra, np.ndarray):
                body[name] = encode_array(extra, key=name)
            else:
                body[name] = to_jsonable(extra)
        if output_format == "npz":
            body["arrays"] = {"boxes_xyxy": encode_array(self.boxes_xyxy, key="boxes_xyxy"),
                              "scores": encode_array(self.scores, key="scores"),
                              "label_ids": encode_array(self.label_ids, key="label_ids")}
        return body


@dataclass
class Segmentation3D:
    """Per-point classes, in input point order after NaN removal (stated in SERVING.md). ``label_ids`` ``[N]``;
    ``scores`` optional ``[N]`` (winning-class probability) or ``[N, C]``."""

    label_ids: np.ndarray
    class_names: Sequence[str] = ()
    scores: Optional[np.ndarray] = None
    model: str = ""
    frame_id: str = "base_link"
    timing_ms: dict = field(default_factory=dict)
    meta: dict = field(default_factory=dict)

    def __post_init__(self) -> None:
        ids = np.asarray(self.label_ids).reshape(-1)
        if ids.size and ids.min() < 0:
            raise ValueError("label ids must be >= 0")
        self.label_ids = ids.astype(np.uint8 if (not ids.size or ids.max() < 256) else np.uint16)
        if self.scores is not None:
            self.scores = np.asarray(self.scores, np.float32)
            if self.scores.shape[0] != ids.shape[0]:
                raise ValueError("scores and label_ids lengths differ")

    def __len__(self) -> int:
        return int(self.label_ids.shape[0])

    def class_counts(self) -> Dict[str, int]:
        counts = np.bincount(self.label_ids.astype(np.int64), minlength=len(self.class_names))
        return {label_name(self.class_names, i): int(c) for i, c in enumerate(counts) if c or i < len(self.class_names)}

    def to_dicts(self) -> List[Dict[str, Any]]:
        counts = np.bincount(self.label_ids.astype(np.int64), minlength=len(self.class_names))
        return [{"label": label_name(self.class_names, i), "label_id": i, "points": int(c)}
                for i, c in enumerate(counts)]

    def to_dict(self, output_format: str = "json") -> Dict[str, Any]:
        body = _envelope(self.model, self.frame_id, self.meta, self.timing_ms)
        body.update(num_points=len(self), labels=encode_array(self.label_ids, key="labels"),
                    class_names=list(self.class_names), class_counts=self.class_counts())
        if self.scores is not None:
            body["scores"] = encode_array(self.scores, key="scores")
        return body


@dataclass
class Trajectory:
    """A planned trajectory ``poses`` ``[T, D]`` whose columns are named by ``columns`` (default x, y, yaw) in
    ``frame_id``; optional ``turn_indicator`` and ``predicted_agents`` ``[A, T, D']``."""

    poses: np.ndarray
    columns: Sequence[str] = ("x", "y", "yaw")
    turn_indicator: Any = None
    predicted_agents: Optional[np.ndarray] = None
    model: str = ""
    frame_id: str = "base_link"
    timing_ms: dict = field(default_factory=dict)
    meta: dict = field(default_factory=dict)

    def __post_init__(self) -> None:
        self.poses = np.asarray(self.poses, np.float32)
        if self.poses.ndim != 2 or self.poses.shape[1] != len(self.columns):
            raise ValueError(f"poses {self.poses.shape} do not match columns {list(self.columns)}")
        if self.predicted_agents is not None:
            self.predicted_agents = np.asarray(self.predicted_agents, np.float32)

    def __len__(self) -> int:
        return int(self.poses.shape[0])

    def to_dicts(self) -> List[Dict[str, Any]]:
        return [{c: round(float(v), 4) for c, v in zip(self.columns, row)} for row in self.poses]

    def to_dict(self, output_format: str = "json") -> Dict[str, Any]:
        body = _envelope(self.model, self.frame_id, self.meta, self.timing_ms)
        body.update(num_poses=len(self), columns=list(self.columns),
                    trajectory=[[round(float(v), 4) for v in row] for row in self.poses],
                    turn_indicator=to_jsonable(self.turn_indicator))
        if self.predicted_agents is not None:
            body["predicted_agents"] = encode_array(self.predicted_agents, key="predicted_agents")
        if output_format == "npz":
            body["arrays"] = {"poses": encode_array(self.poses, key="poses")}
        return body