from ..geometry import unproject_camera_depth
from ..geometry import GeometryFrame


def merge_frame_predictions(predictions):
    import numpy as np

    if not predictions:
        raise ValueError("DA3 returned no frame predictions")
    merged = predictions[0]
    for name in (
        "depth",
        "conf",
        "extrinsics",
        "intrinsics",
        "processed_images",
    ):
        setattr(
            merged,
            name,
            np.concatenate(
                [getattr(prediction, name) for prediction in predictions],
                axis=0,
            ),
        )
    return merged


def infer(
    model_manager,
    parameters,
    images,
    *,
    progress_callback=None,
    cancel_check=None,
):
    from .onnx_da3 import infer as infer_onnx
    from .onnx_runtime import provider_label

    model_directory = model_manager.directory(parameters["model"])
    device = parameters["device"]
    model = model_manager.get_da3(model_directory, device)
    options = {
        "process_res": parameters["process_res"],
        "process_res_method": parameters["process_res_method"],
        "ref_view_strategy": parameters["ref_view_strategy"],
    }

    def run(paths):
        try:
            return infer_onnx(model, paths, **options)
        except UnicodeDecodeError as error:
            if provider_label(model) == "DirectML":
                raise RuntimeError(
                    "DA3 inference failed on DirectML. Switch the device to "
                    "CPU in AnyImage Preferences and retry. "
                    f"Original error: {error}"
                ) from error
            raise

    if parameters.get("inference_mode", "joint") != "per_frame":
        return run([str(path) for path in images])

    predictions = []
    for index, image in enumerate(images):
        if cancel_check is not None:
            cancel_check()
        if progress_callback is not None:
            progress_callback(index, len(images))
        predictions.append(run([str(image)]))
    return merge_frame_predictions(predictions)


def surface_frame(prediction):
    import numpy as np

    depth = np.asarray(prediction.depth, dtype=np.float32)
    confidence = np.asarray(prediction.conf, dtype=np.float32)
    intrinsics = np.asarray(prediction.intrinsics, dtype=np.float32)
    if depth.shape[0] != 1:
        raise ValueError("Convert to Surface currently supports one image")
    points = unproject_camera_depth(depth[0], intrinsics[0])
    valid = (
        np.isfinite(depth[0])
        & (depth[0] > 0.0)
        & np.isfinite(points).all(axis=-1)
    )
    return GeometryFrame(
        depth=depth[0],
        points=points,
        intrinsics=intrinsics[0],
        valid_mask=valid,
        confidence=confidence[0],
    )