"""Result payload lookup and coherent-GPR objects."""

from __future__ import annotations

import numpy as np

from python_app.models.dataset_model import ResultCollection, ResultPayload


def collection_payload_by_name(
    collection: ResultCollection,
    name: str,
    kind: int | None = None,
) -> ResultPayload | None:
    """Return the first collection payload matching name and optional kind."""
    for payload in collection.collection_payloads:
        if payload.processing_name != name:
            continue
        if kind is not None and int(payload.kind) != int(kind):
            continue
        return payload
    return None


def gpr_object_rows(collection: ResultCollection) -> np.ndarray:
    """Return object rows as `[x_m, z_m, score]` from a GPR collection."""
    points_payload = collection_payload_by_name(collection, "gpr_points", kind=4)
    if points_payload is not None:
        points = np.asarray(points_payload.table, dtype=np.float32)
        if points.ndim == 2 and points.shape[1] >= 3:
            return points[:, :3]

    return np.zeros((0, 3), dtype=np.float32)