projection utils

def _compute_new_shape(source_image: Image) -> tuple[tuple[int, ...], int]:
    """Compute the new shape of the image after the projection.

    The new shape is the same as the original one,
    except for the z-axis, which is set to 1.

    Returns:
        - new shape of the image
        - index of the z-axis in the original image
    """
    on_disk_shape = source_image.shape
    logger.info(f"Source {on_disk_shape=}")

    on_disk_z_index = source_image.axes_handler.get_index("z")
    if on_disk_z_index is None:
        raise ValueError(
            "The input image does not contain a z-axis, "
            "projection is only supported for 3D images with a z-axis."
        )

    dest_on_disk_shape = list(on_disk_shape)
    dest_on_disk_shape[on_disk_z_index] = 1
    logger.info(f"Destination {dest_on_disk_shape=}")
    return tuple(dest_on_disk_shape), on_disk_z_index

def max_wrapper(dask_array: da.Array, axis: int = 0) -> da.Array:
    """Perform a da.max on the dask_array.

    Args:
        dask_array (dask.array.Array): The input Dask array.
        axis (int, optional): The axis along which to max the array.
            Defaults to 0.

    Returns:
        dask.array.Array: The result of the max
    """
    return dask_array.max(axis=axis)

def mean_wrapper(dask_array: da.Array, axis: int = 0) -> da.Array:
    """Perform a da.mean on the dask_array & cast it to its original dtype.

    Without casting, the result can change dtype to e.g. float64

    Args:
        dask_array (dask.array.Array): The input Dask array.
        axis (int, optional): The axis along which to mean the array.
            Defaults to 0.

    Returns:
        dask.array.Array: The result of the mean, cast back to the original
            dtype.
    """
    # Determine the original dtype
    original_dtype = dask_array.dtype

    # Perform the mean
    result = da.mean(dask_array, axis=axis)

    # Cast back to the original dtype
    result = result.astype(original_dtype)

    return result

def min_wrapper(dask_array: da.Array, axis: int = 0) -> da.Array:
    """Perform a da.min on the dask_array.

    Args:
        dask_array (dask.array.Array): The input Dask array.
        axis (int, optional): The axis along which to min the array.
            Defaults to 0.

    Returns:
        dask.array.Array: The result of the min
    """
    return dask_array.min(axis=axis)

def projection_core(
    *,
    input_zarr_url: str,
    output_zarr_url: str,
    method: DaskProjectionMethod = DaskProjectionMethod.MIP,
    overwrite: bool = False,
    attributes: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Perform intensity projection along Z axis with a chosen method.

    Note: this task stores the output in a new zarr file.

    Args:
        input_zarr_url: Path or url to the individual OME-Zarr image to be processed.
        output_zarr_url: Path or url to the output OME-Zarr image.
        method: Projection method to be used. See `DaskProjectionMethod`
        overwrite: If `True`, overwrite the task output.
        attributes: Additional attributes to be added to the output image.
    """
    logger.info(f"{input_zarr_url=}")
    logger.info(f"{output_zarr_url=}")
    logger.info(f"{method=}")

    # Read image metadata
    original_ome_zarr = open_ome_zarr_container(input_zarr_url)
    original_image = original_ome_zarr.get_image()

    if original_image.is_2d or original_image.is_2d_time_series:
        raise ValueError(
            "The input image is 2D, projection is only supported for 3D images."
        )

    # Compute the new shape and pixel size
    dest_on_disk_shape, z_axis_index = _compute_new_shape(original_image)
    logger.info(f"New shape: {dest_on_disk_shape=}")

    # Create the new empty image
    ome_zarr_mip = original_ome_zarr.derive_image(
        store=output_zarr_url,
        name=method.value.upper(),
        shape=dest_on_disk_shape,
        pixelsize=original_image.pixel_size.yx,
        z_spacing=1.0,
        time_spacing=original_image.pixel_size.t,
        overwrite=overwrite,
        copy_labels=False,
        copy_tables=True,
    )
    logger.info(f"New Projection image created - {ome_zarr_mip=}")
    proj_image = ome_zarr_mip.get_image()

    # Process the image
    source_dask = original_image.get_as_dask()
    dest_dask = method.apply(dask_array=source_dask, axis=z_axis_index)
    dest_dask = da.expand_dims(dest_dask, axis=z_axis_index)
    proj_image.set_array(dest_dask)
    proj_image.consolidate()
    # Ends

    # Edit the roi tables
    for roi_table_name in ome_zarr_mip.list_roi_tables():
        table = ome_zarr_mip.get_generic_roi_table(roi_table_name)

        for roi in table.rois():
            old_z_slice = roi.get("z")
            if old_z_slice is not None:
                roi = roi.update_slice("z", (0, 1))
            table.add(roi, overwrite=True)

        table.consolidate()
        logger.info(f"Table {roi_table_name} Projection done")

    # Generate image_list_updates
    attributes = attributes or {}
    image_list_update_dict = {
        "image_list_updates": [
            {
                "zarr_url": output_zarr_url,
                "origin": input_zarr_url,
                "attributes": attributes,
                "types": {"is_3D": False},
            }
        ]
    }
    return image_list_update_dict

def safe_sum(dask_array: da.Array, axis: int = 0) -> da.Array:
    """Perform a safe sum on a Dask array to avoid overflow.

    Clips the result of da.sum & casts it to its original dtype.
    Dask.array already correctly handles promotion to uin32 or uint64 when
    necessary internally, but we want to ensure we clip the result.

    Args:
        dask_array (dask.array.Array): The input Dask array.
        axis (int, optional): The axis along which to sum the array.
            Defaults to 0.

    Returns:
        dask.array.Array: The result of the sum, safely clipped and cast
            back to the original dtype.
    """
    # Determine the original dtype
    original_dtype = dask_array.dtype
    if not np.issubdtype(original_dtype, np.integer):
        raise ValueError(
            f"safe_sum only supports integer dtypes, got {original_dtype}. "
            "Use a different projection method for float arrays."
        )
    max_value = np.iinfo(original_dtype).max

    # Perform the sum
    result = da.sum(dask_array, axis=axis)

    # Clip the values to the maximum possible value for the original dtype
    result = da.clip(result, 0, max_value)

    # Cast back to the original dtype
    result = result.astype(original_dtype)

    return result