"""MicroSplit data module for training and validation."""

from collections.abc import Callable
from pathlib import Path
from typing import Union

import numpy as np
import pytorch_lightning as L
import tifffile
from numpy.typing import NDArray
from torch.utils.data import DataLoader

from careamics.dataset.dataset_utils.dataset_utils import reshape_array
from careamics.lvae_training.dataset import (
    DataSplitType,
    LCMultiChDloader,
    MicroSplitDataConfig,
)
from careamics.lvae_training.dataset.types import TilingMode


# TODO refactor
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


# TODO refactor
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 if 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


# TODO refactor
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


# TODO refactor
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)
    return data


class MicroSplitDataModule(L.LightningDataModule):
    """Lightning DataModule for MicroSplit-style datasets.

    Matches the interface of TrainDataModule, but internally uses original MicroSplit
    dataset logic.

    Parameters
    ----------
    data_config : MicroSplitDataConfig
        Configuration for the MicroSplit dataset.
    train_data : str
        Path to training data directory.
    val_data : str, optional
        Path to validation data directory.
    train_data_target : str, optional
        Path to training target data.
    val_data_target : str, optional
        Path to validation target data.
    read_source_func : Callable, optional
        Function to read source data.
    extension_filter : str, optional
        File extension filter.
    val_percentage : float, optional
        Percentage of data to use for validation, by default 0.1.
    val_minimum_split : int, optional
        Minimum number of samples for validation split, by default 5.
    use_in_memory : bool, optional
        Whether to use in-memory dataset, by default True.
    """

    def __init__(
        self,
        data_config: MicroSplitDataConfig,
        train_data: str,
        val_data: str | None = None,
        train_data_target: str | None = None,
        val_data_target: str | None = None,
        read_source_func: Callable | None = None,
        extension_filter: str = "",
        val_percentage: float = 0.1,
        val_minimum_split: int = 5,
        use_in_memory: bool = True,
    ):
        """Initialize MicroSplitDataModule.

        Parameters
        ----------
        data_config : MicroSplitDataConfig
            Configuration for the MicroSplit dataset.
        train_data : str
            Path to training data directory.
        val_data : str, optional
            Path to validation data directory.
        train_data_target : str, optional
            Path to training target data.
        val_data_target : str, optional
            Path to validation target data.
        read_source_func : Callable, optional
            Function to read source data.
        extension_filter : str, optional
            File extension filter.
        val_percentage : float, optional
            Percentage of data to use for validation, by default 0.1.
        val_minimum_split : int, optional
            Minimum number of samples for validation split, by default 5.
        use_in_memory : bool, optional
            Whether to use in-memory dataset, by default True.
        """
        super().__init__()
        # Dataset selection logic (adapted from create_train_val_datasets)
        self.train_config = data_config  # SHould configs be separated?
        self.val_config = data_config
        self.test_config = data_config

        datapath = train_data
        load_data_func = read_source_func

        dataset_class = LCMultiChDloader  # TODO hardcoded for now

        # Create datasets
        self.train_dataset = dataset_class(
            self.train_config,
            datapath,
            load_data_fn=load_data_func,
            val_fraction=val_percentage,
            test_fraction=0.1,
        )
        max_val = self.train_dataset.get_max_val()
        self.val_config.max_val = max_val
        if self.train_config.datasplit_type == DataSplitType.All:
            self.val_config.datasplit_type = DataSplitType.All
            self.test_config.datasplit_type = DataSplitType.All
        self.val_dataset = dataset_class(
            self.val_config,
            datapath,
            load_data_fn=load_data_func,
            val_fraction=val_percentage,
            test_fraction=0.1,
        )
        self.test_config.max_val = max_val
        self.test_dataset = dataset_class(
            self.test_config,
            datapath,
            load_data_fn=load_data_func,
            val_fraction=val_percentage,
            test_fraction=0.1,
        )
        mean_val, std_val = self.train_dataset.compute_mean_std()
        self.train_dataset.set_mean_std(mean_val, std_val)
        self.val_dataset.set_mean_std(mean_val, std_val)
        self.test_dataset.set_mean_std(mean_val, std_val)
        data_stats = self.train_dataset.get_mean_std()

        # Store data statistics
        self.data_stats = (
            data_stats[0],
            data_stats[1],
        )  # TODO repeats old logic, revisit

    def train_dataloader(self):
        """Create a dataloader for training.

        Returns
        -------
        DataLoader
            Training dataloader.
        """
        return DataLoader(
            self.train_dataset,
            batch_size=self.train_config.batch_size,
            # TODO should be inside dataloader params?
            **self.train_config.train_dataloader_params,
        )

    def val_dataloader(self):
        """Create a dataloader for validation.

        Returns
        -------
        DataLoader
            Validation dataloader.
        """
        return DataLoader(
            self.val_dataset,
            batch_size=self.train_config.batch_size,
            **self.val_config.val_dataloader_params,  # TODO duplicated
        )

    def get_data_stats(self):
        """Get data statistics.

        Returns
        -------
        tuple[dict, dict]
            A tuple containing two dictionaries:
            - data_mean: mean values for input and target
            - data_std: standard deviation values for input and target
        """
        return self.data_stats, self.val_config.max_val  # TODO should be in the config?


def create_microsplit_train_datamodule(
    train_data: str,
    patch_size: tuple,
    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,
    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.

    Parameters
    ----------
    train_data : str
        Path to training data.
    patch_size : tuple
        Size of one patch of data.
    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.
    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": "HTLIF24Data",  # TODO remove, hardcoded for this example
        "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,
    )

    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,
    )


class MicroSplitPredictDataModule(L.LightningDataModule):
    """Lightning DataModule for MicroSplit-style prediction datasets.

    Matches the interface of PredictDataModule, but internally uses MicroSplit
    dataset logic for prediction.

    Parameters
    ----------
    pred_config : MicroSplitDataConfig
        Configuration for MicroSplit prediction.
    pred_data : str or Path or numpy.ndarray
        Prediction data, can be a path to a folder, a file or a numpy array.
    read_source_func : Callable, optional
        Function to read custom types.
    extension_filter : str, optional
        Filter to filter file extensions for custom types.
    dataloader_params : dict, optional
        Dataloader parameters.
    """

    def __init__(
        self,
        pred_config: MicroSplitDataConfig,
        pred_data: Union[str, Path, NDArray],
        read_source_func: Callable | None = None,
        extension_filter: str = "",
        dataloader_params: dict | None = None,
    ) -> None:
        """
        Constructor for MicroSplit prediction data module.

        Parameters
        ----------
        pred_config : MicroSplitDataConfig
            Configuration for MicroSplit prediction.
        pred_data : str or Path or numpy.ndarray
            Prediction data, can be a path to a folder, a file or a numpy array.
        read_source_func : Callable, optional
            Function to read custom types, by default None.
        extension_filter : str, optional
            Filter to filter file extensions for custom types, by default "".
        dataloader_params : dict, optional
            Dataloader parameters, by default {}.
        """
        super().__init__()

        if dataloader_params is None:
            dataloader_params = {}
        self.pred_config = pred_config
        self.pred_data = pred_data
        self.read_source_func = read_source_func or get_train_val_data
        self.extension_filter = extension_filter
        self.dataloader_params = dataloader_params
        self.predict_dataset = LCMultiChDloader(
            self.pred_config,
            self.pred_data,
            load_data_fn=self.read_source_func,
            val_fraction=0.0,  # No validation split for prediction
            test_fraction=1.0,  # No test split for prediction
        )
        self.predict_dataset.set_mean_std(*self.pred_config.data_stats)

    def predict_dataloader(self) -> DataLoader:
        """
        Create a dataloader for prediction.

        Returns
        -------
        DataLoader
            Prediction dataloader.
        """
        params = {**self.dataloader_params}
        params["shuffle"] = False
        return DataLoader(
            self.predict_dataset,
            batch_size=self.pred_config.batch_size,
            **params,
        )


def create_microsplit_predict_datamodule(
    pred_data: Union[str, Path, NDArray],
    tile_size: tuple,
    batch_size: int = 1,
    num_channels: int = 2,
    depth3D: int = 1,
    grid_size: tuple | 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.
    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": "HTLIF24Data",  # TODO remove, hardcoded for this example
        "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,
        "enable_random_cropping": False,
        "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,
    )
