"""
Copyright 2022 (C)
    Friedrich Miescher Institute for Biomedical Research and
    University of Zurich

    Original authors:
    Tommaso Comparin <tommaso.comparin@exact-lab.it>

    This file is part of Fractal and was originally developed by eXact lab
    S.r.l.  <exact-lab.it> under contract with Liberali Lab from the Friedrich
    Miescher Institute for Biomedical Research and Pelkmans Lab from the
    University of Zurich.

Functions to use masked loading of ROIs before/after processing
"""
import logging
from pathlib import Path
from typing import Callable
from typing import Optional

import anndata as ad
import dask.array as da
import numpy as np
import zarr

from fractal_tasks_core.lib_upscale_array import convert_region_to_low_res
from fractal_tasks_core.lib_upscale_array import upscale_array

logger = logging.getLogger(__name__)


def _preprocess_input(
    image_array: np.ndarray,
    *,
    region: tuple[slice],
    current_label_path: str,
    ROI_table_path: str,
    ROI_positional_index: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Preprocess a four-dimensional cellpose input

    This involves :

    - Loading the masking label array for the appropriate ROI;
    - Extracting the appropriate label value from the ``ROI_table.obs``
      dataframe;
    - Constructing the background mask, where the masking label matches with a
      specific label value;
    - Setting the background of ``image_array`` to ``0``;
    - Loading the array which will be needed in postprocessing to restore
      background.

    **NOTE 1**: This function relies on a change to OME-NGFF table specs
    (https://github.com/ome/ngff/pull/64) which is still in-progress.

    **NOTE 2**: The pre/post-processing functions and the
    masked_loading_wrapper are currently meant to work as part of the
    cellpose_segmentation task, with the plan of then making them more
    flexible; see
    https://github.com/fractal-analytics-platform/fractal-tasks-core/issues/340.

    Naming of variables refers to a two-steps labeling, as in "first identify
    organoids, then look for nuclei inside each organoid") :

    - ``"masking"`` refers to the labels that are used to identify the object
      vs background (e.g. the organoid labels); these labels already exist.
    - ``"current"`` refers to the labels that are currently being computed in
      the ``cellpose_segmentation`` task, e.g. the nuclear labels.

    :param image_array: The 4D CZYX array with image data for a specific ROI.
    :param region: The ZYX indices of the ROI, in a form like ``(slice(0, 1),
                   slice(1000, 2000), slice(1000, 2000))``.
    :param current_label_path: Path to the image used as current
                               label, in a form like
                               ``/somewhere/plate.zarr/A/01/0/labels/nuclei_in_organoids/0``.
    :param ROI_table_path: Path of the AnnData table for the masking-label
                           ROIs; this is used (together with
                           ``ROI_positional_index``) to extract
                           ``label_value``.
    :param ROI_positional_index: Index of the current ROI, which is used to
                                 extract ``label_value`` from
                                 ``ROI_table_obs``.
    :returns: A tuple with three arrays: the preprocessed image array, the
              background mask, the current label.
    """

    logger.info(f"[_preprocess_input] {image_array.shape=}")
    logger.info(f"[_preprocess_input] {region=}")

    # Check that image data are 4D (CZYX) - FIXME issue 340
    if not image_array.ndim == 4:
        raise ValueError(
            "_preprocess_input requires a 4D "
            f"image_array argument, but {image_array.shape=}"
        )

    # Load the ROI table and its metadata attributes
    ROI_table = ad.read_zarr(ROI_table_path)
    attrs = zarr.group(ROI_table_path).attrs
    logger.info(f"[_preprocess_input] {ROI_table_path=}")
    logger.info(f"[_preprocess_input] {attrs.asdict()=}")
    if not attrs["type"] == "ngff:region_table":
        raise ValueError("Wrong attributes for {ROI_table_path}:\n{attrs}")
    label_relative_path = attrs["region"]["path"]
    column_name = attrs["instance_key"]

    # Check that ROI_table.obs has the right column and extract label_value
    if column_name not in ROI_table.obs.columns:
        raise ValueError(
            'In _preprocess_input, "{column_name}" '
            f" missing in {ROI_table.obs.columns=}"
        )
    label_value = int(ROI_table.obs[column_name][ROI_positional_index])

    # Load masking-label array (lazily)
    masking_label_path = str(
        Path(ROI_table_path).parent / label_relative_path / "0"
    )
    logger.critical(f"{masking_label_path=}")
    masking_label_array = da.from_zarr(masking_label_path)
    logger.info(
        f"[_preprocess_input] {masking_label_path=}, "
        f"{masking_label_array.shape=}"
    )

    # Load current-label array (lazily)
    current_label_array = da.from_zarr(current_label_path)
    logger.info(
        f"[_preprocess_input] {current_label_path=}, "
        f"{current_label_array.shape=}"
    )

    # Load ROI data for current label array
    current_label_region = current_label_array[region].compute()

    # Load ROI data for masking label array, with or without upscaling
    if masking_label_array.shape != current_label_array.shape:
        logger.info("Upscaling of masking label is needed")
        lowres_region = convert_region_to_low_res(
            highres_region=region,
            highres_shape=current_label_array.shape,
            lowres_shape=masking_label_array.shape,
        )
        masking_label_region = masking_label_array[lowres_region].compute()
        masking_label_region = upscale_array(
            array=masking_label_region,
            target_shape=current_label_region.shape,
        )
    else:
        masking_label_region = masking_label_array[region].compute()

    # Check that all shapes match
    shapes = (
        masking_label_region.shape,
        current_label_region.shape,
        image_array.shape[1:],
    )
    if len(set(shapes)) > 1:
        raise ValueError(
            "Shape mismatch:\n"
            f"{current_label_region.shape=}\n"
            f"{masking_label_region.shape=}\n"
            f"{image_array.shape=}"
        )

    # Compute background mask
    background_3D = masking_label_region != label_value
    if (masking_label_region == label_value).sum() == 0:
        raise ValueError(
            f"Label {label_value} is not present in the extracted ROI"
        )

    # Set image background to zero
    background_4D = np.expand_dims(background_3D, axis=0)
    image_array[background_4D] = 0

    return (image_array, background_3D, current_label_region)


def _postprocess_output(
    *,
    modified_array: np.ndarray,
    original_array: np.ndarray,
    background: np.ndarray,
) -> np.ndarray:
    """
    Postprocess cellpose output, mainly to restore its original background

    **NOTE**: The pre/post-processing functions and the
    masked_loading_wrapper are currently meant to work as part of the
    cellpose_segmentation task, with the plan of then making them more
    flexible; see
    https://github.com/fractal-analytics-platform/fractal-tasks-core/issues/340.

    :param modified_array: The 3D (ZYX) array with the correct object data and
                           wrong background data.
    :param original_array: The 3D (ZYX) array with the wrong object data and
                           correct background data.
    :param background: The 3D (ZYX) boolean array that defines the
                       background.
    :returns: The postprocessed array.
    """
    # Restore background
    modified_array[background] = original_array[background]
    return modified_array


def masked_loading_wrapper(
    *,
    function: Callable,
    image_array: np.ndarray,
    kwargs: Optional[dict] = None,
    use_masks: bool,
    preprocessing_kwargs: Optional[dict] = None,
):
    """
    Wrap a function with some pre/post-processing functions

    :param function: The callable function to be wrapped.
    :param image_array: The image array to be preprocessed and then used as
                        positional argument for ``function``.
    :param kwargs: Keyword arguments for ``function``.
    :param use_masks: If ``False``, the wrapper only calls ``function(*args,
                      **kwargs)``.
    :param preprocessing_kwargs: Keyword arguments for the preprocessing
                                 function (see call signature of
                                 ``_preprocess_input()``).
    """
    # Optional preprocessing
    if use_masks:
        preprocessing_kwargs = preprocessing_kwargs or {}
        (
            image_array,
            background_3D,
            current_label_region,
        ) = _preprocess_input(image_array, **preprocessing_kwargs)
    # Run function
    kwargs = kwargs or {}
    new_label_img = function(image_array, **kwargs)
    # Optional postprocessing
    if use_masks:
        new_label_img = _postprocess_output(
            modified_array=new_label_img,
            original_array=current_label_region,
            background=background_3D,
        )
    return new_label_img
