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

Original authors:
Tommaso Comparin <tommaso.comparin@exact-lab.it>
Marco Franzon <marco.franzon@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.
"""
import json
import time
from pathlib import Path
from typing import Any
from typing import Dict
from typing import Iterable
from typing import Optional

import dask.array as da
import numpy as np
import zarr
from cellpose import core
from cellpose import models
from devtools import debug

import fractal_tasks_core
from .lib_pyramid_creation import write_pyramid
from .lib_zattrs_utils import rescale_datasets


__OME_NGFF_VERSION__ = fractal_tasks_core.__OME_NGFF_VERSION__


def image_labeling_whole_well(
    *,
    input_paths: Iterable[Path],
    output_path: Path,
    metadata: Optional[Dict[str, Any]] = None,
    component: str = None,
    labeling_channel: str = None,
    labeling_level: int = 2,
    diameter_level0: float = 80.0,
    cellprob_threshold: float = 0.0,
    flow_threshold: float = 0.4,
    model_type: str = "nuclei",
):

    """
    FIXME
    """

    # Preliminary check
    if len(input_paths) > 1:
        raise NotImplementedError
    in_path = input_paths[0]

    # Read some parameters from metadata
    num_levels = metadata["num_levels"]
    coarsening_xy = metadata["coarsening_xy"]
    chl_list = metadata["channel_list"]
    plate, well = component.split(".zarr/")

    if labeling_level >= num_levels:
        raise Exception(f"ERROR: {labeling_level=} but {num_levels=}")

    # Find well ID
    well_ID = well.replace("/", "_")[:-1]
    logfile = f"LOG_image_labeling_{well_ID}_whole_well"
    debug(well_ID)

    # Find channel index
    if labeling_channel not in chl_list:
        raise Exception(f"ERROR: {labeling_channel} not in {chl_list}")
    ind_channel = chl_list.index(labeling_channel)

    zarrurl = (in_path.parent.resolve() / component).as_posix() + "/"
    debug(zarrurl)

    # Check model_type
    if model_type not in ["nuclei", "cyto2", "cyto"]:
        raise Exception(f"ERROR model_type={model_type} is not allowed.")

    # FIXME: why is this check here?
    # Check that level>=1
    if labeling_level < 1:
        raise NotImplementedError("By now we can only segment levels >= 1")

    # Load ZYX data
    data_zyx = da.from_zarr(f"{zarrurl}{labeling_level}")[ind_channel]

    # Select 2D/3D behavior and set some parameters
    if data_zyx.shape[0] > 1:
        raise Exception(
            f"ERROR shape = {data_zyx.shape}"
            " but there can be only one Z plane."
        )

    # Load zattrs file
    zattrs_file = f"{zarrurl}.zattrs"
    with open(zattrs_file, "r") as jsonfile:
        zattrs = json.load(jsonfile)

    # Preliminary checks on multiscales
    multiscales = zattrs["multiscales"]
    if len(multiscales) > 1:
        raise Exception(f"ERROR: There are {len(multiscales)} multiscales")
    if "coordinateTransformations" in multiscales[0].keys():
        raise NotImplementedError(
            "global coordinateTransformations at the multiscales "
            "level are not currently supported"
        )

    # Extract axes, and remove channel
    new_axes = [ax for ax in multiscales[0]["axes"] if ax["type"] != "channel"]

    # Try to read channel label from OMERO metadata
    try:
        omero_label = zattrs["omero"]["channels"][ind_channel]["label"]
        label_name = f"label_{omero_label}"
    except (KeyError, IndexError):
        label_name = f"label_{ind_channel}"

    # Initialize cellpose
    use_gpu = core.use_gpu()
    model = models.Cellpose(gpu=use_gpu, model_type=model_type)

    # Initialize other things
    with open(logfile, "w") as out:
        out.write(f"Start image_labeling_whole_well task for {zarrurl}\n")
        out.write(f"model_type: {model_type}\n")
        out.write("Total well shape/chunks:\n")
        out.write(f"{data_zyx.shape}\n")
        out.write(f"{data_zyx.chunks}\n\n")

    # Write some debugging info
    with open(logfile, "a") as out:
        out.write(
            f"START Cellpose |" f" well: {type(data_zyx)}, {data_zyx.shape} \n"
        )

    # Actual labeling
    t0 = time.perf_counter()
    mask, flows, styles, diams = model.eval(
        data_zyx,
        channels=[0, 0],
        do_3D=False,
        net_avg=False,
        augment=False,
        diameter=(diameter_level0 / coarsening_xy**labeling_level),
        cellprob_threshold=cellprob_threshold,
        flow_threshold=flow_threshold,
    )
    mask = np.expand_dims(mask, axis=0)
    t1 = time.perf_counter()

    # Write some debugging info
    with open(logfile, "a") as out:
        out.write(
            f"END   Cellpose |"
            f" Elapsed: {t1-t0:.4f} seconds |"
            f" mask shape: {mask.shape},"
            f" mask dtype: {mask.dtype},"
            f" max(mask): {np.max(mask)}\n\n"
        )

    # Convert mask to dask
    mask_da = da.from_array(mask).rechunk((1, 2160, 2560))
    with open(logfile, "a") as out:
        out.write(
            f"da.from_array(upscaled_mask) [with rechunking]: {mask_da}\n\n"
        )

    # Rescale datasets (only relevant for labeling_level>0)
    new_datasets = rescale_datasets(
        datasets=multiscales[0]["datasets"],
        coarsening_xy=coarsening_xy,
        reference_level=labeling_level,
    )

    # Write zattrs for labels and for specific label
    # FIXME deal with: (1) many channels, (2) overwriting
    labels_group = zarr.group(f"{zarrurl}labels")
    labels_group.attrs["labels"] = [label_name]
    label_group = labels_group.create_group(label_name)
    label_group.attrs["image-label"] = {"version": __OME_NGFF_VERSION__}
    label_group.attrs["multiscales"] = [
        {
            "name": label_name,
            "version": __OME_NGFF_VERSION__,
            "axes": new_axes,
            "datasets": new_datasets,
        }
    ]

    # Construct resolution pyramid
    write_pyramid(
        mask_da,
        newzarrurl=f"{zarrurl}labels/{label_name}/",
        overwrite=False,
        coarsening_xy=coarsening_xy,
        num_levels=num_levels,
        chunk_size_x=2560,
        chunk_size_y=2160,
        aggregation_function=np.max,
    )


if __name__ == "__main__":
    from argparse import ArgumentParser

    parser = ArgumentParser(prog="image_labeling_whole_well.py")
    parser.add_argument(
        "-z",
        "--zarrurl",
        help="zarr url, at the merged-FOV level",
        required=True,
    )
    parser.add_argument(
        "-C",
        "--chl_list",
        nargs="+",
        help="list of channel names (e.g. A01_C01)",
    )
    parser.add_argument(
        "-cxy",
        "--coarsening_xy",
        default=2,
        type=int,
        help="coarsening factor along X and Y (optional, defaults to 2)",
    )
    parser.add_argument(
        "-lc",
        "--labeling_channel",
        help="name of channel for labeling (e.g. A01_C01)",
    )

    args = parser.parse_args()
    image_labeling_whole_well(
        args.zarrurl,
        coarsening_xy=args.coarsening_xy,
        chl_list=args.chl_list,
        labeling_channel=args.labeling_channel,
    )
