# -*- coding: utf-8 -*-
# Coordinate the datasets, used to select the right dataset with corresponding setting
#
# @ Fabian Hörst, fabian.hoerst@uk-essen.de
# Institute for Artifical Intelligence in Medicine,
# University Medicine Essen

from typing import Callable

from torch.utils.data import Dataset
from cell_segmentation.datasets.conic import CoNicDataset

from cell_segmentation.datasets.pannuke import PanNukeDataset


def select_dataset(
    dataset_name: str, split: str, dataset_config: dict, transforms: Callable = None
) -> Dataset:
    """Select a cell segmentation dataset from the provided ones, currently just PanNuke is implemented here

    Args:
        dataset_name (str): Name of dataset to use.
            Must be one of: [pannuke, lizzard]
        split (str): Split to use.
            Must be one of: ["train", "val", "validation", "test"]
        dataset_config (dict): Dictionary with dataset configuration settings
        transforms (Callable, optional): PyTorch Image and Mask transformations. Defaults to None.

    Raises:
        NotImplementedError: Unknown dataset

    Returns:
        Dataset: Cell segmentation dataset
    """
    assert split.lower() in [
        "train",
        "val",
        "validation",
        "test",
    ], "Unknown split type!"

    if dataset_name.lower() == "pannuke":
        if split == "train":
            folds = dataset_config["train_folds"]
        if split == "val" or split == "validation":
            folds = dataset_config["val_folds"]
        if split == "test":
            folds = dataset_config["test_folds"]
        dataset = PanNukeDataset(
            dataset_path=dataset_config["dataset_path"],
            folds=folds,
            transforms=transforms,
            stardist=dataset_config.get("stardist", False),
            regression=dataset_config.get("regression_loss", False),
        )
    elif dataset_name.lower() == "conic":
        if split == "train":
            folds = dataset_config["train_folds"]
        if split == "val" or split == "validation":
            folds = dataset_config["val_folds"]
        if split == "test":
            folds = dataset_config["test_folds"]
        dataset = CoNicDataset(
            dataset_path=dataset_config["dataset_path"],
            folds=folds,
            transforms=transforms,
            stardist=dataset_config.get("stardist", False),
            regression=dataset_config.get("regression_loss", False),
            # TODO: Stardist and regression loss
        )
    else:
        raise NotImplementedError(f"Unknown dataset: {dataset_name}")
    return dataset