"""
A place for Datasets and Dataloaders.
"""

from collections import defaultdict
from functools import cache
from pathlib import Path
from typing import Callable, Union

import numpy as np
from skimage.transform import resize

from .config import MicroSplitDataConfig
from .types import DataSplitType, TilingMode
from .utils.empty_patch_fetcher import EmptyPatchFetcher
from .utils.index_manager import GridIndexManagerRef


class MultiChDloaderRef:
    def __init__(
        self,
        data_config: MicroSplitDataConfig,
        fpath: str,
        load_data_fn: Callable,
        val_fraction: float = None,
        test_fraction: float = None,
    ):
        """ """
        self._data_type = data_config.data_type
        self._fpath = Path(fpath)
        self._data = None
        self._3Ddata = False  # TODO wtf it was 5D
        self._tiling_mode = data_config.tiling_mode
        # by default, if the noise is present, add it to the input and target.
        self._depth3D = data_config.depth3D
        self._mode_3D = data_config.mode_3D
        # NOTE: Input is the sum of the different channels. It is not the average of the different channels.
        self._input_is_sum = data_config.input_is_sum
        self._num_channels = data_config.num_channels
        self._input_idx = data_config.input_idx
        self._tar_idx_list = data_config.target_idx_list

        self.load_data(
            data_config,
            data_config.datasplit_type,
            load_data_fn=load_data_fn,
            val_fraction=val_fraction,
            test_fraction=test_fraction,
            allow_generation=data_config.allow_generation,
        )

        self._data_shapes = self.get_data_shapes()
        self._normalized_input = data_config.normalized_input
        self._quantile = 1.0
        self._channelwise_quantile = False
        self._background_quantile = 0.0
        self._clip_background_noise_to_zero = False
        self._skip_normalization_using_mean = False
        self._empty_patch_replacement_enabled = False

        self._background_values = None

        self._overlapping_padding_kwargs = data_config.overlapping_padding_kwargs
        if self._tiling_mode in [TilingMode.TrimBoundary, TilingMode.ShiftBoundary]:
            if (
                self._overlapping_padding_kwargs is None
                or data_config.multiscale_lowres_count is not None
            ):
                # raise warning
                print("Padding is not used with this alignement style")
        else:
            assert (
                self._overlapping_padding_kwargs is not None
            ), "When not trimming boudnary, padding is needed."

        self._is_train = data_config.datasplit_type == DataSplitType.Train

        # input = alpha * ch1 + (1-alpha)*ch2.
        # alpha is sampled randomly between these two extremes
        self._start_alpha_arr = self._end_alpha_arr = self._return_alpha = None

        self._img_sz = self._grid_sz = self._repeat_factor = self.idx_manager = None

        # changed set_img_sz because "grid_size" in data_config returns false
        try:
            grid_size = data_config.grid_size
        except AttributeError:
            grid_size = data_config.image_size

        if self._is_train:
            self._start_alpha_arr = data_config.start_alpha  # TODO why only for train?
            self._end_alpha_arr = data_config.end_alpha

        self.set_img_sz(data_config.image_size, grid_size)

        self._empty_patch_replacement_enabled = (
            data_config.empty_patch_replacement_enabled and self._is_train
        )
        if self._empty_patch_replacement_enabled:
            self._empty_patch_replacement_channel_idx = (
                data_config.empty_patch_replacement_channel_idx
            )
            self._empty_patch_replacement_probab = (
                data_config.empty_patch_replacement_probab
            )
            data_frames = self._data[..., self._empty_patch_replacement_channel_idx]
            # NOTE: This is on the raw data. So, it must be called before removing the background.
            self._empty_patch_fetcher = EmptyPatchFetcher(
                self.idx_manager,
                self._img_sz,
                data_frames,
                max_val_threshold=data_config.empty_patch_max_val_threshold,
            )

        self.rm_bkground_set_max_val_and_upperclip_data(
            data_config.max_val, data_config.datasplit_type
        )

        # For overlapping dloader, image_size and repeat_factors are not related. hence a different function.

        self._mean = None
        self._std = None
        self._use_one_mu_std = data_config.use_one_mu_std

        self._target_separate_normalization = data_config.target_separate_normalization

        self._enable_rotation = data_config.enable_rotation_aug
        flipz_3D = data_config.random_flip_z_3D
        self._flipz_3D = flipz_3D and self._enable_rotation

        self._enable_random_cropping = data_config.enable_random_cropping
        self._uncorrelated_channels = (
            data_config.uncorrelated_channels and self._is_train
        )
        self._uncorrelated_channel_probab = data_config.uncorrelated_channel_probab
        assert self._is_train or self._uncorrelated_channels is False
        assert (
            self._enable_random_cropping is True or self._uncorrelated_channels is False
        )
        # Randomly rotate [-90,90]

        self._rotation_transform = None
        if self._enable_rotation:
            # TODO: fix this import
            import albumentations as A

            self._rotation_transform = A.Compose([A.Flip(), A.RandomRotate90()])

        # TODO: remove print log messages
        # if print_vars:
        #     msg = self._init_msg()
        #     print(msg)

    def get_data_shapes(self):
        if self._3Ddata:  # TODO we assume images don't have a channel dimension
            [
                [
                    im.shape if len(im.shape) == 4 else (1, *im.shape)
                    for im in self._data[ch]
                ]
                for ch in range(len(self._data))
            ]
        else:
            return [
                [
                    im.shape if len(im.shape) == 3 else (1, *im.shape)
                    for im in self._data[ch]
                ]
                for ch in range(len(self._data))
            ]

    def load_data(
        self,
        data_config: MicroSplitDataConfig,
        datasplit_type: DataSplitType,
        load_data_fn: Callable,
        val_fraction=None,
        test_fraction=None,
        allow_generation=None,
    ):
        self._data = load_data_fn(
            data_config,
            self._fpath,
            datasplit_type,
            val_fraction=val_fraction,
            test_fraction=test_fraction,
            allow_generation=allow_generation,
        )

    # TODO check for 2D/3D data consistency with config
    # TODO check number of channels consistency with config

    def save_background(self, channel_idx, frame_idx, background_value):
        self._background_values[frame_idx, channel_idx] = background_value

    def get_background(self, channel_idx, frame_idx):
        return self._background_values[frame_idx, channel_idx]

    def rm_bkground_set_max_val_and_upperclip_data(self, max_val, datasplit_type):
        # self.remove_background() # TODO revisit
        self.set_max_val(max_val, datasplit_type)
        self.upperclip_data()

    def upperclip_data(self):
        for ch_idx, data in enumerate(self._data):
            if self.max_val[ch_idx] is not None:
                for idx in range(len(data)):
                    data[idx][data[idx] > self.max_val[ch_idx]] = self.max_val[ch_idx]

    def compute_max_val(self):
        # TODO add channelwise quantile ?
        return [
            max([np.quantile(im, self._quantile) for im in ch]) for ch in self._data
        ]

    def set_max_val(self, max_val, datasplit_type):
        if max_val is None:
            assert datasplit_type in [DataSplitType.Train, DataSplitType.All]
            self.max_val = self.compute_max_val()
        else:
            assert max_val is not None
            self.max_val = max_val

    def get_max_val(self):
        return self.max_val

    def get_img_sz(self):
        return self._img_sz

    def get_num_frames(self):
        """Returns the number of the longest channel."""
        return max(self.idx_manager.total_grid_count()[0])

    def reduce_data(
        self,
        t_list=None,
        z_start=None,
        z_end=None,
        h_start=None,
        h_end=None,
        w_start=None,
        w_end=None,
    ):
        raise NotImplementedError("Not implemented")

    def get_idx_manager_shapes(
        self, patch_size: int, grid_size: Union[int, tuple[int, int, int]]
    ):
        numC = len(self._data_shapes)
        if self._3Ddata:
            patch_shape = (1, self._depth3D, patch_size, patch_size)
            if isinstance(grid_size, int):
                grid_shape = (1, 1, grid_size, grid_size)
            else:
                assert len(grid_size) == 3
                assert all(
                    [g <= p for g, p in zip(grid_size, patch_shape[1:-1])]
                ), f"Grid size {grid_size} must be less than patch size {patch_shape[1:-1]}"
                grid_shape = (1, grid_size[0], grid_size[1], grid_size[2])
        else:
            assert isinstance(grid_size, int)
            grid_shape = (1, grid_size, grid_size)
            patch_shape = (1, patch_size, patch_size)

        return patch_shape, grid_shape

    def set_img_sz(self, image_size, grid_size: Union[int, tuple[int, int, int]]):
        """
        If one wants to change the image size on the go, then this can be used.
        Args:
            image_size: size of one patch
            grid_size: frame is divided into square grids of this size. A patch centered on a grid having size `image_size` is returned.
        """
        # hacky way to deal with image shape from new conf
        self._img_sz = image_size[-1]  # TODO revisit!
        self._grid_sz = grid_size
        shapes = self._data_shapes

        patch_shape, grid_shape = self.get_idx_manager_shapes(
            self._img_sz, self._grid_sz
        )
        self.idx_manager = GridIndexManagerRef(
            shapes, grid_shape, patch_shape, self._tiling_mode
        )

    def __len__(self):
        # If channel length is not equal, return the longest
        return max(self.idx_manager.total_grid_count()[0])

    def _init_msg(
        self,
    ):
        msg = (
            f"[{self.__class__.__name__}] Train:{int(self._is_train)} Sz:{self._img_sz}"
        )
        dim_sizes = [
            self.idx_manager.get_individual_dim_grid_count(dim)
            for dim in range(len(self._data.shape))
        ]
        dim_sizes = ",".join([str(x) for x in dim_sizes])
        msg += f" N:{self.N} NumPatchPerN:{self._repeat_factor}"
        msg += f"{self.idx_manager.total_grid_count()} DimSz:({dim_sizes})"
        msg += f" TrimB:{self._tiling_mode}"
        # msg += f' NormInp:{self._normalized_input}'
        # msg += f' SingleNorm:{self._use_one_mu_std}'
        msg += f" Rot:{self._enable_rotation}"
        if self._flipz_3D:
            msg += f" FlipZ:{self._flipz_3D}"

        msg += f" RandCrop:{self._enable_random_cropping}"
        msg += f" Channel:{self._num_channels}"
        # msg += f' Q:{self._quantile}'
        if self._input_is_sum:
            msg += f" SummedInput:{self._input_is_sum}"

        if self._empty_patch_replacement_enabled:
            msg += f" ReplaceWithRandSample:{self._empty_patch_replacement_enabled}"
        if self._uncorrelated_channels:
            msg += f" Uncorr:{self._uncorrelated_channels}"
        if self._empty_patch_replacement_enabled:
            msg += f"-{self._empty_patch_replacement_channel_idx}-{self._empty_patch_replacement_probab}"
        if self._background_quantile > 0.0:
            msg += f" BckQ:{self._background_quantile}"

        if self._start_alpha_arr is not None:
            msg += f" Alpha:[{self._start_alpha_arr},{self._end_alpha_arr}]"
        return msg

    def _crop_imgs(self, ch_idx: int, patch_idx: int, img: np.ndarray):
        h, w = img.shape[-2:]
        if self._img_sz is None:
            return (
                img,
                {"h": [0, h], "w": [0, w], "hflip": False, "wflip": False},
            )

        if self._enable_random_cropping:
            # this parameter is ambiguous. It toggles between random/deterministic patching
            patch_start_loc = self._get_random_hw(h, w)
            if self._3Ddata:
                patch_start_loc = (
                    np.random.choice(1 + img.shape[-3] - self._depth3D),
                ) + patch_start_loc
        else:
            # Patch coordinates are calculated by the index manager.
            patch_start_loc = self._get_deterministic_loc(ch_idx, patch_idx)
        cropped_img = self._crop_flip_img(img, patch_start_loc, False, False)

        return cropped_img

    def _crop_img(self, img: np.ndarray, patch_start_loc: tuple):
        if self._tiling_mode in [TilingMode.TrimBoundary, TilingMode.ShiftBoundary]:
            # In training, this is used.
            # NOTE: It is my opinion that if I just use self._crop_img_with_padding, it will work perfectly fine.
            # The only benefit this if else loop provides is that it makes it easier to see what happens during training.
            patch_end_loc = (
                np.array(patch_start_loc, dtype=np.int32)
                + self.idx_manager.patch_shape[1:-1]
            )
            if self._3Ddata:
                z_start, h_start, w_start = patch_start_loc
                z_end, h_end, w_end = patch_end_loc
                new_img = img[..., z_start:z_end, h_start:h_end, w_start:w_end]
            else:
                h_start, w_start = patch_start_loc
                h_end, w_end = patch_end_loc
                new_img = img[..., h_start:h_end, w_start:w_end]

            return new_img
        else:
            # During evaluation, this is used. In this situation, we can have negative h_start, w_start. Or h_start +self._img_sz can be larger than frame
            # In these situations, we need some sort of padding. This is not needed  in the LeftTop alignement.
            return self._crop_img_with_padding(img, patch_start_loc)

    def get_begin_end_padding(self, start_pos, end_pos, max_len):
        """
        The effect is that the image with size self._grid_sz is in the center of the patch with sufficient
        padding on all four sides so that the final patch size is self._img_sz.
        """
        pad_start = 0
        pad_end = 0
        if start_pos < 0:
            pad_start = -1 * start_pos

        pad_end = max(0, end_pos - max_len)

        return pad_start, pad_end

    def _crop_img_with_padding(
        self, img: np.ndarray, patch_start_loc, max_len_vals=None
    ):
        if max_len_vals is None:
            max_len_vals = self.idx_manager.data_shape[1:-1]
        patch_end_loc = np.array(patch_start_loc, dtype=int) + np.array(
            self.idx_manager.patch_shape[1:-1], dtype=int
        )
        boundary_crossed = []
        valid_slice = []
        padding = [[0, 0]]
        for start_idx, end_idx, max_len in zip(
            patch_start_loc, patch_end_loc, max_len_vals
        ):
            boundary_crossed.append(end_idx > max_len or start_idx < 0)
            valid_slice.append((max(0, start_idx), min(max_len, end_idx)))
            pad = [0, 0]
            if boundary_crossed[-1]:
                pad = self.get_begin_end_padding(start_idx, end_idx, max_len)
            padding.append(pad)
        # max() is needed since h_start could be negative.
        if self._3Ddata:
            new_img = img[
                ...,
                valid_slice[0][0] : valid_slice[0][1],
                valid_slice[1][0] : valid_slice[1][1],
                valid_slice[2][0] : valid_slice[2][1],
            ]
        else:
            new_img = img[
                ...,
                valid_slice[0][0] : valid_slice[0][1],
                valid_slice[1][0] : valid_slice[1][1],
            ]

        # print(np.array(padding).shape, img.shape, new_img.shape)
        # print(padding)
        if not np.all(padding == 0):
            new_img = np.pad(new_img, padding, **self._overlapping_padding_kwargs)

        return new_img

    def _crop_flip_img(
        self, img: np.ndarray, patch_start_loc: tuple, h_flip: bool, w_flip: bool
    ):
        new_img = self._crop_img(img, patch_start_loc)
        if h_flip:
            new_img = new_img[..., ::-1, :]
        if w_flip:
            new_img = new_img[..., :, ::-1]

        return new_img.astype(np.float32)

    def _load_img(self, ch_idx: int, patch_idx: int) -> tuple[np.ndarray, np.ndarray]:
        """
        Returns the channels and also the respective noise channels.
        """
        patch_loc_list = self.idx_manager.get_patch_location_from_patch_idx(
            ch_idx, patch_idx
        )
        # TODO we should be adding channel dim here probably
        img = self._data[ch_idx][patch_loc_list[0]]
        return img

    def get_mean_std(self):
        return self._mean, self._std

    def set_mean_std(self, mean_val, std_val):
        self._mean = mean_val
        self._std = std_val

    def normalize_target(self, target):
        mean_dict, std_dict = self.get_mean_std()
        mean_ = mean_dict["target"]  # .squeeze(0)
        std_ = std_dict["target"]  # .squeeze(0)
        return (target - mean_) / std_

    def get_grid_size(self):
        return self._grid_sz

    def get_idx_manager(self):
        return self.idx_manager

    def per_side_overlap_pixelcount(self):
        return (self._img_sz - self._grid_sz) // 2

    def _get_deterministic_loc(self, ch_idx: int, patch_idx: int):
        """
        It returns the top-left corner of the patch corresponding to index.
        """
        loc_list = self.idx_manager.get_patch_location_from_patch_idx(ch_idx, patch_idx)
        # last dim is channel. we need to take the third and the second last element.
        return loc_list[2:]

    @cache
    def crop_probablities(self, ch_idx):
        sizes = np.array([np.prod(x.shape) for x in self._data[ch_idx]])
        return sizes / sizes.sum()

    def sample_crop(self, ch_idx):
        idx = None
        count = 0
        while idx is None:
            count += 1
            idx = np.random.choice(
                len(self._data[ch_idx]), p=self.crop_probablities(ch_idx)
            )
            data = self._data[ch_idx][idx]  # TODO no channel and S dim ?
            # changed for ndim
            if all(
                d >= self._img_sz for d in data.shape[-2:]
            ):  # TODO dims were hardcoded
                h = np.random.randint(0, data.shape[-2] - self._img_sz)
                w = np.random.randint(0, data.shape[-1] - self._img_sz)

                if len(data.shape) > 2 and not self._3Ddata:
                    s = np.random.randint(0, data.shape[0] - 1)
                    return data[s, h : h + self._img_sz, w : w + self._img_sz]
                else:
                    return data[h : h + self._img_sz, w : w + self._img_sz]

            elif count > 100:
                raise ValueError("Cannot find a valid crop")
            else:
                idx = None

        return None

    def _l2(self, x):
        return np.sqrt(np.mean(np.array(x) ** 2))

    def compute_mean_std(self, allow_for_validation_data=False):
        """
        Note that we must compute this only for training data.
        """
        if self._3Ddata:
            raise NotImplementedError("Not implemented for 3D data")

        if self._input_is_sum:
            mean_tar_dict = defaultdict(list)
            std_tar_dict = defaultdict(list)
            mean_inp = []
            std_inp = []
            for _ in range(30000):
                crops = []
                for ch_idx in range(len(self._data)):
                    crop = self.sample_crop(ch_idx)
                    mean_tar_dict[ch_idx].append(np.mean(crop))
                    std_tar_dict[ch_idx].append(np.std(crop))
                    crops.append(crop)

                inp = 0
                for img in crops:
                    inp += img

                mean_inp.append(np.mean(inp))
                std_inp.append(np.std(inp))

            output_mean = defaultdict(list)
            output_std = defaultdict(list)

            NC = len(self._data)
            for ch_idx in range(NC):
                output_mean["target"].append(np.mean(mean_tar_dict[ch_idx]))
                output_std["target"].append(self._l2(std_tar_dict[ch_idx]))

            output_mean["target"] = np.array(output_mean["target"]).reshape(NC, 1, 1)
            output_std["target"] = np.array(output_std["target"]).reshape(NC, 1, 1)

            output_mean["input"] = np.array([np.mean(mean_inp)]).reshape(1, 1, 1)
            output_std["input"] = np.array([self._l2(std_inp)]).reshape(1, 1, 1)
        else:
            raise NotImplementedError("Not implemented for non-summed input")

        return dict(output_mean), dict(output_std)

    def set_mean_std(self, mean_dict, std_dict):
        self._data_mean = mean_dict
        self._data_std = std_dict

    def get_mean_std(self):
        return self._data_mean, self._data_std

    def _get_random_hw(self, h: int, w: int):
        """
        Random starting position for the crop for the img with index `index`.
        """
        if h != self._img_sz:
            h_start = np.random.choice(h - self._img_sz)
            w_start = np.random.choice(w - self._img_sz)
        else:
            h_start = 0
            w_start = 0
        return h_start, w_start

    def replace_with_empty_patch(self, img_tuples):
        """
        Replaces the content of one of the channels with background
        """
        empty_index = self._empty_patch_fetcher.sample()
        empty_img_tuples, empty_img_noise_tuples = self._get_img(empty_index)
        assert (
            len(empty_img_noise_tuples) == 0
        ), "Noise is not supported with empty patch replacement"
        final_img_tuples = []
        for tuple_idx in range(len(img_tuples)):
            if tuple_idx == self._empty_patch_replacement_channel_idx:
                final_img_tuples.append(empty_img_tuples[tuple_idx])
            else:
                final_img_tuples.append(img_tuples[tuple_idx])
        return tuple(final_img_tuples)

    def get_mean_std_for_input(self):
        mean, std = self.get_mean_std()
        return mean["input"], std["input"]

    def _compute_target(self, img_tuples, alpha):
        if self._tar_idx_list is not None and isinstance(self._tar_idx_list, int):
            target = img_tuples[self._tar_idx_list]
        else:
            if self._tar_idx_list is not None:
                assert isinstance(self._tar_idx_list, list) or isinstance(
                    self._tar_idx_list, tuple
                )
                img_tuples = [img_tuples[i] for i in self._tar_idx_list]

            target = np.stack(img_tuples, axis=0)
        return target

    def _compute_input_with_alpha(self, img_tuples, alpha_list):
        # assert self._normalized_input is True, "normalization should happen here"
        if self._input_idx is not None:
            inp = img_tuples[self._input_idx]
        else:
            inp = 0
            for alpha, img in zip(alpha_list, img_tuples):
                inp += img * alpha

            if self._normalized_input is False:
                return inp.astype(np.float32)

        mean, std = self.get_mean_std_for_input()
        mean = mean.squeeze()
        std = std.squeeze()
        if mean.size == 1:
            mean = mean.reshape(
                1,
            )
            std = std.reshape(
                1,
            )

        for i in range(len(mean)):
            assert mean[0] == mean[i]
            assert std[0] == std[i]

        inp = (inp - mean[0]) / std[0]
        return inp.astype(np.float32)

    def _sample_alpha(self):
        alpha_arr = []
        for i in range(self._num_channels):
            alpha_pos = np.random.rand()
            alpha = self._start_alpha_arr[i] + alpha_pos * (
                self._end_alpha_arr[i] - self._start_alpha_arr[i]
            )
            alpha_arr.append(alpha)
        return alpha_arr

    def _compute_input(self, img_tuples):
        alpha = [1 / len(img_tuples) for _ in range(len(img_tuples))]
        if self._start_alpha_arr is not None:
            alpha = self._sample_alpha()

        inp = self._compute_input_with_alpha(img_tuples, alpha)
        if self._input_is_sum:
            inp = len(img_tuples) * inp

        # TODO instead we add channel here
        if len(inp.shape) == 2 or (len(inp.shape) == 3 and self._3Ddata):
            inp = inp[None, ...]

        return inp, alpha

    def _get_index_from_valid_target_logic(self, index):
        if self._validtarget_rand_fract is not None:
            if np.random.rand() < self._validtarget_rand_fract:
                index = self._train_index_switcher.get_valid_target_index()
            else:
                index = self._train_index_switcher.get_invalid_target_index()
        return index

    def _rotate2D(self, img_tuples):
        img_kwargs = {}
        for i, img in enumerate(img_tuples):
            for k in range(len(img)):
                img_kwargs[f"img{i}_{k}"] = img[k]

        keys = list(img_kwargs.keys())
        self._rotation_transform.add_targets({k: "image" for k in keys})
        rot_dic = self._rotation_transform(image=img_tuples[0][0], **img_kwargs)

        rotated_img_tuples = []
        for i, img in enumerate(img_tuples):
            if len(img) == 1:
                rotated_img_tuples.append(rot_dic[f"img{i}_0"][None])
            else:
                rotated_img_tuples.append(
                    np.concatenate(
                        [rot_dic[f"img{i}_{k}"][None] for k in range(len(img))], axis=0
                    )
                )

        return rotated_img_tuples

    def _rotate3D(self, img_tuples):
        img_kwargs = {}
        # random flip in z direction
        flip_z = self._flipz_3D and np.random.rand() < 0.5
        for i, img in enumerate(img_tuples):
            for j in range(self._depth3D):
                for k in range(len(img)):
                    if flip_z:
                        z_idx = self._depth3D - 1 - j
                    else:
                        z_idx = j
                    img_kwargs[f"img{i}_{z_idx}_{k}"] = img[k, j]

        keys = list(img_kwargs.keys())
        self._rotation_transform.add_targets({k: "image" for k in keys})
        rot_dic = self._rotation_transform(image=img_tuples[0][0][0], **img_kwargs)
        rotated_img_tuples = []
        for i, img in enumerate(img_tuples):
            if len(img) == 1:
                rotated_img_tuples.append(
                    np.concatenate(
                        [
                            rot_dic[f"img{i}_{j}_0"][None, None]
                            for j in range(self._depth3D)
                        ],
                        axis=1,
                    )
                )
            else:
                temp_arr = []
                for k in range(len(img)):
                    temp_arr.append(
                        np.concatenate(
                            [
                                rot_dic[f"img{i}_{j}_{k}"][None, None]
                                for j in range(self._depth3D)
                            ],
                            axis=1,
                        )
                    )
                rotated_img_tuples.append(np.concatenate(temp_arr, axis=0))

        return rotated_img_tuples

    def _rotate(self, img_tuples, noise_tuples):

        if self._3Ddata:
            return self._rotate3D(img_tuples, noise_tuples)
        else:
            return self._rotate2D(img_tuples, noise_tuples)

    def _get_img(self, ch_idx: int, patch_idx: int):
        """
        Loads an image.
        Crops the image such that cropped image has content.
        """
        img = self._load_img(ch_idx, patch_idx)
        cropped_img = self._crop_imgs(ch_idx, patch_idx, img)
        return cropped_img

    def get_uncorrelated_img_tuples(self, index):
        """
        Content of channels like actin and nuclei is "correlated" in its
        respective location, this function allows to pick channels' content
        from different patches of the image to make it "uncorrelated".
        """
        img_tuples = []
        for ch_idx in range(len(self._data)):
            if ch_idx == 0:
                # dataset index becomes sample index because all channels have the same
                # length
                img_tuples.append(self._get_img(0, index))
            else:
                # get a random index from corresponding channel
                sample_index = np.random.randint(
                    self.idx_manager.total_grid_count()[0][ch_idx]
                )
                img_tuples.append(self._get_img(ch_idx, sample_index))
        return img_tuples

    def __getitem__(
        self, index: Union[int, tuple[int, int]]
    ) -> tuple[np.ndarray, np.ndarray]:

        # Uncorrelated channels means crops to create the input are taken from different
        # spatial locations of the image.
        if (
            self._uncorrelated_channels
            and np.random.rand() < self._uncorrelated_channel_probab
        ):
            input_tuples = self.get_uncorrelated_img_tuples(index)
        else:
            # 0 is the channel index, because in this case locations are the same for
            # all channels
            # tuple for compatibility with _compute_input. #TODO check
            input_tuples = (self._get_img(0, index),)

        if self._enable_rotation:
            input_tuples = self._rotate(input_tuples)

        # Weight the individual channels, typically alpha is fixed
        inp, alpha = self._compute_input(input_tuples)

        target = self._compute_target(input_tuples, alpha)
        norm_target = self.normalize_target(target)

        return inp, norm_target


class LCMultiChDloaderRef(MultiChDloaderRef):
    def __init__(
        self,
        data_config: MicroSplitDataConfig,
        fpath: str,
        load_data_fn: Callable,
        val_fraction=None,
        test_fraction=None,
    ):
        self._padding_kwargs = (
            data_config.padding_kwargs  # mode=padding_mode, constant_values=constant_value
        )
        self._uncorrelated_channel_probab = data_config.uncorrelated_channel_probab

        super().__init__(
            data_config,
            fpath,
            load_data_fn=load_data_fn,
            val_fraction=val_fraction,
            test_fraction=test_fraction,
        )

        if data_config.overlapping_padding_kwargs is not None:
            assert (
                self._padding_kwargs == data_config.overlapping_padding_kwargs
            ), "During evaluation, overlapping_padding_kwargs should be same as padding_args. \
                It should be so since we just use overlapping_padding_kwargs when it is not None"

        else:
            self._overlapping_padding_kwargs = data_config.padding_kwargs

        self.multiscale_lowres_count = data_config.multiscale_lowres_count
        assert self.multiscale_lowres_count is not None
        self._scaled_data = [self._data]
        self._scaled_noise_data = [self._noise_data]

        assert (
            isinstance(self.multiscale_lowres_count, int)
            and self.multiscale_lowres_count >= 1
        )
        assert isinstance(self._padding_kwargs, dict)
        assert "mode" in self._padding_kwargs

        for _ in range(1, self.multiscale_lowres_count):
            shape = self._scaled_data[-1].shape
            assert len(shape) == 4
            new_shape = (shape[0], shape[1] // 2, shape[2] // 2, shape[3])
            ds_data = resize(
                self._scaled_data[-1].astype(np.float32), new_shape
            ).astype(self._scaled_data[-1].dtype)
            # NOTE: These asserts are important. the resize method expects np.float32. otherwise, one gets weird results.
            assert (
                ds_data.max() / self._scaled_data[-1].max() < 5
            ), "Downsampled image should not have very different values"
            assert (
                ds_data.max() / self._scaled_data[-1].max() > 0.2
            ), "Downsampled image should not have very different values"

            self._scaled_data.append(ds_data)
            # do the same for noise
            if self._noise_data is not None:
                noise_data = resize(self._scaled_noise_data[-1], new_shape)
                self._scaled_noise_data.append(noise_data)

    def reduce_data(
        self, t_list=None, h_start=None, h_end=None, w_start=None, w_end=None
    ):
        assert t_list is not None
        assert h_start is None
        assert h_end is None
        assert w_start is None
        assert w_end is None

        self._data = self._data[t_list].copy()
        self._scaled_data = [
            self._scaled_data[i][t_list].copy() for i in range(len(self._scaled_data))
        ]

        if self._noise_data is not None:
            self._noise_data = self._noise_data[t_list].copy()
            self._scaled_noise_data = [
                self._scaled_noise_data[i][t_list].copy()
                for i in range(len(self._scaled_noise_data))
            ]

        self.N = len(t_list)
        # TODO where tf is self._img_sz defined?
        self.set_img_sz([self._img_sz, self._img_sz], self._grid_sz)
        print(
            f"[{self.__class__.__name__}] Data reduced. New data shape: {self._data.shape}"
        )

    def _init_msg(self):
        msg = super()._init_msg()
        msg += f" Pad:{self._padding_kwargs}"
        if self._uncorrelated_channels:
            msg += f" UncorrChProbab:{self._uncorrelated_channel_probab}"
        return msg

    def _load_scaled_img(
        self, scaled_index, index: Union[int, tuple[int, int]]
    ) -> tuple[np.ndarray, np.ndarray]:
        if isinstance(index, int):
            idx = index
        else:
            idx, _ = index

        # tidx = self.idx_manager.get_t(idx)
        patch_loc_list = self.idx_manager.get_patch_location_from_dataset_idx(idx)
        nidx = patch_loc_list[0]

        imgs = self._scaled_data[scaled_index][nidx]
        imgs = tuple([imgs[None, ..., i] for i in range(imgs.shape[-1])])
        if self._noise_data is not None:
            noisedata = self._scaled_noise_data[scaled_index][nidx]
            noise = tuple([noisedata[None, ..., i] for i in range(noisedata.shape[-1])])
            factor = np.sqrt(2) if self._input_is_sum else 1.0
            imgs = tuple([img + noise[0] * factor for img in imgs])
        return imgs

    def _crop_img(self, img: np.ndarray, patch_start_loc: tuple):
        """
        Here, h_start, w_start could be negative. That simply means we need to pick the content from 0. So,
        the cropped image will be smaller than self._img_sz * self._img_sz
        """
        max_len_vals = list(self.idx_manager.data_shape[1:-1])
        max_len_vals[-2:] = img.shape[-2:]
        return self._crop_img_with_padding(
            img, patch_start_loc, max_len_vals=max_len_vals
        )

    def _get_img(self, index: int):
        """
        Returns the primary patch along with low resolution patches centered on the primary patch.
        """
        # Noise_tuples is populated when there is synthetic noise in training
        # Should have similar type of noise with the noise model
        # Starting with microsplit, dump the noise, use it instead as an augmentation if nessesary
        img_tuples, noise_tuples = self._load_img(index)
        assert self._img_sz is not None
        h, w = img_tuples[0].shape[-2:]
        if self._enable_random_cropping:
            patch_start_loc = self._get_random_hw(h, w)
            if self._3Ddata:
                patch_start_loc = (
                    np.random.choice(img_tuples[0].shape[-3] - self._depth3D),
                ) + patch_start_loc
        else:
            patch_start_loc = self._get_deterministic_loc(index)

        # LC logic is located here, the function crops the image of the highest resolution
        cropped_img_tuples = [
            self._crop_flip_img(img, patch_start_loc, False, False)
            for img in img_tuples
        ]
        cropped_noise_tuples = [
            self._crop_flip_img(noise, patch_start_loc, False, False)
            for noise in noise_tuples
        ]
        patch_start_loc = list(patch_start_loc)
        h_start, w_start = patch_start_loc[-2], patch_start_loc[-1]
        h_center = h_start + self._img_sz // 2
        w_center = w_start + self._img_sz // 2
        allres_versions = {
            i: [cropped_img_tuples[i]] for i in range(len(cropped_img_tuples))
        }
        for scale_idx in range(1, self.multiscale_lowres_count):
            # Returning the image of the lower resolution
            scaled_img_tuples = self._load_scaled_img(scale_idx, index)

            h_center = h_center // 2
            w_center = w_center // 2

            h_start = h_center - self._img_sz // 2
            w_start = w_center - self._img_sz // 2
            patch_start_loc[-2:] = [h_start, w_start]
            scaled_cropped_img_tuples = [
                self._crop_flip_img(img, patch_start_loc, False, False)
                for img in scaled_img_tuples
            ]
            for ch_idx in range(len(img_tuples)):
                allres_versions[ch_idx].append(scaled_cropped_img_tuples[ch_idx])

        output_img_tuples = tuple(
            [
                np.concatenate(allres_versions[ch_idx])
                for ch_idx in range(len(img_tuples))
            ]
        )
        return output_img_tuples, cropped_noise_tuples

    def __getitem__(self, index: Union[int, tuple[int, int]]):
        img_tuples, noise_tuples = self._get_img(index)
        if self._uncorrelated_channels:
            assert (
                self._input_idx is None
            ), "Uncorrelated channels is not implemented when there is a separate input channel."
            if np.random.rand() < self._uncorrelated_channel_probab:
                img_tuples_new = [None] * len(img_tuples)
                img_tuples_new[0] = img_tuples[0]
                for i in range(1, len(img_tuples)):
                    new_index = np.random.randint(len(self))
                    img_tuples_tmp, _ = self._get_img(new_index)
                    img_tuples_new[i] = img_tuples_tmp[i]
                img_tuples = img_tuples_new

        if self._is_train:
            if self._empty_patch_replacement_enabled:
                if np.random.rand() < self._empty_patch_replacement_probab:
                    img_tuples = self.replace_with_empty_patch(img_tuples)

        if self._enable_rotation:
            img_tuples, noise_tuples = self._rotate(img_tuples, noise_tuples)

        # add noise to input, if noise is present combine it with the image
        # factor is for the compute input not to have too much noise because the average of two gaussians
        if len(noise_tuples) > 0:
            factor = np.sqrt(2) if self._input_is_sum else 1.0
            input_tuples = []
            for x in img_tuples:
                x = (
                    x.copy()
                )  # to avoid changing the original image since it is later used for target
                # NOTE: other LC levels already have noise added. So, we just need to add noise to the highest resolution.
                x[0] = x[0] + noise_tuples[0] * factor
                input_tuples.append(x)
        else:
            input_tuples = img_tuples

        # Compute the input by sum / average the channels
        # Alpha is an amount of weight which is applied to the channels when combining them
        # How to sample alpha is still under research
        inp, alpha = self._compute_input(input_tuples)
        target_tuples = [img[:1] for img in img_tuples]
        # add noise to target.
        if len(noise_tuples) >= 1:
            target_tuples = [
                x + noise for x, noise in zip(target_tuples, noise_tuples[1:])
            ]

        target = self._compute_target(target_tuples, alpha)

        norm_target = self.normalize_target(target)

        output = [inp, norm_target]

        if self._return_alpha:
            output.append(alpha)

        if isinstance(index, int):
            return tuple(output)

        _, grid_size = index
        output.append(grid_size)
        return tuple(output)
