microsplit_data_module

def create_microsplit_predict_datamodule(
    pred_data: Union[str, Path, NDArray],
    tile_size: tuple,
    data_type: DataType,
    axes: str,
    batch_size: int = 1,
    num_channels: int = 2,
    depth3D: int = 1,
    grid_size: int | None = None,
    multiscale_count: int | None = None,
    data_stats: tuple | None = None,
    tiling_mode: TilingMode = TilingMode.ShiftBoundary,
    read_source_func: Callable | None = None,
    extension_filter: str = "",
    dataloader_params: dict | None = None,
    **dataset_kwargs,
) -> MicroSplitPredictDataModule:
    """
    Create a MicroSplitPredictDataModule for microSplit-style prediction datasets.

    Parameters
    ----------
    pred_data : str or Path or numpy.ndarray
        Prediction data, can be a path to a folder, a file or a numpy array.
    tile_size : tuple
        Size of one tile of data.
    data_type : DataType
        Type of the dataset (must be a DataType enum value).
    axes : str
        Axes of the data (e.g., 'SYX').
    batch_size : int, default=1
        Batch size for prediction dataloader.
    num_channels : int, default=2
        Number of channels in the input.
    depth3D : int, default=1
        Number of slices in 3D.
    grid_size : tuple, optional
        Grid size for patch extraction.
    multiscale_count : int, optional
        Number of LC scales.
    data_stats : tuple, optional
        Data statistics, by default None.
    tiling_mode : TilingMode, default=ShiftBoundary
        Tiling mode for patch extraction.
    read_source_func : Callable, optional
        Function to read the source data.
    extension_filter : str, optional
        File extension filter.
    dataloader_params : dict, optional
        Parameters for prediction dataloader.
    **dataset_kwargs :
        Additional arguments passed to MicroSplitDataConfig.

    Returns
    -------
    MicroSplitPredictDataModule
        Configured MicroSplitPredictDataModule instance.
    """
    if dataloader_params is None:
        dataloader_params = {}

    # Create prediction config with only valid parameters
    prediction_config_params = {
        "data_type": data_type,
        "image_size": tile_size,
        "num_channels": num_channels,
        "depth3D": depth3D,
        "grid_size": grid_size,
        "multiscale_lowres_count": multiscale_count,
        "data_stats": data_stats,
        "tiling_mode": tiling_mode,
        "batch_size": batch_size,
        "datasplit_type": DataSplitType.Test,  # For prediction, use all data
        **dataset_kwargs,
    }

    pred_config = MicroSplitDataConfig(**prediction_config_params)

    # Remove batch_size from dataloader_params if present
    if "batch_size" in dataloader_params:
        del dataloader_params["batch_size"]

    return MicroSplitPredictDataModule(
        pred_config=pred_config,
        pred_data=pred_data,
        read_source_func=(
            read_source_func if read_source_func is not None else get_train_val_data
        ),
        extension_filter=extension_filter,
        dataloader_params=dataloader_params,
    )

def create_microsplit_train_datamodule(
    train_data: str,
    patch_size: tuple,
    data_type: DataType,
    axes: str,  # TODO should be there after refactoring
    batch_size: int,
    val_data: str | None = None,
    num_channels: int = 2,
    depth3D: int = 1,
    grid_size: tuple | None = None,
    multiscale_count: int | None = None,
    tiling_mode: TilingMode = TilingMode.ShiftBoundary,
    read_source_func: Callable | None = None,  # TODO should be there after refactoring
    extension_filter: str = "",
    val_percentage: float = 0.1,
    val_minimum_split: int = 5,
    use_in_memory: bool = True,
    transforms: list | None = None,  # TODO should it be here?
    train_dataloader_params: dict | None = None,
    val_dataloader_params: dict | None = None,
    **dataset_kwargs,
) -> MicroSplitDataModule:
    """
    Create a MicroSplitDataModule for microSplit-style datasets.

    This includes config creation.

    Parameters
    ----------
    train_data : str
        Path to training data.
    patch_size : tuple
        Size of one patch of data.
    data_type : DataType
        Type of the dataset (must be a DataType enum value).
    axes : str
        Axes of the data (e.g., 'SYX').
    batch_size : int
        Batch size for dataloaders.
    val_data : str, optional
        Path to validation data.
    num_channels : int, default=2
        Number of channels in the input.
    depth3D : int, default=1
        Number of slices in 3D.
    grid_size : tuple, optional
        Grid size for patch extraction.
    multiscale_count : int, optional
        Number of LC scales.
    tiling_mode : TilingMode, default=ShiftBoundary
        Tiling mode for patch extraction.
    read_source_func : Callable, optional
        Function to read the source data.
    extension_filter : str, optional
        File extension filter.
    val_percentage : float, default=0.1
        Percentage of training data to use for validation.
    val_minimum_split : int, default=5
        Minimum number of patches/files for validation split.
    use_in_memory : bool, default=True
        Use in-memory dataset if possible.
    transforms : list, optional
        List of transforms to apply.
    train_dataloader_params : dict, optional
        Parameters for training dataloader.
    val_dataloader_params : dict, optional
        Parameters for validation dataloader.
    **dataset_kwargs :
        Additional arguments passed to DatasetConfig.

    Returns
    -------
    MicroSplitDataModule
        Configured MicroSplitDataModule instance.
    """
    # Create dataset configs with only valid parameters
    dataset_config_params = {
        "data_type": data_type,
        "image_size": patch_size,
        "num_channels": num_channels,
        "depth3D": depth3D,
        "grid_size": grid_size,
        "multiscale_lowres_count": multiscale_count,
        "tiling_mode": tiling_mode,
        "batch_size": batch_size,
        "train_dataloader_params": train_dataloader_params,
        "val_dataloader_params": val_dataloader_params,
        **dataset_kwargs,
    }

    train_config = MicroSplitDataConfig(
        **dataset_config_params,
        datasplit_type=DataSplitType.Train,
    )
    # val_config = MicroSplitDataConfig(
    #     **dataset_config_params,
    #     datasplit_type=DataSplitType.Val,
    # )
    # TODO, data config is duplicated here and in configuration

    return MicroSplitDataModule(
        data_config=train_config,
        train_data=train_data,
        val_data=val_data or train_data,
        train_data_target=None,
        val_data_target=None,
        read_source_func=get_train_val_data,  # Use our wrapped function
        extension_filter=extension_filter,
        val_percentage=val_percentage,
        val_minimum_split=val_minimum_split,
        use_in_memory=use_in_memory,
    )

def get_datasplit_tuples(val_fraction, test_fraction, data_length):
    """Get train/val/test indices for data splitting.

    Parameters
    ----------
    val_fraction : float or None
        Fraction of data to use for validation.
    test_fraction : float or None
        Fraction of data to use for testing.
    data_length : int
        Total length of the dataset.

    Returns
    -------
    tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]
        Training, validation, and test indices.
    """
    indices = np.arange(data_length)
    np.random.shuffle(indices)

    if val_fraction is None:
        val_fraction = 0.0
    if test_fraction is None:
        test_fraction = 0.0

    val_size = int(data_length * val_fraction)
    test_size = int(data_length * test_fraction)
    train_size = data_length - val_size - test_size

    train_idx = indices[:train_size]
    val_idx = indices[train_size : train_size + val_size]
    test_idx = indices[train_size + val_size :]

    return train_idx, val_idx, test_idx

def get_train_val_data(
    data_config,
    datadir,
    datasplit_type: DataSplitType,
    val_fraction=None,
    test_fraction=None,
    allow_generation=None,
    **kwargs,
):
    """Load and split data according to configuration.

    Parameters
    ----------
    data_config : MicroSplitDataConfig
        Data configuration object.
    datadir : str or Path
        Path to the data directory.
    datasplit_type : DataSplitType
        Type of data split to return.
    val_fraction : float, optional
        Fraction of data to use for validation.
    test_fraction : float, optional
        Fraction of data to use for testing.
    allow_generation : bool, optional
        Whether to allow data generation.
    **kwargs
        Additional keyword arguments.

    Returns
    -------
    numpy.ndarray
        Split data array.
    """
    data = load_data(datadir)
    train_idx, val_idx, test_idx = get_datasplit_tuples(
        val_fraction, test_fraction, len(data)
    )

    if datasplit_type == DataSplitType.All:
        data = data.astype(np.float64)
    elif datasplit_type == DataSplitType.Train:
        data = data[train_idx].astype(np.float64)
    elif datasplit_type == DataSplitType.Val:
        data = data[val_idx].astype(np.float64)
    elif datasplit_type == DataSplitType.Test:
        # TODO this is only used for prediction, and only because old dataset uses it
        data = data[test_idx].astype(np.float64)
    else:
        raise Exception("invalid datasplit")

    return data

def load_data(datadir):
    """Load data from a directory containing channel subdirectories with image files.

    Parameters
    ----------
    datadir : str or Path
        Path to the data directory containing channel subdirectories.

    Returns
    -------
    numpy.ndarray
        Stacked array of all channels' data.
    """
    data_path = Path(datadir)

    channel_dirs = sorted(p for p in data_path.iterdir() if p.is_dir())
    channels_data = []

    for channel_dir in channel_dirs:
        image_files = sorted(f for f in channel_dir.iterdir() if f.is_file())
        channel_images = [load_one_file(image_path) for image_path in image_files]

        channel_stack = np.concatenate(
            channel_images, axis=0
        )  # FIXME: this line works iff images have
        # a singleton channel dimension. Specify in the notebook or change with
        # `torch.stack`??
        channels_data.append(channel_stack)

    final_data = np.stack(channels_data, axis=-1)
    return final_data

def load_one_file(fpath):
    """Load a single 2D image file.

    Parameters
    ----------
    fpath : str or Path
        Path to the image file.

    Returns
    -------
    numpy.ndarray
        Reshaped image data.
    """
    data = tifffile.imread(fpath)
    if len(data.shape) == 2:
        axes = "YX"
    elif len(data.shape) == 3:
        axes = "SYX"
    elif len(data.shape) == 4:
        axes = "STYX"
    else:
        raise ValueError(f"Invalid data shape: {data.shape}")
    data = reshape_array(data, axes)
    data = data.reshape(-1, data.shape[-2], data.shape[-1])
    return data