"""
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.

Helper functions to extract parameters from metadata dictionary
"""
from pathlib import Path
from typing import Any
from typing import Dict
from typing import Sequence
from typing import Union

import zarr


def find_omengff_acquisition(image_zarr_path: Path) -> Union[int, None]:
    """
    Discover the acquisition index based on OME-NGFF metadata

    Given the path to a zarr image folder (e.g. ``/path/plate.zarr/B/03/0``),
    extract the acquisition index from the ``.zattrs`` file of the parent
    folder (i.e. at the well level), or return ``None`` if acquisition is not
    specified.

    Notes:

    1. For non-multiplexing datasets, acquisition is not a required
       information in the metadata. If it is not there, this function
       returns ``None``.
    2. This function fails if we use an image that does not belong to
       an OME-NGFF well.

    :param image_zarr_path: full path to an OME-NGFF image folder
    """

    # Identify well plate and attrs
    well_zarr_path = image_zarr_path.parent
    if not (well_zarr_path / ".zattrs").exists():
        raise ValueError(
            f"{str(well_zarr_path)} must be an OME-NGFF plate "
            "folder, but it does not include a .zattrs file."
        )
    well_group = zarr.open_group(str(well_zarr_path))
    attrs_images = well_group.attrs["well"]["images"]

    # Loook for the acquisition of the current image (if any)
    acquisition = None
    for img_dict in attrs_images:
        if (
            img_dict["path"] == image_zarr_path.name
            and "acquisition" in img_dict.keys()
        ):
            acquisition = img_dict["acquisition"]
            break

    return acquisition


def get_parameters_from_metadata(
    *,
    keys: Sequence[str],
    metadata: Dict[str, Any],
    image_zarr_path: Path,
) -> Dict[str, Any]:
    """
    Flexibly extract parameters from metadata dictionary

    This covers both parameters which are acquisition-specific (if the image
    belongs to an OME-NGFF array and its acquisition is specified) or simply
    available in the dictionary.
    The two cases are handled as::

        metadata[acquisition]["some_parameter"]  # acquisition available
        metadata["some_parameter"]               # acquisition not available



    :param keys: list of required parameters
    :param metadata: metadata dictionary
    :param image_zarr_path: full path to image, e.g.
                             ``/path/plate.zarr/B/03/0``
    """

    parameters = {}
    acquisition = find_omengff_acquisition(image_zarr_path)
    if acquisition is not None:
        parameters["acquisition"] = acquisition

    for key in keys:
        if acquisition is None:
            parameter = metadata[key]
        else:
            try:
                parameter = metadata[key][str(acquisition)]
            except TypeError:
                parameter = metadata[key]
            except KeyError:
                parameter = metadata[key]
        parameters[key] = parameter
    return parameters
